From f1e9a74f003c94b7e7fdcfd2576da0e390891224 Mon Sep 17 00:00:00 2001 From: Dhruv Vaishnav Date: Tue, 22 Sep 2026 07:35:07 +0000 Subject: [PATCH 001/133] fix(meta): populate physical metric table column ids (#9286) * fix(meta): populate physical metric table column ids Signed-off-by: dhruvxvaishnav * test(meta): verify physical metric column ids Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav --- .../meta/src/ddl/create_table/executor.rs | 14 ++- tests-integration/tests/repartition.rs | 95 +++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/src/common/meta/src/ddl/create_table/executor.rs b/src/common/meta/src/ddl/create_table/executor.rs index 88d88fe277c..332c9e57eb4 100644 --- a/src/common/meta/src/ddl/create_table/executor.rs +++ b/src/common/meta/src/ddl/create_table/executor.rs @@ -182,7 +182,19 @@ impl CreateTableExecutor { region_wal_options: RegionWalOptions, ) -> Result<()> { if !column_metadatas.is_empty() { - update_table_info_column_ids(&mut table_info, column_metadatas); + // Region metadata can include engine-internal columns that are not visible in + // the table schema. Only visible columns belong in the table's column IDs. + let visible_column_metadatas = column_metadatas + .iter() + .filter(|column| { + table_info + .meta + .schema + .contains_column(&column.column_schema.name) + }) + .cloned() + .collect::>(); + update_table_info_column_ids(&mut table_info, &visible_column_metadatas); } let detecting_regions = convert_region_routes_to_detecting_regions(&table_route.region_routes); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index c243113d928..5fe84b19f24 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -117,6 +117,101 @@ macro_rules! repartition_tests { }; } +#[tokio::test(flavor = "multi_thread")] +async fn test_repartition_physical_metric_without_logical_table() { + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&StorageType::File); + let home_dir = create_temp_dir("repartition_physical_metric_without_logical_table"); + let cluster = + GreptimeDbClusterBuilder::new("test_repartition_physical_metric_without_logical_table") + .await + .with_datanodes(3) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .with_datanode_wal_config(DatanodeWalConfig::Noop) + .with_metasrv_gc_config(GcSchedulerOptions { + enable: true, + gc_cooldown_period: Duration::from_nanos(1), + ..Default::default() + }) + .with_datanode_gc_config(GcConfig { + enable: true, + lingering_time: Some(Duration::from_secs(0)), + unknown_file_lingering_time: Duration::from_secs(0), + ..Default::default() + }) + .build(true) + .await; + + let instance = cluster.fe_instance(); + let query_ctx = QueryContext::arc(); + let sql = r#" + CREATE TABLE `physical_metric_without_logical`( + `ts` TIMESTAMP TIME INDEX, + `val` DOUBLE, + `host` STRING PRIMARY KEY + ) PARTITION ON COLUMNS (`host`) ( + `host` < 'm', + `host` >= 'm' + ) ENGINE = metric + WITH ( + "physical_metric_table" = "true" + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_id = get_table_id(&cluster.metasrv, "physical_metric_without_logical").await; + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + let column_ids_before_split = table_info.table_info.meta.column_ids.clone(); + assert_eq!(column_ids_before_split, vec![0, 1, 2]); + + let sql = r#" + ALTER TABLE `physical_metric_without_logical` SPLIT PARTITION ( + `host` < 'm' + ) INTO ( + `host` < 'g', + `host` >= 'g' AND `host` < 'm' + ); + "#; + run_sql(instance, sql, query_ctx.clone()).await.unwrap(); + + let table_info = cluster + .metasrv + .table_metadata_manager() + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap(); + assert_eq!( + table_info.table_info.meta.column_ids, + column_ids_before_split + ); + + let result = run_sql( + instance, + &query_partitions_sql("physical_metric_without_logical"), + query_ctx, + ) + .await + .unwrap(); + let expected = r#"+---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| table_catalog | table_schema | table_name | partition_name | partition_expression | partition_description | greptime_partition_id | partition_ordinal_position | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+ +| greptime | public | physical_metric_without_logical | p0 | host | host < g | 4398046511104 | 1 | +| greptime | public | physical_metric_without_logical | p1 | host | host >= m | 4398046511105 | 2 | +| greptime | public | physical_metric_without_logical | p2 | host | host >= g AND host < m | 4398046511106 | 3 | ++---------------+--------------+---------------------------------+----------------+----------------------+------------------------+-----------------------+----------------------------+"#; + check_output_stream(result.data, expected).await; +} + /// COUNT must use visible rows rather than the full row counts of shared SSTs. #[tokio::test(flavor = "multi_thread")] async fn test_repartition_append_count_file() { From 614b46e5f52586ad54f617822ff3f026491ee88e Mon Sep 17 00:00:00 2001 From: houyuwushang <180804215+houyuwushang@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:03:40 +0000 Subject: [PATCH 002/133] fix(postgres): return empty responses for comment-only SQL (#9295) fix(postgres): handle parsed empty queries in both protocols Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com> --- src/servers/src/postgres/handler.rs | 7 +++ src/servers/tests/postgres/mod.rs | 86 +++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 204fa7745a5..e618c55614b 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -82,6 +82,10 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { let parsed_query = self.query_parser.compatibility_parser.parse(query); let query = if let Ok(statements) = &parsed_query { + // Comments, whitespace and empty statements also require EmptyQueryResponse. + if statements.is_empty() { + return Ok(vec![Response::EmptyQuery]); + } statements .iter() .map(|s| s.to_string()) @@ -335,6 +339,9 @@ impl QueryParser for DefaultQueryParser { let parsed_statements = self.compatibility_parser.parse(sql); let (sql, copy_to_stdout_format) = if let Ok(mut statements) = parsed_statements { + if statements.is_empty() { + return Ok(None); + } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); (first_stmt.to_string(), format) diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 2e031c0d0db..6b2995283f8 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -587,6 +587,92 @@ async fn start_test_server(server_tls: TlsOption) -> Result { Ok(server_addr.port()) } +const EMPTY_QUERIES: &[&str] = &[ + "", + " \t\r\n", + ";", + ";;;", + "-- ping", + "-- ping\n", + "/* ping */", + "/* outer /* inner */ comment */", + "; -- ping\r\n /* comment */ ;", +]; + +#[tokio::test] +async fn test_simple_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let messages = client + .simple_query(query) + .await + .unwrap_or_else(|err| panic!("query {query:?} failed: {err}")); + // tokio-postgres exposes EmptyQueryResponse as CommandComplete(0). + assert!( + matches!( + messages.as_slice(), + [SimpleQueryMessage::CommandComplete(0)] + ), + "expected an empty query response for {query:?}, got {messages:?}" + ); + } + + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + +#[tokio::test] +async fn test_extended_query_empty_statements() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for query in EMPTY_QUERIES { + let statement = client + .prepare(query) + .await + .unwrap_or_else(|err| panic!("prepare {query:?} failed: {err}")); + assert_eq!(client.execute(&statement, &[]).await.unwrap(), 0); + assert!(client.query(&statement, &[]).await.unwrap().is_empty()); + } + + let row = client + .query_one("SELECT $1::BIGINT", &[&42i64]) + .await + .unwrap(); + assert_eq!(row.get::<_, i64>(0), 42); + Ok(()) +} + +#[tokio::test] +async fn test_simple_query_with_comments() -> Result<()> { + let server_port = start_test_server(Default::default()).await?; + let client = create_plain_connection(server_port, false).await.unwrap(); + + for (query, expected) in [ + ("-- ping\nSELECT 1", vec!["1"]), + ("/* comment */ SELECT 1; -- trailing", vec!["1"]), + ("SELECT '-- ping'", vec!["-- ping"]), + ("SELECT '/* comment */'", vec!["/* comment */"]), + ("SELECT 1; /* between */ SELECT 2;", vec!["1", "2"]), + ] { + let messages = client.simple_query(query).await.unwrap(); + assert_eq!(unwrap_results(&messages), expected, "query: {query}"); + } + + assert!( + client + .simple_query("/* comment */ SELECT missing_column FROM numbers") + .await + .is_err() + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert_eq!(unwrap_results(&messages), vec!["1"]); + Ok(()) +} + async fn do_simple_query(server_tls: TlsOption, client_tls: bool) -> Result<()> { let server_port = start_test_server(server_tls).await?; From 2388257c35beb16e50132d35b077ae26a4183607 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 22 Sep 2026 08:37:21 +0000 Subject: [PATCH 003/133] ci: create docs follow-up issue on PR merge instead of on label (#9237) * ci: create docs follow-up issue on PR merge instead of on label The docbot workflow previously created a docs-repo issue as soon as the 'docs-required' condition was detected (PR opened/edited with the docs checkbox ticked), even if the PR was never merged. Now the workflow also triggers on PR 'closed': - opened/edited: only manage the docs-required/docs-not-required labels - closed: create the docs issue only when the PR was actually merged and carries the docs-required label This also lets maintainers control issue creation by manually adding or removing the docs-required label before merging. Signed-off-by: Ning Sun * fix: address review comments on docs issue creation timing - Only touch docs labels when the docs checkbox state actually changed in an edit. Previously, editing any other part of the PR body while the checkbox stayed checked removed the docs-required label, silently dropping the docs follow-up now that issue creation happens at merge. Unchanged checkbox now leaves labels untouched, which also preserves manual label overrides. - Do not trust the closed event's stale label snapshot at merge time: re-read the live PR via the API and create the docs issue if the docs-required label is present OR the checkbox is ticked in the current body. - Make the workflow concurrency group action-aware so a merge run does not cancel an in-flight label update from an edit run. Signed-off-by: Ning Sun * fix: make docs-required label the single source of truth at merge The label-OR-checkbox merge condition could not distinguish an intentional opt-out from an unfinished label update: removing docs-required while the checkbox stayed checked still produced an issue, and unchecking the box could still produce one if the merge read the stale label before the edit run removed it. At merge time, wait for any pending docbot runs on the PR head SHA to finish their label updates (bounded to 5 minutes), then decide solely by the live docs-required label. Adds actions: read permission for listing workflow runs. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .github/workflows/docbot.yml | 5 +- cyborg/bin/follow-up-docs-issue.ts | 153 +++++++++++++++++++++-------- 2 files changed, 114 insertions(+), 44 deletions(-) diff --git a/.github/workflows/docbot.yml b/.github/workflows/docbot.yml index 250a150c026..53033c6eb5b 100644 --- a/.github/workflows/docbot.yml +++ b/.github/workflows/docbot.yml @@ -1,10 +1,10 @@ name: Follow Up Docs on: pull_request_target: - types: [opened, edited] + types: [opened, edited, closed] concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + group: ${{ github.workflow }}-${{ github.event.action }}-${{ github.head_ref || github.run_id }} cancel-in-progress: true jobs: @@ -13,6 +13,7 @@ jobs: permissions: pull-requests: write contents: read + actions: read timeout-minutes: 10 steps: - uses: actions/checkout@v4 diff --git a/cyborg/bin/follow-up-docs-issue.ts b/cyborg/bin/follow-up-docs-issue.ts index acca21fe935..09ef5db2445 100644 --- a/cyborg/bin/follow-up-docs-issue.ts +++ b/cyborg/bin/follow-up-docs-issue.ts @@ -17,7 +17,12 @@ import * as core from '@actions/core' import {handleError, obtainClient} from "@/common"; import {context} from "@actions/github"; -import {PullRequestEditedEvent, PullRequestEvent, PullRequestOpenedEvent} from "@octokit/webhooks-types"; +import { + PullRequestClosedEvent, + PullRequestEditedEvent, + PullRequestEvent, + PullRequestOpenedEvent +} from "@octokit/webhooks-types"; // @ts-expect-error moduleResolution:nodenext issue 54523 import {RequestError} from "@octokit/request-error"; @@ -30,18 +35,59 @@ async function main() { throw new Error(`Only pull request event supported. ${context.eventName} is unsupported.`) } - const client = obtainClient("GITHUB_TOKEN") - const docsClient = obtainClient("DOCS_REPO_TOKEN") const payload = context.payload as PullRequestEvent - const { owner, repo, number, actor, title, html_url } = { + const { owner, repo, number, actor } = { owner: payload.pull_request.base.user.login, repo: payload.pull_request.base.repo.name, number: payload.pull_request.number, - title: payload.pull_request.title, - html_url: payload.pull_request.html_url, actor: payload.pull_request.user.login, } - const followUpDocs = checkPullRequestEvent(payload) + + switch (payload.action) { + case "opened": { + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent)) + break + } + case "edited": { + const followUpDocs = checkPullRequestEditedEvent(payload as PullRequestEditedEvent) + if (followUpDocs === undefined) { + core.info("Docs checkbox state unchanged; leaving labels as-is.") + return + } + const client = obtainClient("GITHUB_TOKEN") + await updateDocsLabels(client, owner, repo, number, followUpDocs) + break + } + case "closed": { + const event = payload as PullRequestClosedEvent + if (!event.pull_request.merged) { + core.info("PR closed without merging; no docs issue needed.") + return + } + // The docs-required label is the single source of truth. Coordinate + // with pending checkbox-driven label updates before deciding: an edit + // run may still be applying or removing the label, and the event + // payload's label snapshot may be stale. + const client = obtainClient("GITHUB_TOKEN") + const { data: pr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + await waitForPendingLabelUpdates(client, owner, repo, pr.head.sha) + const { data: freshPr } = await client.rest.pulls.get({ owner, repo, pull_number: number }) + const hasDocsLabel = freshPr.labels.some((label) => label.name === labelDocsRequired) + if (!hasDocsLabel) { + core.info(`Label ${labelDocsRequired} not present; no docs issue needed.`) + return + } + const docsClient = obtainClient("DOCS_REPO_TOKEN") + await createDocsIssue(docsClient, freshPr.title, freshPr.html_url, actor) + break + } + default: + throw new Error(`${payload.action} is unsupported.`) + } +} + +async function updateDocsLabels(client: ReturnType, owner: string, repo: string, number: number, followUpDocs: boolean) { if (followUpDocs) { core.info("Follow up docs.") await client.rest.issues.removeLabel({ @@ -55,28 +101,6 @@ async function main() { await client.rest.issues.addLabels({ owner, repo, issue_number: number, labels: [labelDocsRequired], }) - - // Get available assignees for the docs repo - const assigneesResponse = await docsClient.rest.issues.listAssignees({ - owner: 'GreptimeTeam', - repo: 'docs', - }) - const validAssignees = assigneesResponse.data.map(assignee => assignee.login) - core.info(`Available assignees: ${validAssignees.join(', ')}`) - - // Check if the actor is a valid assignee, otherwise fallback to fengjiachun - const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' - core.info(`Assigning issue to: ${assignee}`) - - await docsClient.rest.issues.create({ - owner: 'GreptimeTeam', - repo: 'docs', - title: `Update docs for ${title}`, - body: `A document change request is generated from ${html_url}`, - assignee: assignee, - }).then((res) => { - core.info(`Created issue ${res.data}`) - }) } else { core.info("No need to follow up docs.") await client.rest.issues.removeLabel({ @@ -93,15 +117,52 @@ async function main() { } } -function checkPullRequestEvent(payload: PullRequestEvent) { - switch (payload.action) { - case "opened": - return checkPullRequestOpenedEvent(payload as PullRequestOpenedEvent) - case "edited": - return checkPullRequestEditedEvent(payload as PullRequestEditedEvent) - default: - throw new Error(`${payload.action} is unsupported.`) +async function createDocsIssue(docsClient: ReturnType, title: string, html_url: string, actor: string) { + core.info("Creating follow-up docs issue for merged PR.") + + // Get available assignees for the docs repo + const assigneesResponse = await docsClient.rest.issues.listAssignees({ + owner: 'GreptimeTeam', + repo: 'docs', + }) + const validAssignees = assigneesResponse.data.map(assignee => assignee.login) + core.info(`Available assignees: ${validAssignees.join(', ')}`) + + // Check if the actor is a valid assignee, otherwise fallback to fengjiachun + const assignee = validAssignees.includes(actor) ? actor : 'fengjiachun' + core.info(`Assigning issue to: ${assignee}`) + + await docsClient.rest.issues.create({ + owner: 'GreptimeTeam', + repo: 'docs', + title: `Update docs for ${title}`, + body: `A document change request is generated from ${html_url}`, + assignee: assignee, + }).then((res) => { + core.info(`Created issue ${res.data}`) + }) +} + +// Waits until no other runs of this workflow are pending for the PR head SHA, +// so that any checkbox-driven label update has been applied before the caller +// reads the labels. Bounded by a deadline; on timeout it proceeds with whatever +// the current label state is. +async function waitForPendingLabelUpdates(client: ReturnType, owner: string, repo: string, headSha: string) { + const deadlineMs = 5 * 60 * 1000 + const pollIntervalMs = 5000 + const deadline = Date.now() + deadlineMs + while (Date.now() < deadline) { + const { data } = await client.rest.actions.listWorkflowRuns({ + owner, repo, workflow_id: "docbot.yml", head_sha: headSha, per_page: 20, + }) + const pending = data.workflow_runs.filter((run) => run.id !== context.runId && run.status !== "completed") + if (pending.length === 0) { + return + } + core.info(`Waiting for ${pending.length} pending docbot run(s) to finish label updates...`) + await new Promise((resolve) => setTimeout(resolve, pollIntervalMs)) } + core.warning("Timed out waiting for pending docbot runs; proceeding with the current label state.") } function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { @@ -109,11 +170,19 @@ function checkPullRequestOpenedEvent(event: PullRequestOpenedEvent): boolean { return event.pull_request.body?.includes(needFollowUpDocs) } -function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean { - const previous = event.changes.body?.from.includes(needFollowUpDocs) - const current = event.pull_request.body?.includes(needFollowUpDocs) - // from docs-not-need to docs-required - return (!previous) && current +// Returns undefined when the checkbox state did not change in this edit, so the +// caller leaves the labels untouched (preserving manual label overrides). +function checkPullRequestEditedEvent(event: PullRequestEditedEvent): boolean | undefined { + if (!event.changes.body) { + // The body was not part of this edit (e.g. title-only edit). + return undefined + } + const previous = event.changes.body.from.includes(needFollowUpDocs) + const current = event.pull_request.body?.includes(needFollowUpDocs) ?? false + if (previous === current) { + return undefined + } + return current } main().catch(handleError) From 25bca4609f99fdd46078c528e4ad84aea7f7281f Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 11:25:50 +0000 Subject: [PATCH 004/133] perf(promql): push label filters into grouped join inputs (#9280) * perf(promql): propagate matching filters through grouped joins Signed-off-by: Dennis Zhuang * perf(promql): check matcher safety on the receiving operand Signed-off-by: Dennis Zhuang * refactor(promql): spell out the shapes a filter may cross `preserves_filter` ended in `_ => true`, which was only sound because `selector_matchers` independently rejects label rewriting, `count_values`, subqueries and non-rollup calls on the same operand. Loosening the latter alone would have silently pushed a matcher below a label rewrite. List the shapes that carry a scan filter instead and default to `false`. Cite #9207 for the result labels the grouped cases record: the join projects the right operand's tag set, so `zone` is missing wherever the right side aggregates it away. No behavior change. Signed-off-by: Dennis Zhuang * test(promql): assert the new pushdowns reach the scan The grouped-join unit tests feed tag columns by hand and the SQLness case only checks results, which are identical whether or not the rewrite fires. Nothing would have failed if scalar arithmetic, ranking or grouped matching stopped propagating. Assert through the planner that the matcher reaches both scans, with a global topk one-side as the counter-example. Also state that the duplicate-one-side cases record a cross product Prometheus rejects (#9209), so the baseline is not read as intended semantics. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/planner.rs | 43 ++- .../src/promql/planner/matching_filters.rs | 300 ++++++++++++------ .../common/promql/matching_groups.result | 181 +++++++++++ .../common/promql/matching_groups.sql | 65 ++++ 4 files changed, 484 insertions(+), 105 deletions(-) create mode 100644 tests/cases/standalone/common/promql/matching_groups.result create mode 100644 tests/cases/standalone/common/promql/matching_groups.sql 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/tests/cases/standalone/common/promql/matching_groups.result b/tests/cases/standalone/common/promql/matching_groups.result new file mode 100644 index 00000000000..183ae6d218c --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.result @@ -0,0 +1,181 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); + +Affected Rows: 0 + +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); + +Affected Rows: 5 + +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +Affected Rows: 5 + +-- The result labels are the right operand's tag set rather than what the modifier implies +-- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing +-- only by it are indistinguishable. +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); + ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.Float64(8) * greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| x | 1970-01-01T00:00:00 | 24.0 | +| x | 1970-01-01T00:00:00 | 48.0 | ++------+---------------------+--------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.sum(matching_groups_left.greptime_value) / matching_groups_right.greptime_value * Float64(2) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 9.0 | +| x | b | 1970-01-01T00:00:00 | 4.5 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:00 | 6.0 | +| y | 1970-01-01T00:00:00 | 0.36 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | 1970-01-01T00:00:00 | 6.0 | +| | 1970-01-01T00:00:00 | 6.0 | ++------+---------------------+-------------------------------------------------------------------------------------------------------+ + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); + ++------+----+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+----+-------------------------------------------------------------------------------------------------------+ ++------+----+-------------------------------------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + ++------+----+-------------------------------------------------------------------------------------------------------+ +| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+----+-------------------------------------------------------------------------------------------------------+ ++------+----+-------------------------------------------------------------------------------------------------------+ + +-- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many +-- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no +-- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands +-- alone, so the recorded output is the unchanged baseline. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; + ++------+------+---------------------+----------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | ++------+------+---------------------+----------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 12.0 | +| x | a | 1970-01-01T00:00:00 | 6.0 | +| x | b | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 12.0 | +| x | a | 1970-01-01T00:00:00 | 6.0 | +| x | b | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ + +DROP TABLE matching_groups_left; + +Affected Rows: 0 + +DROP TABLE matching_groups_right; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/matching_groups.sql b/tests/cases/standalone/common/promql/matching_groups.sql new file mode 100644 index 00000000000..97e46da0437 --- /dev/null +++ b/tests/cases/standalone/common/promql/matching_groups.sql @@ -0,0 +1,65 @@ +CREATE TABLE matching_groups_left ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +CREATE TABLE matching_groups_right ( + host STRING NULL, + "zone" STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, "zone") +); +INSERT INTO matching_groups_left VALUES + ('x', 'a', 0, 12), ('x', 'b', 0, 24), ('y', 'a', 0, 36), + ('', 'a', 0, 48), (NULL, 'a', 0, 60); +INSERT INTO matching_groups_right VALUES + ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), + ('', 'a', 0, 8), (NULL, 'a', 0, 10); + +-- The result labels are the right operand's tag set rather than what the modifier implies +-- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing +-- only by it are indistinguishable. +-- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); + +-- bottomk partitions by host when zone is excluded from grouping. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') bottomk without(zone)(1, matching_groups_left) / on(host,zone) matching_groups_right{host="x"}; + +-- A matcher can leave a global ranking operand without changing its candidates. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') topk(1, matching_groups_left{host="x"}) / on(host,zone) matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_groups_right{host="x"}); + +-- The global winner is host y. Pushing host=x below either topk would change the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); + +-- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many +-- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no +-- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands +-- alone, so the recorded output is the unchanged baseline. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); + +DROP TABLE matching_groups_left; +DROP TABLE matching_groups_right; From aa36f74feb5446e0a58237e8483e28187dbef0a8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 22 Sep 2026 14:38:22 +0000 Subject: [PATCH 005/133] fix(ci): build tests-integration lib with meta-srv/mock (#9299) * fix(ci): build tests-integration lib with meta-srv/mock tests-integration's lib code (src/cluster.rs) uses meta_srv::mocks, but the dependency carrying the mock feature sits in [dev-dependencies]. Builds that only touch the lib, such as the apidoc job's cargo doc --workspace, resolve meta-srv without mock and fail with E0432. --all-targets builds unify dev-dependency features, which is why check, clippy and nextest stayed green. Move the mock-enabled meta-srv entry back to [dependencies]. The other testing features moved out in #9072 are not needed by the lib and stay in [dev-dependencies]. Signed-off-by: Dennis Zhuang * test(repartition): split per-case repartition tests test_repartition_metric ran four format/primary-key-encoding cases in a single test function, and test_repartition_mito ran two format cases. Each case builds its own 3-datanode cluster and runs a full repartition plus GC cycle, so on S3 the metric test took 165-178s against the 180s nextest terminate-after. Merge queue runs failed on it at random. Split each case into its own test. Cases were already independent, so they now run in parallel and each stays far inside the timeout, and a failure points at one encoding instead of four. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- tests-integration/Cargo.toml | 3 +- tests-integration/tests/repartition.rs | 52 +++++++++++++++++++++----- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 7a7d0cba72a..a5ac94924e4 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -65,7 +65,7 @@ hyper-util = { workspace = true, features = ["tokio"] } log-query = { workspace = true } loki-proto.workspace = true meta-client.workspace = true -meta-srv.workspace = true +meta-srv = { workspace = true, features = ["mock"] } mito2.workspace = true object-store.workspace = true operator.workspace = true @@ -111,7 +111,6 @@ hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -meta-srv = { workspace = true, features = ["mock"] } mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 5fe84b19f24..1198e6d23e5 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -45,15 +45,25 @@ macro_rules! repartition_tests { $( paste::item! { mod [] { + // Every case below builds its own cluster and runs a full repartition plus + // GC cycle, which takes around a minute on object-store backends. Keep one + // case per test so a case stays well inside the nextest slow timeout and + // cases run in parallel; do not fold several cases back into one test. + #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_mito >]() { + async fn [< test_repartition_mito_flat >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { common_telemetry::init_default_ut_logging(); - // Cover both storage formats for repartition behavior. - // for flat format $crate::repartition::test_repartition_mito(store_type, true).await; - // for primary key format + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_mito_primary_key >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_mito(store_type, false).await; } } @@ -95,19 +105,41 @@ macro_rules! repartition_tests { } #[tokio::test(flavor = "multi_thread")] - async fn [< test_repartition_metric >]() { + async fn [< test_repartition_metric_flat_sparse >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { use store_api::codec::PrimaryKeyEncoding; common_telemetry::init_default_ut_logging(); - // Exercise format + primary key encoding matrix for metric engine. - // for flat format with sparse primary key encoding $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Sparse).await; - // for flat format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_flat_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, true, PrimaryKeyEncoding::Dense).await; - // for primary key format with sparse primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_sparse >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Sparse).await; - // for primary key format with dense primary key encoding + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn [< test_repartition_metric_primary_key_dense >]() { + let store_type = tests_integration::test_util::StorageType::$service; + if store_type.test_on() { + use store_api::codec::PrimaryKeyEncoding; + common_telemetry::init_default_ut_logging(); $crate::repartition::test_repartition_metric(store_type, false, PrimaryKeyEncoding::Dense).await; } } From 045441e3ccce8d474e7e34d9a4e6e5b49f16a74d Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:56:26 +0000 Subject: [PATCH 006/133] feat(json2): support altering JSON2 settings (#9029) * feat(sql): support alter syntax for JSON2 columns Signed-off-by: fys * fix(json2): preserve rows on type hint mismatch during compaction * refactor(json2): simplify alter settings handling * fix(json2): preserve coerced values during compaction * chore: remove unnecessary clone * chor: reduce memory allocations * fix: cargo clippy * chore: update greptime-proto to main branch * refactor(datatypes): unify string handling with other JSON type hints * fix: cr --------- Signed-off-by: fys --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/common/grpc-expr/src/alter.rs | 51 ++- src/common/grpc-expr/src/error.rs | 8 + .../meta/src/ddl/alter_table/executor.rs | 1 + .../src/ddl/alter_table/region_request.rs | 1 + src/common/meta/src/ddl/event/table.rs | 1 + src/datatypes/src/extension/json.rs | 55 ++- src/datatypes/src/json.rs | 384 +++++++++++++++--- src/datatypes/src/vectors/json/array.rs | 105 ++--- src/metric-engine/src/data_region.rs | 1 + src/mito2/src/compaction/json2.rs | 13 +- src/mito2/src/sst/parquet/json_align.rs | 78 +++- src/operator/src/expr_helper.rs | 109 +++-- src/sql/src/parsers/alter_parser.rs | 28 +- src/sql/src/statements/alter.rs | 38 +- src/sql/src/statements/create.rs | 39 +- src/store-api/Cargo.toml | 1 + src/store-api/src/metadata.rs | 36 ++ src/store-api/src/region_request.rs | 166 ++++++++ src/table/src/metadata.rs | 166 +++++++- src/table/src/requests.rs | 11 + .../common/types/json/json2_alter.result | 131 +++++- .../common/types/json/json2_alter.sql | 58 ++- .../json2_alter_type_hints_compaction.result | 83 ++++ .../json2_alter_type_hints_compaction.sql | 36 ++ ...er_type_hints_compaction_conversion.result | 78 ++++ ...alter_type_hints_compaction_conversion.sql | 35 ++ ..._compaction_null_type_hint_mismatch.result | 80 ++++ ...on2_compaction_null_type_hint_mismatch.sql | 35 ++ 30 files changed, 1602 insertions(+), 231 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result create mode 100644 tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql create mode 100644 tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result create mode 100644 tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql 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/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/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/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/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/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/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/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/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index ca38ec0d65e..3e5c27fda73 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -462,19 +462,19 @@ impl ParserContext<'_> { .context(error::SyntaxSnafu)?; self.parse_alter_table_drop_default(column_name) } else { - let (data_type, json2_options) = - if let Some(json2) = parse_json2_type_and_options(&mut self.parser)? { - json2 - } else { - ( - self.parser.parse_data_type().context(error::SyntaxSnafu)?, - None, - ) - }; + if let Some((_, json2_options)) = + parse_json2_type_and_options(&mut self.parser)? + { + return Ok(AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + }); + } + Ok(AlterTableOperation::ModifyColumnType { column_name, - target_type: data_type, - json2_options, + target_type: self.parser.parse_data_type().context(error::SyntaxSnafu)?, + json2_options: None, }) } } @@ -1576,9 +1576,8 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(alter_table) = statements.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { + let AlterTableOperation::SetJsonSettings { column_name, - target_type, json2_options: Some(options), } = alter_table.alter_operation() else { @@ -1586,7 +1585,6 @@ MODIFY COLUMN attrs JSON2 ( }; assert_eq!("attrs", column_name.value); - assert_eq!("JSON2", target_type.to_string()); assert_eq!(Some(2000), options.max_auto_expanded_paths); assert_eq!(4, options.type_hints.len()); assert_eq!(vec!["user", "id"], options.type_hints[1].path); @@ -1610,7 +1608,7 @@ MODIFY COLUMN attrs JSON2 ( let Statement::AlterTable(empty) = empty.remove(0) else { unreachable!() }; - let AlterTableOperation::ModifyColumnType { json2_options, .. } = empty.alter_operation() + let AlterTableOperation::SetJsonSettings { json2_options, .. } = empty.alter_operation() else { unreachable!() }; diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index 2f98d8737ae..4fc45c2e69b 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -88,6 +88,11 @@ pub enum AlterTableOperation { target_type: DataType, json2_options: Option, }, + /// `MODIFY JSON2 [json2_options]` + SetJsonSettings { + column_name: Ident, + json2_options: Option, + }, /// `SET =
` SetTableOptions { options: Vec, @@ -265,6 +270,16 @@ impl Display for AlterTableOperation { } Ok(()) } + AlterTableOperation::SetJsonSettings { + column_name, + json2_options, + } => { + write!(f, r#"MODIFY COLUMN {column_name} JSON2"#)?; + if let Some(options) = json2_options { + write!(f, "{options}")?; + } + Ok(()) + } AlterTableOperation::SetTableOptions { options } => { let kvs = options .iter() @@ -433,8 +448,10 @@ impl Display for AlterDatabaseOperation { mod tests { use std::assert_matches; + use super::AlterTableOperation; use crate::dialect::GreptimeDbDialect; use crate::parser::{ParseOptions, ParserContext}; + use crate::statements::create::Json2Options; use crate::statements::statement::Statement; #[test] @@ -503,7 +520,7 @@ ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY, ADD COLUMN } let sql = r"alter table monitor modify column load_15 string;"; - let stmts = + let mut stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) .unwrap(); assert_eq!(1, stmts.len()); @@ -523,6 +540,25 @@ ALTER TABLE monitor MODIFY COLUMN load_15 STRING"#, } } + let Statement::AlterTable(alter_table) = &mut stmts[0] else { + unreachable!(); + }; + let AlterTableOperation::ModifyColumnType { json2_options, .. } = + alter_table.alter_operation_mut() + else { + unreachable!(); + }; + *json2_options = Some(Json2Options { + max_auto_expanded_paths: Some(1), + type_hints: vec![], + }); + assert_eq!( + r#"ALTER TABLE monitor MODIFY COLUMN load_15 STRING( + max_auto_expanded_paths = 1 + )"#, + alter_table.to_string() + ); + let sql = r"alter table monitor drop column load_15;"; let stmts = ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 46b5b07c9f7..2310531841c 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -157,6 +157,26 @@ impl Display for Json2Options { } } +impl Json2Options { + pub fn build_json_settings(&self) -> Result { + let type_hints = self + .type_hints + .iter() + .map(|hint| { + Ok(datatypes::json::JsonTypeHint { + path: hint.path.clone(), + data_type: json_type_hint_concrete_data_type(&hint.data_type)?, + inverted_index: hint.inverted_index, + }) + }) + .collect::>>()?; + let max_auto_expanded_paths = self + .max_auto_expanded_paths + .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)); + JsonSettings::try_new(type_hints, max_auto_expanded_paths).map_err(Into::into) + } +} + #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub struct JsonTypeHint { pub path: Vec, @@ -351,24 +371,7 @@ impl ColumnExtensions { return Ok(None); }; - let type_hints = options - .type_hints - .iter() - .map(|hint| { - Ok(datatypes::json::JsonTypeHint { - path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - inverted_index: hint.inverted_index, - }) - }) - .collect::>>()?; - let settings = JsonSettings::try_new( - type_hints, - options - .max_auto_expanded_paths - .or(Some(JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS)), - )?; - Ok(Some(settings)) + options.build_json_settings().map(Some) } pub fn set_json_settings(&mut self, settings: JsonSettings) -> Result<()> { diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml index 61ce914fe13..2e5a203912e 100644 --- a/src/store-api/Cargo.toml +++ b/src/store-api/Cargo.toml @@ -10,6 +10,7 @@ workspace = true [dependencies] api.workspace = true aquamarine.workspace = true +arrow-schema.workspace = true async-trait.workspace = true bytes.workspace = true common-base.workspace = true diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 0c663bccc05..21fc3295c44 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -30,6 +30,7 @@ use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use datatypes::arrow; use datatypes::arrow::datatypes::FieldRef; +use datatypes::extension::json::json2_metadata_with_updated_settings; use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef, VectorIndexOptions}; use datatypes::types::TimestampType; use itertools::Itertools; @@ -663,6 +664,10 @@ impl RegionMetadataBuilder { AlterKind::AddColumns { columns } => self.add_columns(columns)?, AlterKind::DropColumns { names } => self.drop_columns(&names), AlterKind::ModifyColumnTypes { columns } => self.modify_column_types(columns)?, + AlterKind::SetJsonSettings { + column_name, + settings, + } => self.set_json_settings(column_name, settings)?, AlterKind::SetIndexes { options } => self.set_indexes(options)?, AlterKind::UnsetIndexes { options } => self.unset_indexes(options)?, AlterKind::SetRegionOptions { options: _ } => { @@ -831,6 +836,37 @@ impl RegionMetadataBuilder { Ok(()) } + fn set_json_settings( + &mut self, + col_name: String, + settings: datatypes::json::JsonSettings, + ) -> Result<()> { + let Some(col_meta) = self + .column_metadatas + .iter_mut() + .find(|col| col.column_schema.name == col_name) + else { + return InvalidRegionRequestSnafu { + region_id: self.region_id, + err: format!("column {col_name} not found"), + } + .fail(); + }; + + let old_metadata = col_meta.column_schema.metadata(); + let new_metadata = + json2_metadata_with_updated_settings(old_metadata, settings).map_err(|err| { + InvalidRegionRequestSnafu { + region_id: self.region_id, + err: err.to_string(), + } + .build() + })?; + + *col_meta.column_schema.mut_metadata() = new_metadata; + Ok(()) + } + fn set_indexes(&mut self, options: Vec) -> Result<()> { let mut set_index_map: HashMap<_, Vec<_>> = HashMap::new(); for option in &options { diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 826ef7d6bc9..113a2393ccc 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -32,6 +32,7 @@ use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint, }; +use arrow_schema::extension::ExtensionType; pub use common_base::AffectedRows; use common_base::readable_size::ReadableSize; use common_grpc::flight::FlightDecoder; @@ -39,6 +40,8 @@ use common_recordbatch::DfRecordBatch; use common_time::range::TimestampRange; use common_time::{TimeToLive, Timestamp}; use datatypes::error::time_index_not_widening_error; +use datatypes::extension::json::Json2ExtensionType; +use datatypes::json::{JsonSettings, JsonTypeHint}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use num_enum::TryFromPrimitive; @@ -795,6 +798,13 @@ pub enum AlterKind { /// Columns to change. columns: Vec, }, + /// Set JSON2 settings of a region column. + SetJsonSettings { + /// Column name. + column_name: String, + /// Target JSON2 settings. + settings: JsonSettings, + }, /// Set region options. SetRegionOptions { options: Vec }, /// Unset region options. @@ -975,6 +985,9 @@ impl AlterKind { col_to_change.validate(metadata)?; } } + AlterKind::SetJsonSettings { column_name, .. } => { + Self::validate_set_json_settings(column_name, metadata)? + } AlterKind::SetRegionOptions { .. } => {} AlterKind::UnsetRegionOptions { .. } => {} AlterKind::SetIndexes { options } => { @@ -1084,6 +1097,18 @@ impl AlterKind { AlterKind::ModifyColumnTypes { columns } => columns .iter() .any(|col_to_change| col_to_change.need_alter(metadata)), + AlterKind::SetJsonSettings { + column_name, + settings, + } => metadata.column_by_name(column_name).is_some_and(|col| { + col.column_schema + .extension_type::() + .ok() + .flatten() + .is_none_or(|extension| { + !extension.metadata().json_settings().equivalent(settings) + }) + }), AlterKind::SetRegionOptions { .. } => true, AlterKind::UnsetRegionOptions { .. } => true, AlterKind::SetIndexes { options, .. } => options @@ -1160,6 +1185,34 @@ impl AlterKind { Ok(()) } + + fn validate_set_json_settings(col_name: &String, metadata: &RegionMetadata) -> Result<()> { + let region_id = metadata.region_id; + + let col = metadata + .column_by_name(col_name) + .with_context(|| InvalidRegionRequestSnafu { + region_id, + err: format!("column {} not found", col_name), + })?; + + ensure!( + col.semantic_type == SemanticType::Field, + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a field column", col_name), + } + ); + ensure!( + col.column_schema.data_type.is_json2(), + InvalidRegionRequestSnafu { + region_id, + err: format!("column {} is not a JSON2 column", col_name), + } + ); + + Ok(()) + } } impl TryFrom for AlterKind { @@ -1183,6 +1236,15 @@ impl TryFrom for AlterKind { .collect::>(); AlterKind::ModifyColumnTypes { columns } } + alter_request::Kind::SetJsonSettings(x) => { + let settings = x.settings.context(InvalidRawRegionRequestSnafu { + err: "missing settings in SetJsonSettings", + })?; + AlterKind::SetJsonSettings { + column_name: x.column_name, + settings: json_settings_from_proto(settings)?, + } + } alter_request::Kind::DropColumns(x) => { let names = x.drop_columns.into_iter().map(|x| x.name).collect(); AlterKind::DropColumns { names } @@ -1484,6 +1546,36 @@ impl From for ModifyColumnType { } } +fn json_settings_from_proto(settings: v1::JsonSettings) -> Result { + let type_hints = settings + .type_hints + .into_iter() + .map(|hint| { + let wrapper = ColumnDataTypeWrapper::try_new(hint.data_type, hint.datatype_extension) + .map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + })?; + let data_type = ConcreteDataType::from(wrapper); + + Ok(JsonTypeHint { + path: hint.path, + data_type, + inverted_index: false, + }) + }) + .collect::>>()?; + + JsonSettings::try_new(type_hints, settings.max_auto_expanded_paths).map_err(|err| { + InvalidRawRegionRequestSnafu { + err: err.to_string(), + } + .build() + }) +} + /// Region option changes used by ALTER requests. /// /// This type is serialized for request persistence. Keep future changes backward @@ -1953,8 +2045,11 @@ mod tests { use api::v1::region::RegionColumnDef; use api::v1::{ColumnDataType, ColumnDef}; use common_time::range::TimestampRange; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend}; + use datatypes::types::JsonType; use super::*; use crate::metadata::RegionMetadataBuilder; @@ -2364,6 +2459,34 @@ mod tests { }, } ); + + let request = RegionAlterRequest::try_from(AlterRequest { + region_id: 0, + schema_version: 1, + kind: Some(alter_request::Kind::SetJsonSettings(v1::SetJsonSettings { + column_name: "payload".to_string(), + settings: Some(v1::JsonSettings { + type_hints: vec![v1::JsonTypeHint { + path: vec!["service".to_string()], + data_type: ColumnDataType::String as i32, + datatype_extension: None, + }], + max_auto_expanded_paths: Some(10), + }), + })), + }) + .unwrap(); + + let AlterKind::SetJsonSettings { + column_name, + settings, + } = request.kind + else { + unreachable!() + }; + assert_eq!("payload", column_name); + assert_eq!(Some(10), settings.max_auto_expanded_paths()); + assert_eq!(1, settings.type_hints().len()); } #[test] @@ -2446,6 +2569,15 @@ mod tests { builder.build().unwrap() } + fn json2_column_schema(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(std::sync::Arc::new( + JsonMetadata::new(settings), + ))); + column_schema + } + #[test] fn test_add_column_validate() { let metadata = new_metadata(); @@ -2730,6 +2862,40 @@ mod tests { assert!(kind.need_alter(&metadata)); } + #[test] + fn test_validate_set_json_settings() { + let mut metadata = new_metadata(); + let current_schema = json2_column_schema("field_0", JsonSettings::new_v2()); + let target_settings = JsonSettings::try_new(vec![], Some(10)).unwrap(); + metadata + .column_metadatas + .iter_mut() + .find(|column| column.column_schema.name == "field_0") + .unwrap() + .column_schema = current_schema.clone(); + + let kind = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: target_settings.clone(), + }; + kind.validate(&metadata).unwrap(); + assert!(kind.need_alter(&metadata)); + + let no_op = AlterKind::SetJsonSettings { + column_name: "field_0".to_string(), + settings: JsonSettings::new_v2(), + }; + no_op.validate(&metadata).unwrap(); + assert!(!no_op.need_alter(&metadata)); + + AlterKind::SetJsonSettings { + column_name: "tag_0".to_string(), + settings: target_settings, + } + .validate(&metadata) + .unwrap_err(); + } + #[test] fn test_validate_add_columns() { let kind = AlterKind::AddColumns { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 5916dd7ab49..d91771b9c5f 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -22,6 +22,7 @@ use common_query::AddColumnLocation; use datafusion_expr::TableProviderFilterPushDown; use datatypes::error::time_index_not_widening_error; pub use datatypes::error::{Error as ConvertError, Result as ConvertResult}; +use datatypes::extension::json::json2_metadata_with_updated_settings; use datatypes::schema::{ ColumnSchema, FulltextOptions, Schema, SchemaBuilder, SchemaRef, SkippingIndexOptions, }; @@ -41,9 +42,9 @@ use crate::error::{self, Result}; use crate::requests::{ AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, REPARTITION_PARTITION_NUM_HINT_KEY, - SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, has_stable_string_form, - parse_entity_columns, parse_entity_option_key, validate_and_normalize_annotation, - validate_annotation_keys, + SetDefaultRequest, SetIndexOption, SetJsonSettingsRequest, TableOptions, UnsetIndexOption, + has_stable_string_form, parse_entity_columns, parse_entity_option_key, + validate_and_normalize_annotation, validate_annotation_keys, }; use crate::table_reference::TableReference; @@ -332,6 +333,7 @@ impl TableMeta { AlterKind::ModifyColumnTypes { columns } => { self.modify_column_types(table_name, columns) } + AlterKind::SetJsonSettings { request } => self.set_json_settings(table_name, request), // No need to rebuild table meta when renaming tables. AlterKind::RenameTable { .. } => Ok(self.new_meta_builder()), AlterKind::SetTableOptions { options } => self.set_table_options(options), @@ -1187,6 +1189,77 @@ impl TableMeta { Ok(meta_builder) } + fn set_json_settings( + &self, + table_name: &str, + req: &SetJsonSettingsRequest, + ) -> Result { + let table_schema = &self.schema; + let idx = table_schema + .column_index_by_name(&req.column_name) + .with_context(|| error::ColumnNotExistsSnafu { + column_name: &req.column_name, + table_name, + })?; + let col = &table_schema.column_schemas()[idx]; + + ensure!( + !self.primary_key_indices.contains(&idx) && table_schema.timestamp_index() != Some(idx), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "Not allowed to change JSON settings for key or timestamp column '{}'", + col.name + ), + } + ); + ensure!( + col.data_type.is_json2(), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!("column '{}' is not a JSON2 column", col.name), + } + ); + + let target_metadata = + json2_metadata_with_updated_settings(col.metadata(), req.settings.clone()).map_err( + |err| { + error::InvalidAlterRequestSnafu { + table: table_name, + err: err.to_string(), + } + .build() + }, + )?; + + let mut cols = table_schema.column_schemas().to_vec(); + *cols[idx].mut_metadata() = target_metadata; + + let mut builder = SchemaBuilder::try_from_columns(cols) + .with_context(|_| error::SchemaBuildSnafu { + msg: format!("Failed to convert column schemas into schema for table {table_name}"), + })? + .version(table_schema.version() + 1); + + for (k, v) in table_schema.metadata().iter() { + builder = builder.add_metadata(k, v); + } + + let new_schema = builder.build().with_context(|_| error::SchemaBuildSnafu { + msg: format!( + "Table {table_name} cannot change JSON settings for column {}", + req.column_name + ), + })?; + + let mut meta_builder = self.new_meta_builder(); + let _ = meta_builder + .schema(Arc::new(new_schema)) + .primary_key_indices(self.primary_key_indices.clone()); + + Ok(meta_builder) + } + /// Split requests into different groups using column location info. fn split_requests_by_column_location<'a>( &self, @@ -1603,9 +1676,12 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use datatypes::data_type::ConcreteDataType; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; use datatypes::schema::{ ColumnSchema, FulltextAnalyzer, FulltextBackend, Schema, SchemaBuilder, }; + use datatypes::types::JsonType; use super::*; use crate::Error; @@ -1700,6 +1776,90 @@ mod tests { builder.build().unwrap() } + #[test] + fn test_set_json_settings() { + let schema = Arc::new( + SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + json2_column_schema_v1("payload", JsonSettings::new_v2()), + ]) + .unwrap() + .version(123) + .build() + .unwrap(), + ); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "payload".to_string(), + settings: JsonSettings::try_new(vec![], Some(10)).unwrap(), + }, + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + + let payload = new_meta.schema.column_schema_by_name("payload").unwrap(); + let json_metadata: JsonMetadata = + serde_json::from_str(payload.metadata().get("ARROW:extension:metadata").unwrap()) + .unwrap(); + assert!(json_metadata.is_version_2()); + assert_eq!( + Some(10), + json_metadata.json_settings().max_auto_expanded_paths() + ); + assert_eq!(124, new_meta.schema.version()); + assert_eq!(meta.primary_key_indices, new_meta.primary_key_indices); + } + + fn json2_column_schema_v1(name: &str, settings: JsonSettings) -> ColumnSchema { + let mut column_schema = + ColumnSchema::new(name, ConcreteDataType::Json(JsonType::null()), true); + column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new( + JsonMetadata::new_v1(settings), + ))); + column_schema + } + + #[test] + fn test_set_json_settings_rejects_non_json2_column() { + let schema = Arc::new(new_test_schema()); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let alter_kind = AlterKind::SetJsonSettings { + request: SetJsonSettingsRequest { + column_name: "col2".to_string(), + settings: JsonSettings::new_v2(), + }, + }; + + let err = match meta.builder_with_alter_kind("my_table", &alter_kind) { + Ok(_) => panic!("expected modifying JSON settings on a non-JSON2 column to fail"), + Err(err) => err, + }; + assert!(err.to_string().contains("is not a JSON2 column")); + } + #[test] fn test_modify_time_index_column_type() { let schema = Arc::new(new_test_schema()); diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 1ca99268d67..2c519b32a5e 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -25,6 +25,7 @@ use common_query::AddColumnLocation; use common_time::TimeToLive; use common_time::range::TimestampRange; use datatypes::data_type::ConcreteDataType; +use datatypes::json::JsonSettings; use datatypes::prelude::VectorRef; use datatypes::schema::{ ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, @@ -371,6 +372,13 @@ pub struct ModifyColumnTypeRequest { pub target_type: ConcreteDataType, } +/// Set JSON2 settings request. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SetJsonSettingsRequest { + pub column_name: String, + pub settings: JsonSettings, +} + /// A family of annotation table options: pure metadata markers that no region /// consumes. Setting or unsetting them only rewrites the table's /// `extra_options`, so the alter skips region dispatch entirely. @@ -661,6 +669,9 @@ pub enum AlterKind { ModifyColumnTypes { columns: Vec, }, + SetJsonSettings { + request: SetJsonSettingsRequest, + }, RenameTable { new_table_name: String, }, diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index c23f9e2533f..ab9e39dc777 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -2,37 +2,110 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); Affected Rows: 0 -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type - -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2 (); - -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 2 ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration INT64, + trace_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + ++-----------------------+ +| alter_flushed_sst_num | ++-----------------------+ +| 1 | ++-----------------------+ + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +Affected Rows: 2 + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.service"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.duration"),Int64(NULL)) | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | frontend | 12 | before-1 | +| 1970-01-01T00:00:00.002 | ingester | 34 | before-2 | +| 1970-01-01T00:00:00.003 | frontend | 56 | after-1 | +| 1970-01-01T00:00:00.004 | compactor | | after-2 | ++-------------------------+-------------------------------------------------------------------+-----------------------------------------------------------------+--------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( @@ -40,10 +113,38 @@ ALTER TABLE application_logs trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); -Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type +Affected Rows: 0 + +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +Affected Rows: 1 + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| ts | json_get(application_logs.attrs,Utf8("$.trace_id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.id"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.user.name"),Utf8View(NULL)) | json_get(application_logs.attrs,Utf8("$.request_id"),Utf8View(NULL)) | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | before-1 | | | | +| 1970-01-01T00:00:00.002 | before-2 | | | | +| 1970-01-01T00:00:00.003 | after-1 | | | | +| 1970-01-01T00:00:00.004 | after-2 | | | | +| 1970-01-01T00:00:00.005 | after-3 | u1 | alice | r1 | +| 1970-01-01T00:00:00.006 | after-4 | u2 | | r2 | ++-------------------------+--------------------------------------------------------------------+-------------------------------------------------------------------+---------------------------------------------------------------------+----------------------------------------------------------------------+ + +ADMIN FLUSH_TABLE('application_logs'); + ++---------------------------------------+ +| ADMIN FLUSH_TABLE('application_logs') | ++---------------------------------------+ +| 0 | ++---------------------------------------+ DROP TABLE application_logs; diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql index c89556c3485..7fce02b76ee 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -2,35 +2,75 @@ CREATE TABLE application_logs ( ts TIMESTAMP TIME INDEX, attrs JSON2 ) WITH ( - 'append_mode' = 'true' + 'append_mode' = 'true', + 'sst_format' = 'flat' ); -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2; - -ALTER TABLE application_logs - MODIFY COLUMN attrs JSON2 (); +INSERT INTO application_logs VALUES + (1, '{"service":"frontend","duration":12,"trace_id":"before-1"}'), + (2, '{"service":"ingester","duration":34,"trace_id":"before-2"}'); ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( - max_auto_expanded_paths = 2000 + max_auto_expanded_paths = 2000, + service STRING, + duration INT64, + trace_id STRING ); +SELECT COUNT(*) AS alter_flushed_sst_num +FROM information_schema.ssts_manifest +WHERE visible + AND table_id IN ( + SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'application_logs' + ); + +INSERT INTO application_logs VALUES + (3, '{"service":"frontend","duration":56,"trace_id":"after-1"}'), + (4, '{"service":"compactor","trace_id":"after-2"}'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + +SELECT ts, attrs.service, attrs.duration, attrs.trace_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (5, '{"trace_id":"after-3","user":{"id":"u1","name":"alice"},"request_id":"r1"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, trace_id STRING, user.id STRING, user.name STRING, - request_id STRING INVERTED INDEX + request_id STRING ); +INSERT INTO application_logs VALUES + (6, '{"trace_id":"after-4","user":{"id":"u2"},"request_id":"r2"}'); + +SELECT ts, attrs.trace_id, attrs.user.id, attrs.user.name, attrs.request_id +FROM application_logs +ORDER BY ts; + +ADMIN FLUSH_TABLE('application_logs'); + DROP TABLE application_logs; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result new file mode 100644 index 00000000000..284b7e6f850 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result @@ -0,0 +1,83 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + ++--------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + ++---------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400') | ++---------------------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_alter_type_hints_compaction.attrs,Utf8("$.message")) | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | before-has | +| 1970-01-01T00:00:00.002 | | before-missing | +| 1970-01-01T00:00:00.003 | 3 | after-has | +| 1970-01-01T00:00:00.004 | | after-missing | ++-------------------------+------------------------------------------------------------------------------+---------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql new file mode 100644 index 00000000000..45d859d4ece --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql @@ -0,0 +1,36 @@ +CREATE TABLE json2_alter_type_hints_compaction ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (1, '{"kind":1,"message":"before-has"}'), + (2, '{"message":"before-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +ALTER TABLE json2_alter_type_hints_compaction + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +INSERT INTO json2_alter_type_hints_compaction VALUES + (3, '{"kind":3,"message":"after-has"}'), + (4, '{"message":"after-missing"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_alter_type_hints_compaction +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction; diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result new file mode 100644 index 00000000000..5ba2cf4700a --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result @@ -0,0 +1,78 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + ++-------------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion') | ++-------------------------------------------------------------------+ +| 0 | ++-------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + ++--------------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400') | ++--------------------------------------------------------------------------------------+ +| 0 | ++--------------------------------------------------------------------------------------+ + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| ts | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_string"),Utf8View(NULL)) | json_get(json2_alter_type_hints_compaction_conversion.attrs,Utf8("$.to_int"),Int64(NULL)) | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 123 | 456 | +| 1970-01-01T00:00:00.002 | after | 789 | ++-------------------------+-------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------+ + +DROP TABLE json2_alter_type_hints_compaction_conversion; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql new file mode 100644 index 00000000000..de9f211cfb5 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_alter_type_hints_compaction_conversion ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (1, '{"to_string":123,"to_int":"456"}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +ALTER TABLE json2_alter_type_hints_compaction_conversion + MODIFY COLUMN attrs JSON2 ( + to_string STRING, + to_int INT64 + ); + +INSERT INTO json2_alter_type_hints_compaction_conversion VALUES + (2, '{"to_string":"after","to_int":789}'); + +ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_alter_type_hints_compaction_conversion', 'swcs', '86400'); + +SELECT ts, attrs.to_string, attrs.to_int +FROM json2_alter_type_hints_compaction_conversion +ORDER BY ts; + +DROP TABLE json2_alter_type_hints_compaction_conversion; diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result new file mode 100644 index 00000000000..6e0f833ccd6 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result @@ -0,0 +1,80 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +Affected Rows: 2 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +Affected Rows: 0 + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + ++---------------------------------------------------------------+ +| ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch') | ++---------------------------------------------------------------+ +| 0 | ++---------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + ++----------------------------------------------------------------------------------+ +| ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400') | ++----------------------------------------------------------------------------------+ +| 0 | ++----------------------------------------------------------------------------------+ + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| ts | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.kind"),Int64(NULL)) | json_get(json2_compaction_null_type_hint_mismatch.attrs,Utf8("$.message")) | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00.001 | 1 | | +| 1970-01-01T00:00:00.002 | | keep | +| 1970-01-01T00:00:00.003 | 3 | | ++-------------------------+-------------------------------------------------------------------------------------+----------------------------------------------------------------------------+ + +DROP TABLE json2_compaction_null_type_hint_mismatch; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql new file mode 100644 index 00000000000..7ebff9cceb1 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql @@ -0,0 +1,35 @@ +CREATE TABLE json2_compaction_null_type_hint_mismatch ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true', + 'sst_format' = 'flat' +); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (1, '{"kind":1}'), + (2, '{"kind":"invalid","message":"keep"}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +ALTER TABLE json2_compaction_null_type_hint_mismatch + MODIFY COLUMN attrs JSON2 ( + kind INT64 + ); + +INSERT INTO json2_compaction_null_type_hint_mismatch VALUES + (3, '{"kind":3}'); + +ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +ADMIN COMPACT_TABLE('json2_compaction_null_type_hint_mismatch', 'swcs', '86400'); + +SELECT ts, attrs.kind, attrs.message +FROM json2_compaction_null_type_hint_mismatch +ORDER BY ts; + +DROP TABLE json2_compaction_null_type_hint_mismatch; From a9e2a89b7b61877875e522dc32bc134eaf08ed39 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:12:39 +0000 Subject: [PATCH 007/133] fix: keep compaction pruning, metadata, and index work on compact runtime (#9304) * fix: run compaction pruner tasks on compact runtime Signed-off-by: Lei, HUANG * fix: keep compaction metadata and index work on compact runtime Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito2/src/access_layer.rs | 8 +- src/mito2/src/cache.rs | 210 +++++++++++++++++++--------- src/mito2/src/cache/file_cache.rs | 7 +- src/mito2/src/read/pruner.rs | 12 +- src/mito2/src/region/opener.rs | 8 +- src/mito2/src/sst/index.rs | 20 ++- src/mito2/src/sst/parquet/reader.rs | 9 +- 7 files changed, 193 insertions(+), 81 deletions(-) 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/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..ed3ec36a473 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -57,8 +57,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, @@ -774,7 +774,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/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) => { From d0f8f4b80c25c3f379acf52bd51399313372ba20 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:12:24 +0000 Subject: [PATCH 008/133] feat: add AI matching, classification, and scoring functions (#9300) * feat: return matching scores from jev Replace the experimental three-argument Boolean function with jev(text, prompt) returning a Float64 probability in [0, 1]. Move threshold comparisons into SQL and update tests and migration examples. Signed-off-by: Lei, HUANG * feat: add Jev choice and score functions Share asynchronous execution across Noul, Choice, and Score. Validate JSON criteria before requests and return typed scalar answers. Add SQL and HTTP mock coverage with usage examples. Signed-off-by: Lei, HUANG * refactor: use generic AI SQL function names Expose ai_match, ai_choose, and ai_score and move their implementation, tests, and usage guide under generic AI names. Document the current unreleased interface without migration history. Signed-off-by: Lei, HUANG * fix: share constant AI criteria within each batch Borrow scalar string arguments and lazily parse constant criteria once per batch. Share the parsed allocation across requests while preserving NULL propagation and batch validation before HTTP calls. Signed-off-by: Lei, HUANG * feat: preserve AI score uncertainty in JSONB results Return score, confidence, and probabilities in criteria-level order from one evaluation. Validate the distribution and preserve provider precision. Add JSON extraction, uncertainty, and single-request regressions, and document confidence-aware ranking. Signed-off-by: Lei, HUANG * docs: explain reuse of volatile AI evaluations Document repeated SELECT and WHERE evaluation costs as N + M requests, and show subquery aliases for reusing scalar or structured AI results without additional model calls. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- docs/how-to/ai-functions.md | 247 ++++++ docs/how-to/jev.md | 119 --- src/common/function/src/function_registry.rs | 22 +- src/common/function/src/scalars.rs | 4 +- src/common/function/src/scalars/ai.rs | 454 +++++++++++ src/common/function/src/scalars/ai/tests.rs | 766 ++++++++++++++++++ src/common/function/src/scalars/jev.rs | 209 ----- src/common/function/src/scalars/jev/tests.rs | 250 ------ .../common/function/ai_functions.result | 144 ++++ .../common/function/ai_functions.sql | 65 ++ .../standalone/common/function/jev.result | 66 -- .../cases/standalone/common/function/jev.sql | 32 - 12 files changed, 1692 insertions(+), 686 deletions(-) create mode 100644 docs/how-to/ai-functions.md delete mode 100644 docs/how-to/jev.md create mode 100644 src/common/function/src/scalars/ai.rs create mode 100644 src/common/function/src/scalars/ai/tests.rs delete mode 100644 src/common/function/src/scalars/jev.rs delete mode 100644 src/common/function/src/scalars/jev/tests.rs create mode 100644 tests/cases/standalone/common/function/ai_functions.result create mode 100644 tests/cases/standalone/common/function/ai_functions.sql delete mode 100644 tests/cases/standalone/common/function/jev.result delete mode 100644 tests/cases/standalone/common/function/jev.sql 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/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/tests/cases/standalone/common/function/ai_functions.result b/tests/cases/standalone/common/function/ai_functions.result new file mode 100644 index 00000000000..057c9ecd990 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.result @@ -0,0 +1,144 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT ai_match('message', NULL) AS score; + ++-------+ +| score | ++-------+ +| | ++-------+ + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + ++------------+ +| score_type | ++------------+ +| Float64 | ++------------+ + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + ++-----------+-------------+---------------+ +| null_text | null_prompt | null_criteria | ++-----------+-------------+---------------+ +| | | | ++-----------+-------------+---------------+ + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + ++-------------+------------+ +| choice_type | score_type | ++-------------+------------+ +| Utf8 | BinaryView | ++-------------+------------+ + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + ++--------+-------+------------+------------------+ +| rating | score | confidence | high_probability | ++--------+-------+------------+------------------+ +| | | | | ++--------+-------+------------+------------------+ + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +Affected Rows: 0 + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +Affected Rows: 2 + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + ++-------------+---------+---------+ +| occurred_at | service | message | ++-------------+---------+---------+ ++-------------+---------+---------+ + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + ++---------------------+------+--------+ +| occurred_at | team | rating | ++---------------------+------+--------+ +| 2026-09-19T01:00:00 | | | +| 2026-09-19T02:00:00 | | | ++---------------------+------+--------+ + +DROP TABLE ai_events; + +Affected Rows: 0 + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_match' function: coercion from Utf8, Utf8, Float64 to the signature Exact(Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_match(Utf8, Utf8, Float64)'. You might need to add explicit type casts. + Candidate functions: + ai_match(Utf8View, Utf8View) + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria is not valid JSON: expected ident at line 1 column 2 + +SELECT ai_choose('message', 'prompt', '{}'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_choose criteria must be a JSON object with 1 to 255 options; descriptions must be strings, objects, arrays, or null + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +Error: 3001(EngineExecuteQuery), Execution error: ai_score criteria must be a JSON array with 2 to 10 levels; descriptions must be strings, objects, or arrays + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_choose' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_choose(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_choose(Utf8View, Utf8View, Utf8View) + +SELECT ai_score('message', 'prompt'); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Failed to coerce arguments to satisfy a call to 'ai_score' function: coercion from Utf8, Utf8 to the signature Exact(Utf8View, Utf8View, Utf8View) failed. No function matches the given name and argument types 'ai_score(Utf8, Utf8)'. You might need to add explicit type casts. + Candidate functions: + ai_score(Utf8View, Utf8View, Utf8View) + diff --git a/tests/cases/standalone/common/function/ai_functions.sql b/tests/cases/standalone/common/function/ai_functions.sql new file mode 100644 index 00000000000..fd3e80d1e18 --- /dev/null +++ b/tests/cases/standalone/common/function/ai_functions.sql @@ -0,0 +1,65 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT ai_match(NULL, 'The event reports a failed payment.') AS score; + +SELECT ai_match('message', NULL) AS score; + +SELECT arrow_typeof(ai_match(NULL, 'condition')) AS score_type; + +SELECT ai_choose(NULL, 'Route the ticket', '{"billing":null}') AS null_text, + ai_choose('message', NULL, 'invalid JSON') AS null_prompt, + ai_choose('message', 'Route the ticket', NULL) AS null_criteria; + +SELECT ai_score(NULL, 'Rate severity', '["low","high"]') AS null_text, + ai_score('message', NULL, 'invalid JSON') AS null_prompt, + ai_score('message', 'Rate severity', NULL) AS null_criteria; + +SELECT arrow_typeof(ai_choose(NULL, 'prompt', '{"billing":null}')) AS choice_type, + arrow_typeof(ai_score(NULL, 'prompt', '["low","high"]')) AS score_type; + +SELECT json_to_string(rating) AS rating, + json_get_float(rating, 'score') AS score, + json_get_float(rating, 'confidence') AS confidence, + json_get_float(rating, 'probabilities[2]') AS high_probability +FROM (SELECT ai_score(NULL, 'prompt', '["low","medium","high"]') AS rating) AS rated; + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE ai_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +INSERT INTO ai_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +SELECT occurred_at, service, message FROM ai_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND ai_match((message), 'The event reports that a payment still failed after retries.') >= 0.8 +ORDER BY occurred_at; + +SELECT occurred_at, + ai_choose(message, 'Route the ticket', '{"billing":"Payments","technical":"Errors"}') AS team, + ai_score(message, 'Rate severity', '["low","medium","high"]') AS rating +FROM ai_events +ORDER BY occurred_at; + +DROP TABLE ai_events; + +-- Matching uses two arguments; thresholds are expressed as SQL comparisons. +SELECT ai_match('message', 'condition', 0.8); + +-- Invalid criteria fail locally before any HTTP request. +SELECT ai_choose('message', 'prompt', 'not json'); + +SELECT ai_choose('message', 'prompt', '{}'); + +SELECT ai_score('message', 'prompt', '["only one level"]'); + +-- Both additional modes require criteria. +SELECT ai_choose('message', 'prompt'); + +SELECT ai_score('message', 'prompt'); diff --git a/tests/cases/standalone/common/function/jev.result b/tests/cases/standalone/common/function/jev.result deleted file mode 100644 index 4f6d91d2867..00000000000 --- a/tests/cases/standalone/common/function/jev.result +++ /dev/null @@ -1,66 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', NULL, 0.8) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - -+---------+ -| matched | -+---------+ -| | -+---------+ - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -Affected Rows: 0 - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -Affected Rows: 2 - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -+-------------+---------+---------+ -| occurred_at | service | message | -+-------------+---------+---------+ -+-------------+---------+---------+ - -DROP TABLE jev_events; - -Affected Rows: 0 - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - -SELECT jev('message', 'condition', 1.1); - -Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] - diff --git a/tests/cases/standalone/common/function/jev.sql b/tests/cases/standalone/common/function/jev.sql deleted file mode 100644 index f4c252c9273..00000000000 --- a/tests/cases/standalone/common/function/jev.sql +++ /dev/null @@ -1,32 +0,0 @@ --- No credentials or external service needed: NULL propagates without an API call. -SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; - -SELECT jev('message', NULL, 0.8) AS matched; - -SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; - --- Validate SQL registration, coercion, and asynchronous filtering on a real table. -CREATE TABLE jev_events ( - occurred_at TIMESTAMP TIME INDEX, - "service" STRING, - "message" STRING, - PRIMARY KEY ("service") -); - -INSERT INTO jev_events VALUES - ('2026-09-19T01:00:00Z', 'payments', NULL), - ('2026-09-19T02:00:00Z', 'auth', NULL); - -SELECT occurred_at, service, message FROM jev_events -WHERE occurred_at >= '2026-09-19T00:00:00Z' - AND occurred_at < '2026-09-20T00:00:00Z' - AND service = 'payments' - AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) -ORDER BY occurred_at; - -DROP TABLE jev_events; - --- Invalid thresholds fail locally before any HTTP request. -SELECT jev('message', 'condition', -0.1); - -SELECT jev('message', 'condition', 1.1); From 723da69b2182b747a82574779db32e33ffb63111 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 22 Sep 2026 14:38:17 +0000 Subject: [PATCH 009/133] feat: share logical table batching with OTLP metrics (#9288) * feat: share logical table batching with OTLP metrics Signed-off-by: WenyXu * fix: unify pending rows batch acknowledgement policy Signed-off-by: WenyXu * fix: align logical batcher example configuration expectations Signed-off-by: WenyXu * fix: align batcher worker channel defaults to 65536 Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- config/config.md | 38 ++- config/frontend.example.toml | 52 +-- config/standalone.example.toml | 52 +-- src/cmd/src/bin/query_perf_fixture/case.rs | 2 +- src/cmd/tests/load_config_test.rs | 9 + src/frontend/AGENTS.md | 5 + src/frontend/src/frontend.rs | 49 ++- src/frontend/src/instance.rs | 5 +- src/frontend/src/instance/builder.rs | 48 +-- src/frontend/src/instance/logical_batcher.rs | 102 ++++++ src/frontend/src/instance/otlp.rs | 57 +++- src/frontend/src/server.rs | 171 +++++++--- src/frontend/src/service_config.rs | 2 +- .../service_config/pending_rows_batcher.rs | 212 ++++++++++-- src/frontend/src/service_config/prom_store.rs | 2 +- src/operator/src/batcher.rs | 3 +- src/operator/src/insert.rs | 315 +++++++++++++++++- src/servers/src/batcher.rs | 18 + src/servers/src/batcher/logical_table.rs | 45 +-- .../src/batcher/logical_table/tables.rs | 44 +-- src/servers/src/batcher/table.rs | 117 ++++++- src/servers/src/batcher/table/batch.rs | 2 + src/servers/src/http.rs | 63 +++- src/session/src/context.rs | 24 +- src/standalone/src/options.rs | 73 +++- tests-integration/src/otlp.rs | 227 +++++++++++++ tests-integration/src/standalone.rs | 14 + tests-integration/tests/http.rs | 4 +- 28 files changed, 1518 insertions(+), 237 deletions(-) create mode 100644 src/frontend/src/instance/logical_batcher.rs 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/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/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..a3cfbcdaffe 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(), @@ -386,6 +385,7 @@ impl FrontendBuilder { 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/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/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/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> + Send, + { let (table_batches, total_rows) = { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_build_and_align"]) .start_timer(); - self.build_and_align_table_batches(requests, &ctx).await? + self.build_and_align_table_batches(&requests, &ctx).await? }; + let prepared = after_prepare(requests).await?; if total_rows == 0 { - return Ok(0); + return Ok((0, prepared)); } // Flushes dispatch directly to datanodes, so admit once before enqueueing. @@ -269,9 +270,9 @@ impl LogicalTablePendingRowsBatcher { }; result .context(error::SubmitBatchSnafu) - .map(|()| total_rows as u64) + .map(|()| (total_rows as u64, prepared)) } else { - Ok(total_rows as u64) + Ok((total_rows as u64, prepared)) } } } diff --git a/src/servers/src/batcher/logical_table/tables.rs b/src/servers/src/batcher/logical_table/tables.rs index 4ed07f73a6e..87110366a8c 100644 --- a/src/servers/src/batcher/logical_table/tables.rs +++ b/src/servers/src/batcher/logical_table/tables.rs @@ -75,7 +75,7 @@ impl LogicalTablePendingRowsBatcher { /// RecordBatches. pub(in crate::batcher::logical_table) async fn build_and_align_table_batches( &self, - requests: RowInsertRequests, + requests: &RowInsertRequests, ctx: &QueryContextRef, ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { let catalog = ctx.current_catalog().to_string(); @@ -113,13 +113,13 @@ impl LogicalTablePendingRowsBatcher { /// Extracts non-empty `(table_name, rows)` pairs and computes total row /// count across the retained entries. pub(in crate::batcher::logical_table) fn collect_non_empty_table_rows( - requests: RowInsertRequests, - ) -> (Vec<(String, Rows)>, usize) { - let mut table_rows: Vec<(String, Rows)> = Vec::with_capacity(requests.inserts.len()); + requests: &RowInsertRequests, + ) -> (Vec<(&str, &Rows)>, usize) { + let mut table_rows: Vec<(&str, &Rows)> = Vec::with_capacity(requests.inserts.len()); let mut total_rows = 0; - for request in requests.inserts { - let Some(rows) = request.rows else { + for request in &requests.inserts { + let Some(rows) = &request.rows else { continue; }; if rows.rows.is_empty() { @@ -127,7 +127,7 @@ impl LogicalTablePendingRowsBatcher { } total_rows += rows.rows.len(); - table_rows.push((request.table_name, rows)); + table_rows.push((request.table_name.as_str(), rows)); } (table_rows, total_rows) @@ -137,15 +137,15 @@ impl LogicalTablePendingRowsBatcher { impl LogicalTablePendingRowsBatcher { /// Returns unique `(table_name, proto_schema)` pairs while keeping the /// first-seen schema for duplicate table names. - pub(in crate::batcher::logical_table) fn collect_unique_table_schemas( - table_rows: &[(String, Rows)], - ) -> Result> { + pub(in crate::batcher::logical_table) fn collect_unique_table_schemas<'a>( + table_rows: &[(&'a str, &'a Rows)], + ) -> Result> { let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); let mut seen = HashSet::new(); for (table_name, rows) in table_rows { - if seen.insert(table_name.as_str()) { - unique_tables.push((table_name.as_str(), &rows.schema)); + if seen.insert(*table_name) { + unique_tables.push((*table_name, &rows.schema)); } else { // table_rows should group rows by table name. return error::InvalidPromRemoteRequestSnafu { @@ -233,7 +233,7 @@ impl LogicalTablePendingRowsBatcher { catalog: &str, schema: &str, ctx: &QueryContextRef, - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { if plan.tables_to_create.is_empty() { @@ -298,7 +298,7 @@ impl LogicalTablePendingRowsBatcher { /// For newly created tables, re-checks all row schemas and appends alter /// operations when additional tag columns are still missing. pub(in crate::batcher::logical_table) fn enqueue_alter_for_new_tables( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], plan: &mut TableResolutionPlan, ) -> Result<()> { let created_tables: HashSet<&str> = plan @@ -308,11 +308,11 @@ impl LogicalTablePendingRowsBatcher { .collect(); for (table_name, rows) in table_rows { - if !created_tables.contains(table_name.as_str()) { + if !created_tables.contains(table_name) { continue; } - let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { + let Some((region_schema, _)) = plan.region_schemas.get(*table_name) else { continue; }; @@ -321,13 +321,13 @@ impl LogicalTablePendingRowsBatcher { || plan .tables_to_alter .iter() - .any(|(existing_name, _)| existing_name == table_name) + .any(|(existing_name, _)| existing_name == *table_name) { continue; } plan.tables_to_alter - .push((table_name.clone(), missing_columns)); + .push((table_name.to_string(), missing_columns)); } Ok(()) @@ -397,13 +397,13 @@ impl LogicalTablePendingRowsBatcher { /// Converts proto rows to `RecordBatch` values aligned to resolved region /// schemas and returns `(table_name, table_id, batch)` tuples. pub(in crate::batcher::logical_table) fn build_aligned_batches( - table_rows: &[(String, Rows)], + table_rows: &[(&str, &Rows)], region_schemas: &HashMap, u32)>, ) -> Result> { let mut aligned_batches = Vec::with_capacity(table_rows.len()); for (table_name, rows) in table_rows { let (region_schema, table_id) = - region_schemas.get(table_name).cloned().with_context(|| { + region_schemas.get(*table_name).cloned().with_context(|| { error::UnexpectedResultSnafu { reason: format!("Region schema not resolved for table: {}", table_name), } @@ -415,7 +415,7 @@ impl LogicalTablePendingRowsBatcher { .start_timer(); rows_to_aligned_record_batch(rows, region_schema.as_ref())? }; - aligned_batches.push((table_name.clone(), table_id, record_batch)); + aligned_batches.push((table_name.to_string(), table_id, record_batch)); } Ok(aligned_batches) @@ -531,7 +531,7 @@ mod tests { }; let (table_rows, total_rows) = - LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); + LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(&requests); assert_eq!(2, total_rows); assert_eq!(1, table_rows.len()); diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 1a46a7f10f3..67b3431a810 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -38,6 +38,7 @@ use snafu::ResultExt; use table::metadata::TableInfoRef; use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot}; +use crate::batcher::pending_rows_batch_sync_enabled; use crate::batcher::table::flow_notifier::FlowNotifier; use crate::batcher::table::metrics::PENDING_WORKERS; use crate::batcher::table::pending_batch::PendingBatch; @@ -70,6 +71,7 @@ pub struct TablePendingRowsBatcher { flush_limiter: FlushLimiter, request_limiter: RequestLimiter, worker_channel_capacity: usize, + pending_rows_batch_sync: bool, worker_idle_timeout: Duration, inserter: Arc, flow_notifier: FlowNotifier, @@ -110,6 +112,7 @@ impl TablePendingRowsBatcher { flush_limiter, request_limiter, worker_channel_capacity, + pending_rows_batch_sync: pending_rows_batch_sync_enabled(), worker_idle_timeout: flush_interval.checked_mul(3).unwrap_or(flush_interval), inserter, flow_notifier, @@ -160,8 +163,8 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { }) } - /// Waits for completed bulk writes. Cancellation does not retract an admitted - /// submission. Combined failures affect all waiters and may be partial writes. + /// Acknowledges according to the global batching policy. Cancellation does + /// not retract an admitted submission. Flush failures may be partial writes. async fn submit( &self, table_info: TableInfoRef, @@ -210,6 +213,9 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } .fail(); } + if !self.pending_rows_batch_sync { + return Ok(total_rows); + } response_rx .await .map_err(|_| { @@ -236,7 +242,7 @@ mod tests { use api::v1::region::{RegionRequest, bulk_insert_request, region_request}; use api::v1::value::ValueData; use api::v1::{ColumnDataType, Row, Rows, Value}; - use arrow::array::{Int32Array, TimestampMillisecondArray}; + use arrow::array::{ArrayRef, Int32Array, TimestampMillisecondArray}; use arrow::datatypes::Schema as ArrowSchema; use arrow::record_batch::RecordBatch; use catalog::memory::MemoryCatalogManager; @@ -320,6 +326,111 @@ mod tests { } } + #[tokio::test] + async fn test_acknowledgement_preserves_admission() { + use operator::error::UnexpectedSnafu; + + use crate::batcher::pending_rows_batch_sync_enabled; + use crate::batcher::table::batch_key_from_ctx; + use crate::batcher::table::pending_batch::notify_batches; + + for sync in [false, true] { + for fail in [false, true] { + let backend = prepare_mocked_backend().await; + let nodes = Arc::new(MockDatanodeManager::new(BulkHandler { + requests: Arc::new(Mutex::new(Vec::new())), + report_missing_row: false, + expected_skip_wal: false, + expected_schema: "public".to_string(), + })); + let inserter = Arc::new(Inserter::new( + MemoryCatalogManager::new(), + create_partition_rule_manager(backend).await, + nodes, + mock_table_flownode_cache(1, vec![]).await, + true, + )); + let mut batcher = TablePendingRowsBatcher::try_new( + Duration::from_secs(3600), + 1, + 1, + 1, + 1, + NonZeroUsize::new(1).unwrap(), + inserter, + ) + .unwrap(); + assert_eq!( + batcher.pending_rows_batch_sync, + pending_rows_batch_sync_enabled() + ); + Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync; + let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter())); + let ctx = QueryContext::arc(); + let key = batch_key_from_ctx(&table.name, &ctx); + // Hold the worker command to control completion independently of scheduling. + let (_, receiver) = batcher.workers.get_or_create(key, 1).await; + let mut receiver = receiver.unwrap(); + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]) + .unwrap(); + let permit = batcher.acquire().await.unwrap(); + let submitter = batcher.clone(); + let submitted = + tokio::spawn(async move { submitter.submit(table, batch, ctx, permit).await }); + let WorkerCommand::Submit(pending) = + timeout(Duration::from_secs(5), receiver.recv()) + .await + .unwrap() + .unwrap(); + let mut submitted = Some(submitted); + if sync { + assert!(!submitted.as_ref().unwrap().is_finished()); + } else { + assert_eq!( + timeout(Duration::from_secs(5), submitted.take().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(), + 1 + ); + } + // Early acknowledgement must not release capacity before completion. + let acquire = batcher.acquire(); + tokio::pin!(acquire); + assert!(futures::poll!(acquire.as_mut()).is_pending()); + let result = if fail { + Err(Arc::new( + UnexpectedSnafu { + violated: "flush failed".to_string(), + } + .build(), + )) + } else { + Ok(()) + }; + notify_batches(vec![pending], result); + if let Some(submitted) = submitted { + let result = timeout(Duration::from_secs(5), submitted) + .await + .unwrap() + .unwrap(); + assert_eq!(result.is_err(), fail); + if !fail { + assert_eq!(result.unwrap(), 1); + } + } + timeout(Duration::from_secs(5), acquire) + .await + .unwrap() + .unwrap(); + } + } + } + fn rows(value: i32) -> Rows { Rows { schema: vec![ diff --git a/src/servers/src/batcher/table/batch.rs b/src/servers/src/batcher/table/batch.rs index 21ebb94a33a..e9ab339a6ac 100644 --- a/src/servers/src/batcher/table/batch.rs +++ b/src/servers/src/batcher/table/batch.rs @@ -17,6 +17,7 @@ use std::time::Instant; use arrow::compute::concat_batches; use arrow::record_batch::RecordBatch; +use common_telemetry::error; use operator::error::{ComputeArrowSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; use operator::metrics::DIST_INGEST_ROW_COUNT; @@ -55,6 +56,7 @@ pub(in crate::batcher::table) async fn flush_batch( notifier.notify(table, &combined); } Err(error) => { + error!(error; "Failed to flush table batch, rows: {}", batch.total_rows); FLUSH_FAILURES.inc(); FLUSH_DROPPED_ROWS.inc_by(batch.total_rows as u64); notify_batches(batch.submissions, Err(Arc::new(error))); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 7f2688d30cf..4d168111ade 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -215,6 +215,7 @@ pub struct HttpServer { // server configs options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, bind_addr: Option, /// What this server instance exposes. See [`HttpServerKind`]. kind: HttpServerKind, @@ -235,6 +236,9 @@ pub fn is_api_listener_path(path: &str) -> bool { is_namespace(path, HTTP_API_PREFIX_WITHOUT_TRAILING_SLASH) || is_namespace(path, "/dashboard") } +#[derive(Clone)] +struct LogicalBatchingProtocols(Vec); + /// Sets a local-only write selector after authentication creates the context. async fn set_http_write_batching( State(protocol): State, @@ -245,8 +249,13 @@ async fn set_http_write_batching( .extensions() .get::>>() .is_some_and(|protocols| protocols.contains(&protocol)); + let logical_enabled = req + .extensions() + .get::() + .is_some_and(|protocols| protocols.0.contains(&protocol)); if let Some(ctx) = req.extensions_mut().get_mut::() { ctx.set_batching_enabled(enabled); + ctx.set_logical_batching_enabled(logical_enabled); } next.run(req).await } @@ -650,6 +659,7 @@ pub struct DashboardState { pub struct HttpServerBuilder { options: HttpOptions, batching_protocols: Vec, + logical_batching_protocols: Vec, user_provider: Option, router: Router, memory_limiter: ServerMemoryLimiter, @@ -660,13 +670,20 @@ impl HttpServerBuilder { Self { options, batching_protocols: Vec::new(), + logical_batching_protocols: Vec::new(), user_provider: None, router: Router::new(), memory_limiter: ServerMemoryLimiter::default(), } } - /// Selects HTTP write protocols allowed to use the shared batcher. + /// Selects HTTP protocols allowed to use logical-table batching. + pub fn with_logical_batching_protocols(mut self, protocols: Vec) -> Self { + self.logical_batching_protocols = protocols; + self + } + + /// Selects HTTP protocols allowed to use ordinary-table batching. pub fn with_batching_protocols(mut self, protocols: Vec) -> Self { self.batching_protocols = protocols; self @@ -894,6 +911,7 @@ impl HttpServerBuilder { HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider, shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -928,6 +946,7 @@ impl HttpServerBuilder { let internal = HttpServer { options: self.options, batching_protocols: self.batching_protocols.clone(), + logical_batching_protocols: self.logical_batching_protocols.clone(), user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router.clone()), @@ -946,6 +965,7 @@ impl HttpServerBuilder { Some(HttpServer { options: api_options, batching_protocols: self.batching_protocols, + logical_batching_protocols: self.logical_batching_protocols, user_provider: self.user_provider.clone(), shutdown_tx: Mutex::new(None), router: StdMutex::new(self.router), @@ -1104,6 +1124,9 @@ impl HttpServer { authorize::check_http_auth, )) .layer(Extension(Arc::new(self.batching_protocols.clone()))) + .layer(Extension(LogicalBatchingProtocols( + self.logical_batching_protocols.clone(), + ))) .layer(middleware::from_fn(hints::extract_hints)) .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( @@ -2462,6 +2485,44 @@ mod batching_tests { use crate::opentsdb::codec::DataPoint; use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler}; + #[tokio::test] + async fn test_batching_selectors_are_independent() { + use axum::routing::post; + use axum::{Extension, Json, Router, middleware}; + use session::context::{QueryContext, QueryContextBuilder}; + + use crate::http::{LogicalBatchingProtocols, set_http_write_batching}; + for table in [false, true] { + for logical in [false, true] { + let app = Router::new() + .route( + "/", + post(|Extension(ctx): Extension| async move { + Json([ctx.batching_enabled(), ctx.logical_batching_enabled()]) + }), + ) + .route_layer(middleware::from_fn_with_state( + BatchingProtocol::Otlp, + set_http_write_batching, + )) + .layer(Extension(QueryContextBuilder::default().build())) + .layer(Extension(Arc::new(if table { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))) + .layer(Extension(LogicalBatchingProtocols(if logical { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }))); + let client = TestClient::new(app).await; + let actual: [bool; 2] = client.post("/").send().await.json().await; + assert_eq!(actual, [table, logical]); + } + } + } + #[test] fn test_protocol_names_reject_unknown_values() { assert_eq!( diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 0bb3f0f6a31..5ebf8b36d21 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -82,6 +82,9 @@ pub struct QueryContext { /// Local-only write batching selection; never transported in protobuf extensions. #[builder(default)] batching_enabled: bool, + /// Local-only opt-in for logical metric writes, independent of ordinary-table batching. + #[builder(default)] + logical_batching_enabled: bool, /// Track which protocol the query comes from. #[builder(default)] channel: Channel, @@ -458,7 +461,17 @@ impl QueryContext { &self.configuration_parameter } - /// Whether the local HTTP entry point selected write batching. + /// Whether the local HTTP entry point selected logical-table batching. + pub fn logical_batching_enabled(&self) -> bool { + self.logical_batching_enabled + } + + /// Sets local logical-table batching selection without adding a wire-visible extension. + pub fn set_logical_batching_enabled(&mut self, enabled: bool) { + self.logical_batching_enabled = enabled; + } + + /// Whether the local HTTP entry point selected ordinary-table batching. pub fn batching_enabled(&self) -> bool { self.batching_enabled } @@ -634,6 +647,7 @@ impl QueryContextBuilder { channel, batching_enabled: self.batching_enabled.unwrap_or_default(), admitted_write: None, + logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(), process_id: self.process_id.unwrap_or_default(), conn_info: self.conn_info.unwrap_or_default(), protocol_ctx: self.protocol_ctx.unwrap_or_default(), @@ -923,11 +937,17 @@ mod test { fn test_batching_selection_is_local_only() { let mut ctx = QueryContextBuilder::default().build(); assert!(!ctx.batching_enabled()); + assert!(!ctx.logical_batching_enabled()); + ctx.set_logical_batching_enabled(true); + assert!(ctx.clone().logical_batching_enabled()); + assert!(ctx.fork().logical_batching_enabled()); ctx.set_batching_enabled(true); assert!(ctx.clone().batching_enabled()); assert!(ctx.fork().batching_enabled()); let wire: api::v1::QueryContext = ctx.into(); - assert!(!QueryContext::from(wire).batching_enabled()); + let restored = QueryContext::from(wire); + assert!(!restored.batching_enabled()); + assert!(!restored.logical_batching_enabled()); let ctx = QueryContextBuilder::default() .set_extension("batching_enabled".to_string(), "true".to_string()) .build(); diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 80ba4cf11a5..d0de7408789 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -58,7 +58,7 @@ pub struct StandaloneOptions { pub postgres: PostgresOptions, pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, - /// Shared experimental ordinary-table batching; independent of Prom batching. + /// Ordinary-table batching with independent logical-table controls. pub pending_rows_batcher: PendingRowsBatcherOptions, pub jaeger: JaegerOptions, pub otlp: OtlpOptions, @@ -137,6 +137,7 @@ impl Configurable for StandaloneOptions { "wal.broker_endpoints", "event_recorder.event_types", "pending_rows_batcher.protocols", + "pending_rows_batcher.logical_table.protocols", ]) } } @@ -219,6 +220,61 @@ mod tests { use crate::options::*; + #[test] + fn test_logical_batcher_config_forwarding() { + let opts: StandaloneOptions = toml::from_str("[pending_rows_batcher]\nprotocols = ['influxdb']\n[pending_rows_batcher.logical_table]\nprotocols = ['otlp', 'prom']\npending_rows_flush_interval = '10ms'").unwrap(); + let frontend = opts.frontend_options(); + assert_eq!(frontend.pending_rows_batcher, opts.pending_rows_batcher); + assert_eq!( + frontend.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + let restored: StandaloneOptions = toml::from_str(&toml::to_string(&opts).unwrap()).unwrap(); + assert_eq!( + restored.pending_rows_batcher.logical_table, + opts.pending_rows_batcher.logical_table + ); + assert!( + toml::from_str::( + "[pending_rows_batcher.logical_table]\nprotocols = ['logs']" + ) + .is_err() + ); + } + + #[test] + fn test_logical_batcher_protocols_from_env() { + temp_env::with_vars( + [ + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__PROTOCOLS", + Some("otlp,influxdb"), + ), + ( + "STANDALONE_LOGICAL_TEST__PENDING_ROWS_BATCHER__LOGICAL_TABLE__PROTOCOLS", + Some("prom,otlp"), + ), + ], + || { + let options = + StandaloneOptions::load_layered_options(None, "STANDALONE_LOGICAL_TEST") + .unwrap(); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .protocols, + vec![ + servers::http::BatchingProtocol::Prom, + servers::http::BatchingProtocol::Otlp + ] + ); + }, + ); + } + #[test] fn test_batcher_protocols_from_env() { temp_env::with_vars( @@ -231,7 +287,7 @@ mod tests { StandaloneOptions::load_layered_options(None, "STANDALONE_BATCHER_TEST") .unwrap(); assert_eq!( - options.pending_rows_batcher.protocols, + options.pending_rows_batcher.table.protocols, vec![ servers::http::BatchingProtocol::Influxdb, servers::http::BatchingProtocol::HttpSql @@ -247,6 +303,7 @@ mod tests { assert!( !defaults .pending_rows_batcher + .table .pending_rows_batching_enabled() ); let options: StandaloneOptions = toml::from_str( @@ -259,16 +316,22 @@ flow_notification_queue_capacity = 17 "#, ) .unwrap(); - assert_eq!(options.pending_rows_batcher.max_batch_rows, 25); - assert_eq!(options.pending_rows_batcher.protocols.len(), 2); + assert_eq!(options.pending_rows_batcher.table.max_batch_rows, 25); + assert_eq!(options.pending_rows_batcher.table.protocols.len(), 2); assert_eq!( options .pending_rows_batcher + .table .flow_notification_queue_capacity .get(), 17 ); - assert!(options.pending_rows_batcher.pending_rows_batching_enabled()); + assert!( + options + .pending_rows_batcher + .table + .pending_rows_batching_enabled() + ); let serialized = toml::to_string(&options).unwrap(); let parsed: StandaloneOptions = toml::from_str(&serialized).unwrap(); assert_eq!(options.influxdb, parsed.influxdb); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index dacab7ad6bd..338b03ff68f 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -503,6 +503,233 @@ WITH( Ok(()) } + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_alignment() { + use std::time::Duration; + + use common_base::Plugins; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, exponential_histogram_data_point, + }; + use prost::Message; + use servers::batcher::pending_rows_batch_sync_enabled; + use servers::http::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + // Fixed workload, only the batcher opt-in changes. Compare stored rows, + // schema evolution and visibility under the configured acknowledgement policy. + let mut results = Vec::new(); + for enabled in [false, true] { + let standalone = GreptimeDbStandaloneBuilder::new(&format!("otlp_logical_{enabled}")) + .with_logical_batcher(BatcherOptions { + protocols: if enabled { + vec![BatchingProtocol::Otlp] + } else { + vec![] + }, + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let mut options = standalone.opts.clone(); + options.otlp.experimental_enable_exponential_histogram = true; + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_wait_flush_result"]); + let before = submissions.get_sample_count(); + for (ts, value) in [(60, 10), (120, 20)] { + let mut request = build_sum_request( + "batch.alignment", + AggregationTemporality::Cumulative, + &[(ts, value)], + ); + if ts == 120 + && let Some(metric::Data::Sum(sum)) = + &mut request.resource_metrics[0].scope_metrics[0].metrics[0].data + { + sum.data_points[0].attributes.push(keyvalue("extra", "new")); + } + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + } + if enabled { + assert_eq!( + submissions.get_sample_count() - before, + if pending_rows_batch_sync_enabled() { + 2 + } else { + 0 + }, + "logical submissions must follow the global acknowledgement policy" + ); + } + let sql = "SELECT greptime_timestamp, greptime_value, stream, extra FROM batch_alignment_total ORDER BY greptime_timestamp"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if !enabled + || pending_rows_batch_sync_enabled() + || batches.iter().map(|batch| batch.num_rows()).sum::() == 2 + { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 2 + ); + results.push(batches.pretty_print().unwrap()); + + let before_histograms = submissions.get_sample_count(); + // First create a histogram table, then reuse it in a mixed export. + // Neither request may enter the scalar-only logical batcher. + for mixed in [false, true] { + let mut request = build_sum_request( + "batch.mixed", + AggregationTemporality::Cumulative, + &[(180, 30)], + ); + let metrics = &mut request.resource_metrics[0].scope_metrics[0].metrics; + let mut histogram = metrics[0].clone(); + histogram.name = "batch.histogram".to_string(); + histogram.data = Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + aggregation_temporality: AggregationTemporality::Cumulative as i32, + data_points: vec![ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: if mixed { 4_000_000_000 } else { 3_000_000_000 }, + count: 4, + sum: Some(8.0), + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![1, 2], + }), + ..Default::default() + }], + })); + if !mixed { + metrics.clear(); + } + metrics.push(histogram); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!( + response.status().as_u16(), + 200, + "enabled={enabled}, mixed={mixed}" + ); + } + assert_eq!(submissions.get_sample_count(), before_histograms); + for (sql, expected_rows) in [ + ( + "SELECT greptime_timestamp, greptime_native_histogram FROM batch_histogram ORDER BY greptime_timestamp", + 2, + ), + ( + "SELECT greptime_timestamp, greptime_value FROM batch_mixed_total ORDER BY greptime_timestamp", + 1, + ), + ] { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + expected_rows + ); + results.push(batches.pretty_print().unwrap()); + } + + // An existing metric attached to another physical table must use + // the original routing, not the request's default physical table. + for sql in [ + "CREATE TABLE custom_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='')", + "CREATE TABLE custom_total (greptime_timestamp TIMESTAMP(3) TIME INDEX, greptime_value DOUBLE, \"stream\" STRING PRIMARY KEY) ENGINE=metric WITH ('on_physical_table'='custom_physical')", + ] { + instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + } + instance + .metrics( + build_sum_request("custom", AggregationTemporality::Cumulative, &[(60, 7)]), + ctx.clone(), + ) + .await + .unwrap(); + let output = instance + .do_query("SELECT greptime_value FROM custom_total", ctx.clone()) + .await + .remove(0) + .unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + assert!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() + .contains("7.0") + ); + + // Request-level schema policy cannot be bypassed by batching. + let mut fixed = ctx.fork(); + fixed.set_extension("auto_create_table", "false"); + assert!( + instance + .metrics( + build_sum_request( + "missing", + AggregationTemporality::Cumulative, + &[(60, 1)] + ), + Arc::new(fixed) + ) + .await + .is_err() + ); + } + assert_eq!(results[..3], results[3..]); + } + #[tokio::test(flavor = "multi_thread")] pub async fn test_otlp_on_standalone() { let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_otlp") diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 307a7f6daa2..38f546e9ccb 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -52,6 +52,7 @@ use frontend::frontend::Frontend; use frontend::instance::Instance; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; +use frontend::service_config::{BatcherOptions, PendingRowsBatcherOptions}; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; use standalone::options::StandaloneOptions; @@ -88,6 +89,7 @@ pub struct GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions, auto_create_table: bool, experimental_metric_export: bool, + logical_batcher: Option, } impl GreptimeDbStandaloneBuilder { @@ -108,6 +110,7 @@ impl GreptimeDbStandaloneBuilder { event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, experimental_metric_export: false, + logical_batcher: None, } } @@ -118,6 +121,13 @@ impl GreptimeDbStandaloneBuilder { self } + /// Configures logical-table batching for integration tests. + #[must_use] + pub fn with_logical_batcher(mut self, options: BatcherOptions) -> Self { + self.logical_batcher = Some(options); + self + } + #[must_use] pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self { self.auto_create_table = auto_create_table; @@ -377,6 +387,10 @@ impl GreptimeDbStandaloneBuilder { event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, experimental_metric_export: self.experimental_metric_export, + pending_rows_batcher: PendingRowsBatcherOptions { + logical_table: self.logical_batcher.clone(), + ..Default::default() + }, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { experimental_enable_resource_info: true, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index e4e30baeba8..020e9c023c6 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2432,7 +2432,7 @@ protocols = [] pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 @@ -2453,7 +2453,7 @@ experimental_enable_prometheus_native_histogram = false pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 -worker_channel_capacity = 65526 +worker_channel_capacity = 65536 max_inflight_requests = 3000 flow_notification_queue_capacity = 1024 From e91faa9df8885506954053a2cbc9959a7e975ab6 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:52:54 +0000 Subject: [PATCH 010/133] perf(mito2): lazily decode dense primary key columns (#9226) * perf(mito2): lazily decode dense primary key columns Signed-off-by: Lei, HUANG * perf(mito2): bypass lazy decoding for full primary keys Signed-off-by: Lei, HUANG * fix(mito-codec): preserve prefix decoding errors Signed-off-by: Lei, HUANG * refactor(mito-codec): align encoded length helper naming Rename encoded_length to encoded_len and update all callers to match the other length helpers in the module. Signed-off-by: Lei, HUANG * refactor(mito2): clarify conditional dense key decoding Rename decode_dense_pk to ensure_dense_pk_decoded so callers can see that existing decoded values are preserved and only missing caches are populated. Signed-off-by: Lei, HUANG * refactor(mito-codec): share string framing in row converter Move encoded_string_len to the parent module so Dense and Sparse use the same framing helper without depending on each other. Preserve its implementation and visibility. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/mito-codec/src/row_converter.rs | 30 ++ src/mito-codec/src/row_converter/dense.rs | 481 ++++++++++++++---- .../src/row_converter/sparse/checked.rs | 45 +- src/mito2/benches/bench_dense_index_update.rs | 162 ++++-- src/mito2/benches/bench_pk_tag_column.rs | 117 ++++- src/mito2/src/read.rs | 127 ++++- src/mito2/src/sst/index.rs | 53 ++ src/mito2/src/sst/index/column_test.rs | 90 +++- src/mito2/src/sst/index/indexer/abort.rs | 1 + src/mito2/src/sst/index/indexer/finish.rs | 1 + src/mito2/src/sst/index/indexer/update.rs | 22 + src/mito2/src/sst/parquet/flat_format.rs | 247 ++++++++- src/mito2/src/test_util/bench_util.rs | 20 +- 13 files changed, 1167 insertions(+), 229 deletions(-) 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/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/sst/index.rs b/src/mito2/src/sst/index.rs index ed3ec36a473..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; @@ -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()); 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/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, From edc81c23559e97788d01741f07bc1308b0cf4a7e Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 01:41:27 +0000 Subject: [PATCH 011/133] ci: update cargo fuzz command to use nightly toolchain explicitly (#9298) * ci: update cargo fuzz command to use nightly toolchain explicitly * ci: honor RUSTUP_TOOLCHAIN pin in fuzz orchestration script An explicit `+toolchain` argument overrides the RUSTUP_TOOLCHAIN env var in rustup precedence, so the hard-coded `cargo +nightly` in run-fuzz-targets.sh bypassed the pinned FUZZ_RUST_TOOLCHAIN (nightly-2026-03-21) configured in the workflow. - Invoke `cargo +"${RUSTUP_TOOLCHAIN:-nightly}" fuzz run` in the script so CI uses the pinned toolchain and local runs fall back to the floating nightly - Pass RUSTUP_TOOLCHAIN through to all four fuzz-test action invocations, covering the no-prebuilt-binaries path and making the reproduce command in the summary print the exact pinned toolchain - Add test_rustup_toolchain_env_is_honored covering the pinned-env scenario for both the cargo invocation args and the summary text Addresses #9298 (review). Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .github/scripts/run-fuzz-targets-test.sh | 21 ++++++++++++++++++--- .github/scripts/run-fuzz-targets.sh | 4 ++-- .github/workflows/integration.yml | 17 +++++++++++++++-- Makefile | 8 ++++---- tests-fuzz/README.md | 6 ++++++ 5 files changed, 45 insertions(+), 11 deletions(-) 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/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/Makefile b/Makefile index 2526f3e31eb..29801465351 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. diff --git a/tests-fuzz/README.md b/tests-fuzz/README.md index cc9d7eb84ee..7945bae1583 100644 --- a/tests-fuzz/README.md +++ b/tests-fuzz/README.md @@ -6,6 +6,12 @@ cargo install cargo-fuzz ``` +Note: `cargo-fuzz` instruments targets with `-Zsanitizer=fuzzer`, so fuzz +targets must be built with a **nightly** toolchain (the workspace default +toolchain does not need to be nightly — `make fuzz` / `make fuzz-ls` +invoke `cargo +nightly` for you; CI pins its own nightly in +`FUZZ_RUST_TOOLCHAIN`). + 2. Start GreptimeDB 3. Copy the `.env.example`, which is at project root, to `.env` and change the values on need. From f3eb8e6c72e2c1583d5d4483ddc149c134966234 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 23 Sep 2026 01:52:52 +0000 Subject: [PATCH 012/133] test: cut integration test time and make the storage matrix meaningful (#9308) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cut integration test time and make the storage matrix meaningful tests-integration is ~85% of workspace test CPU, and 81% of that is the S3/S3WithCache variants of the HTTP and gRPC suites. Those suites do not touch the object store: of the 70 matrix HTTP tests only one flushed and read back an SST, so the matrix was paying real AWS round trips to re-prove protocol parsing. - Point the PR CI object-store matrix at the MinIO already started by tests-integration/fixtures. Three GT_S3_* consumers did not read GT_S3_ENDPOINT_URL and would have hit real AWS with MinIO credentials; they now do. - Add a nightly Linux job against real AWS S3, and pass GT_S3_* into the release integration-test container. The release previously ran every remote-backend case as a skip and only exercised the file backend. - Give each S3WithCache test its own read cache directory. They shared /tmp/greptimedb_cache, which the datanode wipes on startup, so a starting test deleted the read cache of a running one. - Add flush -> read-back assertions to the tests whose columns have a non-trivial SST representation: JSON/JSON2 columns, native histograms, metric-engine logical tables, and tables carrying fulltext or skipping indexes whose puffin files only exist after a flush. - Move eight tests that create no table out of the storage matrix. - Make the event recorder flush interval a constructor parameter and shorten it in the event tests, which otherwise wait a 5s window per DDL they assert on. It is skipped by serde and never reaches config files. - Drop duplicates: test_grpc_zstd_compression was a verbatim copy of test_grpc_message_size_ok and is now rewritten to assert the negotiated grpc-encoding; test_execute_copy_to_{s3,oss,gcs,azblob} were strict prefixes of their copy_from siblings; two standalone/distributed event test pairs shared one assertion body. - Fix and un-ignore stddev_by_label. stddev_pop merges partial aggregates in a parallelism-dependent order, so its last digits are unstable; the test now compares values with a tolerance. - Rebase the jaeger v1 fixture on the current instant. It carries ttl=7d with 2025 timestamps, so its rows were only readable as long as they stayed in the memtable. Signed-off-by: Dennis Zhuang * test: address review — wire nightly real-S3 job into check-status, keep the short event interval The nightly `check-status` job did not depend on the new real-S3 job, so a failure there would not have reached the status or Slack notification. In database_ddl_event the short interval was set by a first `with_event_recorder_options` call and then overwritten by the pre-existing one, which carries `..Default::default()`. Merged into a single call. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../actions/build-linux-artifacts/action.yml | 23 + .github/workflows/dev-build.yml | 8 + .github/workflows/nightly-build.yml | 8 + .github/workflows/nightly-ci.yml | 53 ++- .github/workflows/release.yml | 8 + .github/workflows/rust.yml | 34 +- Makefile | 1 + src/common/event-recorder/src/recorder.rs | 30 +- src/frontend/src/instance/builder.rs | 1 + src/meta-srv/src/metasrv/builder.rs | 1 + src/object-store/src/test_util.rs | 29 -- src/object-store/tests/object_store_test.rs | 10 +- tests-integration/src/cluster.rs | 14 + tests-integration/src/test_util.rs | 79 ++-- tests-integration/src/tests/instance_test.rs | 198 +-------- tests-integration/src/tests/promql_test.rs | 49 ++- .../tests/admin_function_event.rs | 25 +- tests-integration/tests/database_ddl_event.rs | 7 +- .../tests/export_logical_tables.rs | 9 +- tests-integration/tests/flow_ddl_event.rs | 5 +- tests-integration/tests/gc_event.rs | 3 + tests-integration/tests/grpc.rs | 147 ++++++- tests-integration/tests/http.rs | 396 ++++++++++++++++-- tests-integration/tests/main.rs | 11 + .../tests/reconciliation_event.rs | 4 + tests-integration/tests/repartition_event.rs | 2 + tests-integration/tests/table_ddl_event.rs | 5 +- tests-integration/tests/view_ddl_event.rs | 16 +- tests-integration/tests/wal_prune_event.rs | 2 + 29 files changed, 805 insertions(+), 373 deletions(-) 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/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/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 1a093a2d2a5..e4483b4a42b 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -206,15 +206,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 @@ -263,15 +262,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/Makefile b/Makefile index 29801465351..7176eafd0bb 100644 --- a/Makefile +++ b/Makefile @@ -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/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/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index a3cfbcdaffe..d9b86ac5a05 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -381,6 +381,7 @@ 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); 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/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/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 3e53b6104a0..bc2ed6f20a4 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -36,6 +36,7 @@ use client::client_manager::NodeClients; use cmd::frontend::create_heartbeat_task; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; +use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; @@ -174,6 +175,7 @@ pub struct GreptimeDbClusterBuilder { shared_home_dir: Option>, meta_selector: Option, local_file_access: LocalFileAccess, + event_recorder_options: EventRecorderOptions, } impl GreptimeDbClusterBuilder { @@ -209,6 +211,7 @@ impl GreptimeDbClusterBuilder { shared_home_dir: None, meta_selector: None, local_file_access: LocalFileAccess::default(), + event_recorder_options: EventRecorderOptions::default(), } } @@ -253,6 +256,15 @@ impl GreptimeDbClusterBuilder { self } + #[must_use] + pub fn with_event_recorder_options( + mut self, + event_recorder_options: EventRecorderOptions, + ) -> Self { + self.event_recorder_options = event_recorder_options; + self + } + /// Sets whether the test frontend automatically creates tables on write. #[must_use] pub fn with_frontend_auto_create_table(mut self, auto_create_table: bool) -> Self { @@ -303,6 +315,7 @@ impl GreptimeDbClusterBuilder { ..Default::default() }, gc: self.metasrv_gc_config.clone(), + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; @@ -550,6 +563,7 @@ impl GreptimeDbClusterBuilder { fn build_frontend_options(&self) -> FrontendOptions { let mut fe_opts = FrontendOptions { auto_create_table: self.frontend_auto_create_table, + event_recorder: self.event_recorder_options.clone(), ..Default::default() }; diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 562bd48cdb0..fd02db2e5a1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -31,6 +31,7 @@ use client::{Client, Database}; use common_base::Plugins; use common_catalog::consts::MIN_USER_TABLE_ID; use common_config::Configurable; +use common_event_recorder::EventRecorderOptions; #[cfg(test)] use common_meta::DatanodeId; use common_meta::key::TableMetadataManager; @@ -328,6 +329,18 @@ impl StorageType { } } +/// Event recorder options for tests. +/// +/// The production flush interval is 5s, and the event tests interleave "run a +/// DDL, wait for its event, run the next DDL", so each barrier costs a full +/// window. A short interval removes that wait without changing what is asserted. +pub fn test_event_recorder_options() -> EventRecorderOptions { + EventRecorderOptions { + flush_interval: Duration::from_millis(100), + ..Default::default() + } +} + fn s3_test_config() -> S3Config { S3Config { connection: S3Connection { @@ -363,7 +376,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Gcs::from(&gcs_config.connection); let config = ObjectStoreConfig::Gcs(gcs_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Gcs(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Azblob => { let azblob_config = AzblobConfig { @@ -381,7 +394,7 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Azblob::from(&azblob_config.connection); let config = ObjectStoreConfig::Azblob(azblob_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Azblob(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::Oss => { let oss_config = OssConfig { @@ -398,32 +411,57 @@ pub fn get_test_store_config(store_type: &StorageType) -> (ObjectStoreConfig, Te let builder = Oss::from(&oss_config.connection); let config = ObjectStoreConfig::Oss(oss_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::Oss(TempFolder::new(&store, "/"))) + (config, TempDirGuard::remote(TempFolder::new(&store, "/"))) } StorageType::S3 | StorageType::S3WithCache => { let mut s3_config = s3_test_config(); - if *store_type == StorageType::S3WithCache { - s3_config.cache.cache_path = "/tmp/greptimedb_cache".to_string(); + // The datanode wipes `/cache/object/read` on startup, so a + // path shared between concurrently running tests lets a starting test + // delete the read cache of a running one. + let cache_dir = if *store_type == StorageType::S3WithCache { + let dir = create_temp_dir("gt_s3_read_cache"); + s3_config.cache.cache_path = dir.path().to_string_lossy().to_string(); + Some(dir) } else { s3_config.cache.enable_read_cache = false; - } + None + }; let builder = S3::from(&s3_config.connection); let config = ObjectStoreConfig::S3(s3_config); let store = ObjectStore::new(builder).unwrap(); - (config, TempDirGuard::S3(TempFolder::new(&store, "/"))) + ( + config, + TempDirGuard { + remote: Some(TempFolder::new(&store, "/")), + local_cache: cache_dir, + }, + ) } - StorageType::File => (ObjectStoreConfig::File(FileConfig {}), TempDirGuard::None), + StorageType::File => ( + ObjectStoreConfig::File(FileConfig {}), + TempDirGuard::default(), + ), } } -pub enum TempDirGuard { - None, - S3(TempFolder), - Oss(TempFolder), - Azblob(TempFolder), - Gcs(TempFolder), +#[derive(Default)] +pub struct TempDirGuard { + /// Prefix to wipe from the remote object store, absent for the file backend. + remote: Option, + /// Local read cache directory. Only held so it is removed when the guard drops. + #[allow(dead_code)] + local_cache: Option, +} + +impl TempDirGuard { + fn remote(folder: TempFolder) -> Self { + Self { + remote: Some(folder), + local_cache: None, + } + } } pub struct TestGuard { @@ -446,11 +484,7 @@ pub struct StorageGuard(pub TempDirGuard); impl TestGuard { pub async fn remove_all(&mut self) { for storage_guard in self.storage_guards.iter_mut() { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = &mut storage_guard.0 - { + if let Some(guard) = &mut storage_guard.0.remote { guard.remove_all().await.unwrap() } } @@ -465,11 +499,8 @@ impl Drop for TestGuard { common_runtime::spawn_global(async move { let mut errors = vec![]; for guard in guards { - if let TempDirGuard::S3(guard) - | TempDirGuard::Oss(guard) - | TempDirGuard::Azblob(guard) - | TempDirGuard::Gcs(guard) = guard.0 - && let Err(e) = guard.remove_all().await + if let Some(remote) = guard.0.remote + && let Err(e) = remote.remove_all().await { errors.push(e); } diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index dc4d3898860..caac6d33238 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -2195,47 +2195,6 @@ async fn test_delete(instance: Arc) { check_output_stream(output, expect).await; } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_s3(instance: Arc) { - if let Ok(bucket) = env::var("GT_S3_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - // setups - assert!(matches!(execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await.data, OutputData::AffectedRows(0))); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - let key_id = env::var("GT_S3_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_S3_ACCESS_KEY").unwrap(); - let region = env::var("GT_S3_REGION").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - // exports - let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_s3(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2268,9 +2227,21 @@ async fn test_execute_copy_from_s3(instance: Arc) { let key = env::var("GT_S3_ACCESS_KEY").unwrap(); let region = env::var("GT_S3_REGION").unwrap(); + let mut connection = format!( + "ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}'", + key_id, key, region + ); + // Honors an S3-compatible endpoint (MinIO in CI) so the test does not + // fall back to resolving the bucket against real AWS. + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") + && !endpoint.is_empty() + { + connection = format!("{},ENDPOINT='{}'", connection, endpoint); + } + let copy_to_stmt = format!( - "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - bucket, root, key_id, key, region + "Copy demo TO 's3://{}/{}/export/demo.parquet' CONNECTION ({})", + bucket, root, connection ); let output = execute_sql(&instance, ©_to_stmt).await.data; @@ -2315,10 +2286,7 @@ async fn test_execute_copy_from_s3(instance: Arc) { .data, OutputData::AffectedRows(0) )); - let sql = format!( - "{} CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", - test.sql, key_id, key, region, - ); + let sql = format!("{} CONNECTION ({})", test.sql, connection); let output = execute_sql(&instance, &sql).await.data; assert!(matches!(output, OutputData::AffectedRows(2))); @@ -2340,50 +2308,6 @@ async fn test_execute_copy_from_s3(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_oss(instance: Arc) { - if let Ok(bucket) = env::var("GT_OSS_BUCKET") - && !bucket.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let key_id = env::var("GT_OSS_ACCESS_KEY_ID").unwrap(); - let key = env::var("GT_OSS_ACCESS_KEY").unwrap(); - let endpoint = env::var("GT_OSS_ENDPOINT").unwrap(); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'oss://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", - bucket, root, key_id, key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_oss(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2490,54 +2414,6 @@ async fn test_execute_copy_from_oss(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_gcs(instance: Arc) { - if let (Ok(bucket), Ok(scope), Ok(credential)) = ( - env::var("GT_GCS_BUCKET"), - env::var("GT_GCS_SCOPE"), - env::var("GT_GCS_CREDENTIAL"), - ) && !bucket.is_empty() - { - let endpoint = env::var("GT_GCS_ENDPOINT").unwrap_or_default(); - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - let mut connection = format!("SCOPE='{}',CREDENTIAL='{}'", scope, credential); - if !endpoint.is_empty() { - connection = format!("{},ENDPOINT='{}'", connection, endpoint); - } - - let copy_to_stmt = format!( - "Copy demo TO 'gcs://{}/{}/export/demo.parquet' CONNECTION ({})", - bucket, root, connection - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_gcs(instance: Arc) { common_telemetry::init_default_ut_logging(); @@ -2646,50 +2522,6 @@ async fn test_execute_copy_from_gcs(instance: Arc) { } } -#[apply(both_instances_cases)] -async fn test_execute_copy_to_azblob(instance: Arc) { - if let (Ok(container), Ok(account_name), Ok(account_key), Ok(endpoint)) = ( - env::var("GT_AZBLOB_CONTAINER"), - env::var("GT_AZBLOB_ACCOUNT_NAME"), - env::var("GT_AZBLOB_ACCOUNT_KEY"), - env::var("GT_AZBLOB_ENDPOINT"), - ) && !container.is_empty() - { - let instance = instance.frontend(); - - assert!(matches!( - execute_sql( - &instance, - "create table demo(host string, cpu double, memory double, ts timestamp time index);", - ) - .await - .data, - OutputData::AffectedRows(0) - )); - - let output = execute_sql( - &instance, - r#"insert into demo(host, cpu, memory, ts) values - ('host1', 66.6, 1024, 1655276557000), - ('host2', 88.8, 333.3, 1655276558000) - "#, - ) - .await - .data; - assert!(matches!(output, OutputData::AffectedRows(2))); - - let root = uuid::Uuid::new_v4().to_string(); - - let copy_to_stmt = format!( - "Copy demo TO 'azblob://{}/{}/export/demo.parquet' CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", - container, root, account_name, account_key, endpoint - ); - - let output = execute_sql(&instance, ©_to_stmt).await.data; - assert!(matches!(output, OutputData::AffectedRows(2))); - } -} - #[apply(both_instances_cases)] async fn test_execute_copy_from_azblob(instance: Arc) { common_telemetry::init_default_ut_logging(); diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index c816583f135..74114bbd6c4 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -644,28 +644,57 @@ async fn two_aggregators_combined_aggrs(instance: Arc) { // eval instant at 50m stddev by (instance)(http_requests) // {instance="0"} 223.60679774998 // {instance="1"} 223.60679774998 +// +// The values are compared with a tolerance rather than against a rendered table: +// `stddev_pop` merges partial aggregates in a parallelism-dependent order, so the +// last digits of its output are not stable across runs. #[apply(both_instances_cases)] -#[ignore = "TODO(ruihang): fix this case"] async fn stddev_by_label(instance: Arc) { let instance = instance.frontend(); - create_insert_query_assert( + execute_all(&instance, AGGREGATORS_CREATE_TABLE, QueryContext::arc()).await; + execute_all(&instance, AGGREGATORS_INSERT_DATA, QueryContext::arc()).await; + + let batches = promql_query_as_batches( instance, - AGGREGATORS_CREATE_TABLE, - AGGREGATORS_INSERT_DATA, r#"stddev by (instance)(http_requests)"#, + None, + QueryContext::arc(), UNIX_EPOCH, unix_epoch_plus_100s(), Duration::from_secs(60), Duration::from_secs(0), - "+----------+---------------------+--------------------------------+\ - \n| instance | ts | STDDEVPOP(http_requests.value) |\ - \n+----------+---------------------+--------------------------------+\ - \n| 0 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n| 1 | 1970-01-01T00:00:00 | 223.606797749979 |\ - \n+----------+---------------------+--------------------------------+", ) .await; + + let batches = batches.iter().collect::>(); + assert_eq!(batches.len(), 1); + // Tag columns come back dictionary-encoded. + let instances = cast( + batches[0].column_by_name("instance").unwrap(), + &DataType::Utf8, + ) + .unwrap(); + let instances = instances.as_any().downcast_ref::().unwrap(); + let values = batches[0] + .column_by_name("stddev_pop(http_requests.value)") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + let mut actual = (0..values.len()) + .map(|i| (instances.value(i), values.value(i))) + .collect::>(); + actual.sort_by(|a, b| a.0.cmp(b.0)); + assert_eq!(actual.len(), 2, "{}", batches[0].pretty_print()); + for (expected_instance, (instance, value)) in ["0", "1"].into_iter().zip(actual) { + assert_eq!(instance, expected_instance); + assert!( + (value - 223.606_797_749_979).abs() < 1e-9, + "instance {instance}: {value}" + ); + } } // This is not derived from prometheus diff --git a/tests-integration/tests/admin_function_event.rs b/tests-integration/tests/admin_function_event.rs index d669d2900ad..43f7445ca2f 100644 --- a/tests-integration/tests/admin_function_event.rs +++ b/tests-integration/tests/admin_function_event.rs @@ -14,7 +14,6 @@ use std::sync::Arc; -use common_test_util::temp_dir::create_temp_dir; use frontend::instance::Instance; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; @@ -23,9 +22,8 @@ use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::admin::Admin; use sql::statements::statement::Statement; -use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::assert_eventually_eq; @@ -35,31 +33,12 @@ const TABLE_NAME: &str = "admin_function_event_test"; async fn test_standalone_admin_function_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_admin_function_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; assert_admin_function_events(standalone.fe_instance()).await; } -#[tokio::test(flavor = "multi_thread")] -async fn test_distributed_admin_function_events() { - let store_type = StorageType::File; - if !store_type.test_on() { - return; - } - - common_telemetry::init_default_ut_logging(); - let (store_config, _guard) = get_test_store_config(&store_type); - let home_dir = create_temp_dir("test_distributed_admin_function_events_data_home"); - let cluster = GreptimeDbClusterBuilder::new("test_distributed_admin_function_events") - .await - .with_datanodes(1) - .with_store_config(store_config) - .with_shared_home_dir(Arc::new(home_dir)) - .build(true) - .await; - assert_admin_function_events(cluster.fe_instance()).await; -} - async fn assert_admin_function_events(instance: &Arc) { execute( instance, diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index c08b81211de..1c6d3f39633 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -21,6 +21,7 @@ use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, execute_sql, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -42,7 +43,7 @@ async fn test_event_table_auto_creation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -94,7 +95,7 @@ async fn test_event_table_schema_reconciliation_with_auto_create_disabled() { event_types: Arc::new(EventTypeFilter::Only( [String::from("create_database")].into_iter().collect(), )), - ..Default::default() + ..test_event_recorder_options() }) .build() .await; @@ -185,6 +186,7 @@ async fn test_database_ddl_events() { let home_dir = create_temp_dir("test_database_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_database_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_frontend_auto_create_table(false) .with_store_config(store_config) @@ -208,6 +210,7 @@ async fn test_database_ddl_events() { async fn test_standalone_database_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_database_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 92888cced3c..41fe8382c82 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1453,9 +1453,14 @@ async fn metric_export_v2_cli_resume_roundtrip() { } #[tokio::test] -#[ignore = "requires GT_S3_* credentials and an S3-compatible test bucket"] async fn metric_export_v2_cli_s3_resume_roundtrip() { - metric_export_v2_cli_roundtrip(true).await; + // The CLI takes an explicit `--s3-endpoint`, so this only runs against an + // S3-compatible endpoint (MinIO in CI), not against bucket+region alone. + if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) + && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) + { + metric_export_v2_cli_roundtrip(true).await; + } } async fn metric_export_v2_cli_roundtrip(s3: bool) { diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 122c36c84f4..9d2884468d2 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -15,7 +15,9 @@ use std::sync::Arc; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use uuid::Uuid; use crate::event_recorder_test_util::{ @@ -31,6 +33,7 @@ const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; async fn test_standalone_flow_ddl_events() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_flow_ddl_events") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let instance = standalone.fe_instance(); diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 4721a442a80..48749b784a0 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -43,6 +43,9 @@ async fn test_batch_gc_event() { common_telemetry::init_default_ut_logging(); let (store_config, _guard) = get_test_store_config(&store_type); let home_dir = create_temp_dir("test_batch_gc_event_data_home"); + // Keeps the production event-recorder flush interval: `assert_sst_count` + // counts SST files across every table, so eagerly flushed event rows would + // show up in the count. let cluster = GreptimeDbClusterBuilder::new("test_batch_gc_event") .await .with_datanodes(1) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 025d56b959a..970f65a2383 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -19,12 +19,13 @@ use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request as RequestBody; use api::v1::greptime_response::Response as ResponseBody; use api::v1::promql_request::Promql; +use api::v1::query_request::Query; use api::v1::value::ValueData; use api::v1::{ AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, ColumnSchema, CreateTableExpr, GreptimeRequest, InsertRequest, InsertRequests, PromInstantQuery, - PromRangeQuery, PromqlRequest, RequestHeader, Row, RowInsertRequest, RowInsertRequests, Rows, - SemanticType, Value, column, + PromRangeQuery, PromqlRequest, QueryRequest, RequestHeader, Row, RowInsertRequest, + RowInsertRequests, Rows, SemanticType, Value, column, }; use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -72,6 +73,7 @@ use tests_integration::test_util::{ setup_grpc_server_with_auto_create_table_disabled, setup_grpc_server_with_user_provider, }; use tonic::Request; +use tonic::codec::CompressionEncoding; use tonic::metadata::MetadataValue; use crate::both_deployment_cases; @@ -89,9 +91,13 @@ macro_rules! grpc_test { async fn [< $test >]() { let store_type = tests_integration::test_util::StorageType::$service; if store_type.test_on() { - let _ = $crate::grpc::$test(store_type).await; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::grpc::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); } - } )* } @@ -198,29 +204,128 @@ pub async fn test_grpc_message_size_ok(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Both the server (`servers::grpc::builder`) and the standard client +/// (`configure_tonic_client!`) enable zstd, so a plain round trip never shows +/// whether negotiation actually happened. This drives raw tonic clients instead +/// and asserts on the `grpc-encoding` the server answers with. pub async fn test_grpc_zstd_compression(store_type: StorageType) { - // server and client both support gzip - let config = GrpcServerConfig { - max_recv_message_size: 1024, - max_send_message_size: 1024, - ..Default::default() - }; - let (_db, fe_grpc_server) = setup_grpc_server_with( - store_type, - "test_grpc_zstd_compression", - None, - Some(config), - None, - ) - .await; + let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_grpc_zstd_compression").await; let addr = fe_grpc_server.bind_addr().unwrap().to_string(); - let grpc_client = Client::with_urls(vec![addr]); + let ddl = |sql: &str| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + })), + }; + + // Sends zstd and accepts zstd: the server has to decode a compressed request + // body and compress its response. + let mut zstd_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd) + .accept_compressed(CompressionEncoding::Zstd); + let response = zstd_client + .handle(Request::new(ddl( + "CREATE TABLE zstd_compression (ts TIMESTAMP TIME INDEX, payload STRING)", + ))) + .await + .unwrap(); + assert_eq!( + response + .metadata() + .get("grpc-encoding") + .map(|v| v.to_str().unwrap()), + Some("zstd") + ); + + // A payload far above the zstd frame overhead, so the request body really is + // compressed rather than passed through. + let payload = "compressible-".repeat(4096); + let insert = |ts: i64| GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "zstd_compression".to_string(), + rows: Some(Rows { + schema: vec![ + ColumnSchema { + column_name: "ts".to_string(), + semantic_type: SemanticType::Timestamp as i32, + datatype: ColumnDataType::TimestampMillisecond as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "payload".to_string(), + semantic_type: SemanticType::Field as i32, + datatype: ColumnDataType::String as i32, + ..Default::default() + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(ts)), + }, + Value { + value_data: Some(ValueData::StringValue(payload.clone())), + }, + ], + }], + }), + }], + })), + }; + + let response = zstd_client + .handle(Request::new(insert(1000))) + .await + .unwrap(); + let ResponseBody::AffectedRows(rows) = response.into_inner().response.unwrap(); + assert_eq!(rows.value, 1); + + // A client that does not advertise zstd gets an uncompressed response. + let mut plain_client = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let response = plain_client + .handle(Request::new(insert(2000))) + .await + .unwrap(); + assert!(response.metadata().get("grpc-encoding").is_none()); + + // Both payloads survived their respective paths intact. let db = Database::new_with_dbname( format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), - grpc_client, + Client::with_urls(vec![addr]), ); - db.sql("show tables;").await.unwrap(); + let sql = format!( + "SELECT count(*) AS c FROM zstd_compression WHERE length(payload) = {}", + payload.len() + ); + let OutputData::Stream(stream) = db.sql(&sql).await.unwrap().data else { + panic!("expected a stream"); + }; + let recordbatches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + recordbatches.pretty_print().unwrap(), + "\ ++---+ +| c | ++---+ +| 2 | ++---+" + ); + let _ = fe_grpc_server.shutdown().await; } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 020e9c023c6..4cc409deec4 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -93,6 +93,31 @@ use yaml_rust::YamlLoader; use crate::both_deployment_cases; use crate::event_recorder_test_util::assert_procedure_actor_by_table; +/// Registers tests that exercise nothing below the HTTP layer, so running them +/// against a remote object store only pays bucket setup and teardown. +#[macro_export] +macro_rules! http_local_tests { + ($($(#[$meta:meta])* $test:ident),*,) => { + mod integration_http_local_test { + $( + #[tokio::test(flavor = "multi_thread")] + $( + #[$meta] + )* + async fn $test() { + let store_type = tests_integration::test_util::StorageType::File; + // Support both unit tests and fallible tests without discarding errors. + let result = $crate::http::$test(store_type).await; + assert_eq!( + std::process::Termination::report(result), + std::process::ExitCode::SUCCESS, + ); + } + )* + } + }; +} + #[macro_export] macro_rules! http_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -134,12 +159,7 @@ macro_rules! http_tests { test_prometheus_promql_api, test_prometheus_label_replace_response, test_prom_http_api, - test_metrics_api, - test_health_api, - test_status_api, test_config_api, - test_dynamic_tracer_toggle, - test_dashboard_path, test_dashboard_api, test_prometheus_remote_write, test_prometheus_remote_write_v2, @@ -151,7 +171,6 @@ macro_rules! http_tests { test_vm_proto_remote_write, test_pipeline_api, - test_test_pipeline_api, test_pipeline_name_in_header, test_plain_text_ingestion, test_pipeline_auto_transform, @@ -187,8 +206,6 @@ macro_rules! http_tests { test_loki_json_logs_with_pipeline, test_elasticsearch_logs, test_elasticsearch_logs_with_index, - test_splunk_health, - test_splunk_health_is_public, test_splunk_logs, test_splunk_raw, test_log_query, @@ -1810,15 +1827,21 @@ pub async fn test_splunk_logs(store_type: StorageType) { .await; let client = TestClient::new(app).await; - // Authenticated SQL query (the user-provider harness requires auth on /v1/sql). + // Authenticated SQL statement (the user-provider harness requires auth on + // /v1/sql). Posted as a form so quotes and parentheses survive. async fn query(client: &TestClient, sql: &str) -> String { let res = client - .get(format!("/v1/sql?sql={sql}").as_str()) + .post("/v1/sql") .header("Authorization", basic_auth("greptime_user", "greptime_pwd")) + .form(&[("sql", sql)]) .send() .await; - assert_eq!(res.status(), StatusCode::OK, "query failed: {sql}"); - res.text().await + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "query failed: {sql}: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "query failed: {sql}: {body}"); + body } // HEC `Authorization: Splunk ` + JSON content type. @@ -1969,6 +1992,34 @@ transform: "custom pipeline should have dropped sourcetype (identity would keep it): {create}" ); + // 6b. Ingest into a pre-created table with a fulltext index, then flush: the + // index file is only written on flush, so this is the one path that reads it + // back. `matches_term` is the predicate the index applier recognises. + query( + &client, + "create table splunk_indexed (greptime_timestamp timestamp(9) time index, \ + \"event\" string fulltext index with(backend='bloom'))", + ) + .await; + let res = send_req( + &client, + splunk_headers(), + "/v1/splunk/services/collector/event?table=splunk_indexed", + br#"{"event":"disk pressure on web-01","time":1700000020}"#.to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + query(&client, "ADMIN FLUSH_TABLE('splunk_indexed')").await; + let rows = get_rows_from_output( + &query( + &client, + "select \"event\" from splunk_indexed where matches_term(\"event\", 'pressure')", + ) + .await, + ); + assert_eq!(rows, r#"[["disk pressure on web-01"]]"#); + // 7. Auth failures return HEC codes: missing token -> 2 (401), bad token -> 4 (403). let res = send_req( &client, @@ -2911,6 +2962,18 @@ pub async fn test_prometheus_remote_write(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "metric2").await; + validate_data( + "prometheus_remote_write_after_flush", + &client, + "select * from metric2", + table_val, + ) + .await; + // Prom RW tables carry the metric identity; type is inferred from naming. validate_data( "prometheus_remote_write_semantic_identity", @@ -3217,11 +3280,24 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + let histogram_sql = "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;"; + let histogram_rows = "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]"; validate_data( "prometheus_remote_write_v2_native_histogram_rows", &client, - "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + histogram_sql, + histogram_rows, + ) + .await; + + // The native histogram column is a struct of nested lists, so its SST + // encoding differs from its memtable representation. + flush_table(&client, "remote_write_v2_latency_seconds").await; + validate_data( + "prometheus_remote_write_v2_native_histogram_rows_after_flush", + &client, + histogram_sql, + histogram_rows, ) .await; @@ -4283,6 +4359,31 @@ transform: ) .await; + // The inverted, skipping and fulltext index files are only written on flush, + // so the queries below are the first thing that reads them back. + flush_table(&client, "test_db.logs1").await; + validate_data( + "pipeline_db_inverted_index", + &client, + "select id2 from test_db.logs1 where id1 = 2436", + "[[2528]]", + ) + .await; + validate_data( + "pipeline_db_skipping_index", + &client, + "select id1 from test_db.logs1 where type = 'I'", + "[[2436]]", + ) + .await; + validate_data( + "pipeline_db_fulltext_index", + &client, + "select id1 from test_db.logs1 where matches(log, 'ClusterAdapter')", + "[[2436]]", + ) + .await; + // 5. remove pipeline let encoded_ver_str: String = url::form_urlencoded::byte_serialize(version_str.as_bytes()).collect(); @@ -5079,6 +5180,24 @@ transform: ) .await; + // The fulltext and skipping index files are only written on flush, so the + // queries below are the first thing that reads them back. + flush_table(&client, "pipeline_index_options").await; + validate_data( + "pipeline_index_options_fulltext", + &client, + "select trace_id from pipeline_index_options where matches(message, 'greptime')", + "[[42]]", + ) + .await; + validate_data( + "pipeline_index_options_skipping", + &client, + "select message from pipeline_index_options where trace_id = 42", + r#"[["hello greptime"]]"#, + ) + .await; + guard.remove_all().await; } @@ -6995,6 +7114,18 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // A metric-engine logical table stores its rows in the shared physical + // table, so reading them back from SST files goes through a different + // column mapping than the memtable read does. + flush_table(&client, "claude_code_cost_usage_USD_total").await; + validate_data( + "otlp_metrics_all_select_after_flush", + &client, + "select * from `claude_code_cost_usage_USD_total` order by model desc;", + expected, + ) + .await; + // The synthesized resource descriptor: a plain mito info table keyed by // the raw OTel attribute keys (service.name is the only allowlisted // resource attribute in this payload), independent of the promote/ignore @@ -7846,6 +7977,17 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { ) .await; + // The attribute columns are JSON and the event/link columns are lists, whose + // SST encodings differ from their memtable representations. + flush_table(&client, "opentelemetry_traces").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from opentelemetry_traces;", + expected, + ) + .await; + guard.remove_all().await; } @@ -7972,6 +8114,30 @@ pub(crate) async fn test_otlp_traces_v2( ) .await; + // JSON2 path extraction has to resolve the same way against SST files as it + // does against the memtable. + flush_table(&client, "trace_v2_spans").await; + validate_data( + "otlp_traces_v2_json2_after_flush", + &client, + "select service_name, span_attributes.\"http.status_code\"::BIGINT, \ + resource_attributes.\"deployment.environment\"::STRING, \ + scope_attributes.enabled::BOOLEAN from trace_v2_spans;", + r#"[["frontend",200,"production",true]]"#, + ) + .await; + validate_data( + "otlp_traces_v2_events_links_after_flush", + &client, + r#"select json_get_string(span_events, '$[0].name'), + json_get_int(span_events, '$[0].attributes."event.code"'), + json_get_string(span_links, '$[0].trace_id'), + json_get_string(span_links, '$[0].attributes."link.type"') + from trace_v2_spans;"#, + r#"[["cache.hit",7,"cc9e0991a2e63d274984bd44ee669203","follows_from"]]"#, + ) + .await; + guard.remove_all().await; Ok(()) } @@ -8125,6 +8291,17 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { let expected = r#"[[1736480942444376000,1736480942444499000,123000,null,"c05d7a4ec8e1f231f02ed6e8da8655b4","d24f921c75f68e23","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444376000,1736480942444499000,123000,"d24f921c75f68e23","c05d7a4ec8e1f231f02ed6e8da8655b4","9630f2916e2f7909","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]],[1736480942444589000,1736480942444712000,123000,null,"cc9e0991a2e63d274984bd44ee669203","eba7be77e3558179","SPAN_KIND_CLIENT","lets-go","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-server",[],[]],[1736480942444589000,1736480942444712000,123000,"eba7be77e3558179","cc9e0991a2e63d274984bd44ee669203","8f847259b0f6e1ab","SPAN_KIND_SERVER","okey-dokey-0","STATUS_CODE_UNSET","","","telemetrygen","","telemetrygen","1.2.3.4","telemetrygen-client",[],[]]]"#; validate_data("otlp_traces", &client, "select * from mytable;", expected).await; + // `span_events` / `span_links` are JSON columns and the table is partitioned + // on `trace_id`, so the read back has to reassemble SST files across regions. + flush_table(&client, "mytable").await; + validate_data( + "otlp_traces_after_flush", + &client, + "select * from mytable;", + expected, + ) + .await; + // The trace v1 main table carries the trace identity (events/links preserved as // JSON columns by the v1 model). validate_data( @@ -8961,6 +9138,17 @@ pub async fn test_otlp_logs(store_type: StorageType) { ) .await; + // The attribute columns are JSON, whose SST encoding differs from their + // memtable representation. + flush_table(&client, "opentelemetry_logs").await; + validate_data( + "otlp_logs_after_flush", + &client, + "select * from opentelemetry_logs;", + expected, + ) + .await; + // The auto-created log table carries the log identity. validate_data( "otlp_logs_semantic_identity", @@ -9343,6 +9531,17 @@ pub async fn test_loki_pb_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_pb_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9548,6 +9747,17 @@ pub async fn test_loki_json_logs(store_type: StorageType) { ) .await; + // `structured_metadata` is a JSON column, whose SST encoding differs from its + // memtable representation. + flush_table(&client, "loki_table_name").await; + validate_data( + "loki_json_content_after_flush", + &client, + "select * from loki_table_name;", + expected, + ) + .await; + guard.remove_all().await; } @@ -9767,7 +9977,7 @@ pub async fn test_log_query(store_type: StorageType) { // prepare data with SQL API let res = client - .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string);") + .get("/v1/sql?sql=create table logs (`ts` timestamp time index, `message` string fulltext index with(backend='bloom'));") .send() .await; assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); @@ -9842,6 +10052,35 @@ pub async fn test_log_query(store_type: StorageType) { ] ); + // The fulltext index file is only written on flush. Replay the log query to + // cover the SST read path, and add a `matches_term` query, which is the + // predicate the index applier actually recognises. + flush_table(&client, "logs").await; + let res = client + .post("/v1/logs") + .header("Content-Type", "application/json") + .body(serde_json::to_string(&log_query).unwrap()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "{:?}", res.text().await); + let resp = res.text().await; + let output = serde_json::from_str::(&resp).unwrap(); + let rows = output["output"][0]["records"]["rows"].as_array().unwrap(); + let mut flushed_messages = rows + .iter() + .map(|row| row[1].as_str().unwrap()) + .collect::>(); + flushed_messages.sort_unstable(); + assert_eq!(flushed_messages, messages); + + validate_data( + "log_query_fulltext_index", + &client, + "select message from logs where matches_term(message, 'at-explicit-end')", + r#"[["at-explicit-end"]]"#, + ) + .await; + let res = client .get("/v1/sql?sql=create table logs_limit (`ts` timestamp time index, `message` string);") .send() @@ -10865,11 +11104,20 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let mut req: ExportTraceServiceRequest = serde_json::from_str(content).unwrap(); - // Modify timestamp fields + // The fixture is written against a fixed base instant. Move it next to now: + // this table is created with `ttl=7d`, and stale spans would be dropped as + // expired the moment they are flushed into SST files. let now = Utc::now().timestamp_nanos_opt().unwrap() as u64; + let delta_micros = (now / 1_000) as i64 - 60_000_000 - JAEGER_V1_FIXTURE_BASE_MICROS as i64; + let delta_nanos = delta_micros * 1_000; for span in req.resource_spans.iter_mut() { for scope_span in span.scope_spans.iter_mut() { - // Only modify the timestamp fields for the span with the name "test-jaeger-get-operations" to current time. + for span in scope_span.spans.iter_mut() { + span.start_time_unix_nano = (span.start_time_unix_nano as i64 + delta_nanos) as u64; + span.end_time_unix_nano = (span.end_time_unix_nano as i64 + delta_nanos) as u64; + } + // `/api/operations` is queried without a time range, so these spans + // stay pinned to the last few seconds. if scope_span.scope.as_ref().unwrap().name == "test-jaeger-get-operations" { for span in scope_span.spans.iter_mut() { span.start_time_unix_nano = now - 5_000_000_000; // 5 seconds ago @@ -10946,7 +11194,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/operations` API. @@ -10995,12 +11244,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/services/{service_name}/operations` API. let res = client - .get("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/services/test-jaeger-query-api/operations?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11023,7 +11273,8 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API without start and end. @@ -11149,12 +11400,26 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); + assert_eq!(resp, expected); + + // Replay the same lookup against SST files: `trace_id` and `service_name` + // carry BLOOM skipping indexes whose puffin files only exist after a flush, + // and `span_events` / `span_links` are JSON columns. + flush_table(&client, trace_table_name).await; + let res = client + .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea") + .header("x-greptime-trace-table-name", trace_table_name) + .send() + .await; + assert_eq!(StatusCode::OK, res.status()); + let resp: Value = serde_json::from_str(&res.text().await).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API with start and end in microseconds. let res = client - .get("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces/5611dce1bc9ebed65352d99a027b08ea?start=1738726754492421&end=1738726754642422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11275,12 +11540,16 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces/{trace_id}` API for non-existent trace. let res = client - .get("/v1/jaeger/api/traces/0000000000000000000000000000dead") + .get(&shift_fixture_micros( + "/v1/jaeger/api/traces/0000000000000000000000000000dead", + delta_micros, + )) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11299,12 +11568,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { } "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&operation=access-mysql&start=1738726754492421&end=1738726754642422&tags=%7B%22operation.type%22%3A%22access-mysql%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11426,13 +11696,14 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with tags. // 1. first query without tags, get 2 results let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11443,12 +11714,13 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // 2. second query with tags, get 1 result let res = client -.get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D") +.get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754492422&end=1738726754592422&tags=%7B%22operation.type%22%3A%22access-pg%22%7D", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11459,14 +11731,15 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { "#; let resp: Value = serde_json::from_str(&res.text().await).unwrap(); - let expected: Value = serde_json::from_str(expected).unwrap(); + let expected: Value = + serde_json::from_str(&shift_fixture_micros(expected, delta_micros)).unwrap(); assert_eq!(resp, expected); // Test `/api/traces` API with Grafana User-Agent. // When user agent is Grafana, only return at most 3 spans per trace (earliest by timestamp). // Trace `5611dce1bc9ebed65352d99a027b08fb` has 4 spans, so only 3 should be returned. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Grafana/8.0.0") .send() @@ -11514,7 +11787,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API without User-Agent (default behavior). // All 4 spans should be returned for the trace. let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .send() .await; @@ -11534,7 +11807,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { // Test `/api/traces` API with Jaeger User-Agent (should return all spans like default). let res = client - .get("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004") + .get(&shift_fixture_micros("/v1/jaeger/api/traces?service=test-jaeger-query-api&start=1738726754600000&end=1738726754700004", delta_micros)) .header("x-greptime-trace-table-name", trace_table_name) .header("User-Agent", "Jaeger-Query/1.0.0") .send() @@ -11613,6 +11886,59 @@ pub async fn test_influxdb_write(store_type: StorageType) { guard.remove_all().await; } +/// Flushes `table` so the rows written so far move from the memtable into SST +/// files. Replaying a query afterwards covers the encode/decode round trip and, +/// for indexed columns, the puffin files that only exist after a flush. +/// Base microsecond instant the `test_jaeger_query_api_for_trace_v1` fixture and +/// its expected responses are written against. +const JAEGER_V1_FIXTURE_BASE_MICROS: u64 = 1_738_726_754_492_421; + +/// Shifts every fixture timestamp in `text` by `delta` microseconds, so expected +/// responses and query bounds follow the payload after it is moved next to now. +/// +/// Only integer literals within a day of [`JAEGER_V1_FIXTURE_BASE_MICROS`] are +/// rewritten; durations, counts and ids are left alone. +fn shift_fixture_micros(text: &str, delta: i64) -> String { + let mut out = String::with_capacity(text.len()); + let mut rest = text; + while let Some(start) = rest.find(|c: char| c.is_ascii_digit()) { + out.push_str(&rest[..start]); + let digits_end = rest[start..] + .find(|c: char| !c.is_ascii_digit()) + .map_or(rest.len(), |offset| start + offset); + let digits = &rest[start..digits_end]; + match digits.parse::() { + Ok(value) if (value - JAEGER_V1_FIXTURE_BASE_MICROS as i64).abs() < 86_400_000_000 => { + out.push_str(&(value + delta).to_string()) + } + _ => out.push_str(digits), + } + rest = &rest[digits_end..]; + } + out.push_str(rest); + out +} + +async fn flush_table(client: &TestClient, table: &str) { + // The argument is parsed as a SQL table name, so an unquoted identifier is + // lower-cased and will not resolve a mixed-case table. + let quoted = table + .split('.') + .map(|part| format!("\"{part}\"")) + .collect::>() + .join("."); + let res = client + .post("/v1/sql") + .form(&[("sql", format!("ADMIN FLUSH_TABLE('{quoted}')"))]) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!(status, StatusCode::OK, "flush {table} fail: {body}"); + // A failing statement still comes back as 200 with an `error` envelope. + assert!(!body.contains("\"error\""), "flush {table} fail: {body}"); +} + async fn validate_data(test_name: &str, client: &TestClient, sql: &str, expected: &str) { let res = client .get(format!("/v1/sql?sql={sql}").as_str()) diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index b61b6cd3e19..278a7ba3851 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -53,6 +53,17 @@ grpc_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); http_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); +http_local_tests!( + test_health_api, + test_status_api, + test_metrics_api, + test_splunk_health, + test_splunk_health_is_public, + test_dynamic_tracer_toggle, + test_dashboard_path, + test_test_pipeline_api, +); + #[tokio::test(flavor = "multi_thread")] async fn test_http_auth_from_standalone_user_provider_config() { http::test_http_auth_from_standalone_user_provider_config().await; diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs index 19eaea6ff2c..a366217a7d9 100644 --- a/tests-integration/tests/reconciliation_event.rs +++ b/tests-integration/tests/reconciliation_event.rs @@ -20,6 +20,7 @@ use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use table::table_reference::TableReference; use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; @@ -33,6 +34,7 @@ async fn test_catalog_and_database_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("catalog_database_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -243,6 +245,7 @@ async fn test_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; @@ -351,6 +354,7 @@ async fn test_logical_table_reconciliation_events() { let cluster = GreptimeDbClusterBuilder::new("logical_table_reconciliation_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .build(true) .await; diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 906eb3cfee3..aaefcc2d49c 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -22,6 +22,7 @@ use mito2::gc::GcConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use crate::event_recorder_test_util::{ @@ -44,6 +45,7 @@ async fn test_repartition_event() { let home_dir = create_temp_dir("test_repartition_event_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_repartition_event") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_shared_home_dir(Arc::new(home_dir)) diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 8bad66dae3c..6cce1736151 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -41,7 +41,9 @@ use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::test_util::setup_authenticated_grpc_database; +use tests_integration::test_util::{ + setup_authenticated_grpc_database, test_event_recorder_options, +}; use crate::event_recorder_test_util::{ assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, @@ -79,6 +81,7 @@ async fn test_table_ddl_procedure_events() { } let cluster = GreptimeDbClusterBuilder::new("table_ddl_procedure_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_metasrv_gc_config(gc_options) .with_datanode_gc_config(GcConfig { diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index b21838995c3..8d17e9b4e25 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -17,9 +17,9 @@ use std::sync::Arc; use common_test_util::temp_dir::create_temp_dir; use common_wal::config::DatanodeWalConfig; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::standalone::GreptimeDbStandaloneBuilder; use tests_integration::test_util::{ StorageType, get_test_store_config, setup_authenticated_grpc_database, + test_event_recorder_options, }; use uuid::Uuid; @@ -44,6 +44,7 @@ async fn test_view_ddl_events() { let home_dir = create_temp_dir("test_view_ddl_events_data_home"); let cluster = GreptimeDbClusterBuilder::new("test_view_ddl_events") .await + .with_event_recorder_options(test_event_recorder_options()) .with_datanodes(1) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Noop) @@ -58,19 +59,6 @@ async fn test_view_ddl_events() { execute_view_ddl(instance, &source_table, &view).await; } -#[tokio::test(flavor = "multi_thread")] -async fn test_standalone_view_ddl_events() { - common_telemetry::init_default_ut_logging(); - let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_view_ddl_events") - .build() - .await; - let suffix = Uuid::new_v4().simple(); - let source_table = format!("view_ddl_event_source_{suffix}"); - let view = format!("view_ddl_event_{suffix}"); - - execute_view_ddl(standalone.fe_instance(), &source_table, &view).await; -} - async fn execute_view_ddl( instance: &Arc, source_table: &str, diff --git a/tests-integration/tests/wal_prune_event.rs b/tests-integration/tests/wal_prune_event.rs index db86ade4d9d..3879fb3e653 100644 --- a/tests-integration/tests/wal_prune_event.rs +++ b/tests-integration/tests/wal_prune_event.rs @@ -26,6 +26,7 @@ use meta_srv::procedure::wal_prune::{Context as WalPruneContext, WalPruneProcedu use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::record::Record; use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::test_event_recorder_options; use crate::event_recorder_test_util::assert_single_event; @@ -35,6 +36,7 @@ async fn test_standalone_wal_prune_event() { common_telemetry::init_default_ut_logging(); let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_wal_prune_event") + .with_event_recorder_options(test_event_recorder_options()) .build() .await; let topic_name = format!("test_standalone_wal_prune_event-{}", uuid::Uuid::new_v4()); From 0f118bc5ba3ad591824562f3ede4b8fd536fb86d Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 01:55:46 +0000 Subject: [PATCH 013/133] fix: serialize struct to json in postgres (#9170) * feat: serialize struct to json in postgres * fix: support view scalars and preserve null structs in scalar-to-value conversion Address PR review: - Utf8View/BinaryView ScalarValues now convert like their non-view forms instead of failing row extraction for struct columns - a null struct scalar converts to Value::Null so a null struct inside a list stays null in the serialized JSON Signed-off-by: Ning Sun * fix: return errors instead of panics for unsupported arrow field types Struct-typed query results with arrow field types greptimedb cannot represent (e.g. Decimal256) used to panic during schema conversion and row extraction, dropping the client connection. They now surface as query errors: - ConcreteDataType::try_from builds struct types fallibly via the new StructType::try_from_arrow_fields - Value::try_from(ScalarValue::Struct) uses the same fallible path - new try_value_from_array converts an arrow element to Value with error propagation, used by the postgres struct encoding Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/datatypes/src/data_type.rs | 4 +- src/datatypes/src/types/struct_type.rs | 19 ++ src/datatypes/src/value.rs | 99 +++++++- src/servers/src/postgres/types.rs | 315 ++++++++++++++++++++++++- tests-integration/src/test_util.rs | 55 +++++ tests-integration/tests/sql.rs | 233 +++++++++++++++++- 6 files changed, 711 insertions(+), 14 deletions(-) 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/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/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 033a1c665e2..c83debd39a1 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -19,7 +19,7 @@ use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; -use arrow::array::{Array, AsArray}; +use arrow::array::{Array, ArrayRef, AsArray}; use arrow_pg::encoder::{Encoder, encode_value}; use arrow_pg::list_encoder::encode_list; use arrow_schema::{DataType, TimeUnit}; @@ -37,7 +37,7 @@ use datatypes::json::JsonSettings; use datatypes::prelude::{ConcreteDataType, DataType as _, Value}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; -use datatypes::value::StructValue; +use datatypes::value::try_value_from_array; use futures::Stream; use pg_interval::Interval as PgInterval; use pgwire::api::Type; @@ -234,7 +234,8 @@ impl ToSqlText for OidAliasValue<'_> { } } -/// this function will encode greptime's `StructValue` into PostgreSQL jsonb type +/// this function will encode greptime's structured values (`StructValue`, or a +/// list of them) into PostgreSQL jsonb type /// /// Note that greptimedb has different types of StructValue for storing json data, /// based on policy defined in `JsonSettings`. But here the `StructValue` @@ -245,18 +246,35 @@ impl ToSqlText for OidAliasValue<'_> { /// arrays: element in array must be the same type fn encode_struct( _query_ctx: &QueryContextRef, - struct_value: StructValue, + value: Value, builder: &mut S, pg_field: &FieldInfo, ) -> PgWireResult<()> { let encoding_setting = JsonSettings::default(); let json_value = encoding_setting - .decode(Value::Struct(struct_value)) + .decode(value) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; builder.encode_field(&json_value, pg_field) } +/// Encodes one row of a structured column (`Struct`, or `List` of `Struct`) as +/// a PostgreSQL `json` value. +/// +/// Returns an error for arrow field types greptimedb cannot represent +/// (e.g. `Decimal256`) instead of panicking, so the client receives a proper +/// error response. +fn encode_structured_json( + query_ctx: &QueryContextRef, + column: &ArrayRef, + i: usize, + encoder: &mut S, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + let value = try_value_from_array(column.as_ref(), i).map_err(convert_err)?; + encode_struct(query_ctx, value, encoder, pg_field) +} + pub(crate) struct RecordBatchRowStream where S: Encoder, @@ -369,6 +387,11 @@ where } } + DataType::List(list_field) + if matches!(list_field.data_type(), DataType::Struct(_)) => + { + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; + } DataType::List(_) => { let array = column.as_list::(); let items = array.value(i); @@ -376,7 +399,7 @@ where encode_list(encoder, items, pg_field)?; } DataType::Struct(_) => { - encode_struct(query_ctx, Default::default(), encoder, pg_field)?; + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; } DataType::Utf8 => { let arrow_field = arrow_schema.field(j); @@ -445,7 +468,9 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC_ARRAY), &ConcreteDataType::Json(_) => Ok(Type::JSON_ARRAY), &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL_ARRAY), - &ConcreteDataType::Struct(_) => Ok(Type::JSON_ARRAY), + // a list of structs is encoded as a single JSON array, matching the + // struct-to-json mapping above + &ConcreteDataType::Struct(_) => Ok(Type::JSON), &ConcreteDataType::Dictionary(_) | &ConcreteDataType::Vector(_) | &ConcreteDataType::List(_) => server_error::UnsupportedDataTypeSnafu { @@ -1921,6 +1946,282 @@ mod test { } } + #[test] + fn test_encode_list_of_struct_as_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::{ListValue, StructValue}; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + // struct type: {a: int32, b: float64} + let fields = vec![ + StructField::new("a".to_string(), ConcreteDataType::int32_datatype(), true), + StructField::new("b".to_string(), ConcreteDataType::float64_datatype(), true), + ]; + let struct_type = + ConcreteDataType::struct_datatype(StructType::new(Arc::new(fields.clone()))); + let list_type = ConcreteDataType::list_datatype(Arc::new(struct_type.clone())); + + let make_struct = |a: i32, b: f64| { + Value::Struct(StructValue::new( + vec![Value::Int32(a), Value::Float64(b.into())], + StructType::new(Arc::new(fields.clone())), + )) + }; + let rows = vec![ + Value::List(ListValue::new( + vec![make_struct(1, 1.5), Value::Null, make_struct(3, 3.5)], + Arc::new(struct_type.clone()), + )), + Value::List(ListValue::new(vec![], Arc::new(struct_type.clone()))), + Value::Null, + ]; + let vector = Helper::try_from_row_into_vector(&rows, &list_type).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "l".to_string(), + list_type.clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + // The list-of-struct column is declared as a single json value. + let pg_schema = Arc::new(vec![FieldInfo::new( + "l".into(), + None, + None, + type_gt_to_pg(&list_type).unwrap(), + FieldFormat::Text, + )]); + assert_eq!(pg_schema[0].datatype(), &Type::JSON); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL list row stays NULL. + assert_eq!(decoded[2], None); + + // A list of structs is a JSON array; null elements stay null. + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!([ + {"a": 1, "b": 1.5}, + null, + {"a": 3, "b": 3.5}, + ]) + ); + + // An empty list is an empty JSON array. + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!(json, serde_json::json!([])); + } + + #[test] + fn test_encode_struct_value_as_json() { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use api::helper::pb_value_to_value_ref; + use api::v1::Value as PbValue; + use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, + }; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + fn histogram_value(histogram: &Histogram) -> Value { + let value_data = encode_native_histogram(histogram).unwrap(); + let pb_value = PbValue { + value_data: Some(value_data), + }; + let column_schema = native_histogram_column_schema().unwrap(); + Value::from(pb_value_to_value_ref( + &pb_value, + column_schema.datatype_extension.as_ref(), + )) + } + + // One integer-count histogram and one float-count histogram, mirroring + // what prometheus remote-write v2 stores. The expected JSON matches the + // HTTP query output for the same payloads. + let int_histogram = Histogram { + schema: 1, + sum: 10.0, + zero_threshold: 0.001, + reset_hint: 2, + start_timestamp: 1500, + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + count: Some(Count::CountInt(8)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }; + let float_histogram = Histogram { + schema: 2, + sum: 20.0, + zero_threshold: 0.002, + reset_hint: 3, + start_timestamp: 2500, + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + ..Default::default() + }; + + let rows = vec![ + histogram_value(&int_histogram), + histogram_value(&float_histogram), + Value::Null, + ]; + let vector = + Helper::try_from_row_into_vector(&rows, native_histogram_value_type()).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "greptime_native_histogram", + native_histogram_value_type().clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + let pg_schema = Arc::new(vec![FieldInfo::new( + "greptime_native_histogram".into(), + None, + None, + Type::JSON, + FieldFormat::Text, + )]); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL struct row stays NULL. + assert_eq!(decoded[2], None); + + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + } + #[test] fn test_encode_utf8_oid_alias_data() { let aliases = [ diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index fd02db2e5a1..37a1740ccd1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1264,6 +1264,61 @@ pub async fn setup_pg_server_with_user_provider( (instance.guard, Arc::new(pg_server)) } +/// Sets up a standalone instance with both a Prometheus remote-write HTTP app +/// (native histograms enabled) and a Postgres server attached, so native +/// histogram data written via remote-write can be queried over the Postgres +/// protocol. +pub async fn setup_pg_server_with_prom_native_histogram( + store_type: StorageType, + name: &str, +) -> (TestGuard, Router, Arc>) { + unsafe { + std::env::set_var("TZ", "UTC"); + } + + let instance = setup_standalone_instance(name, store_type).await; + + // Prometheus remote-write HTTP app with native histograms enabled. + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, true) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); + let app = http_server.build(http_server.make_app()).unwrap(); + + // Postgres server on the same instance. + let runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("pg-runtime") + .build() + .unwrap(); + + let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port()); + let opts = PostgresOptions { + addr: fe_pg_addr.clone(), + ..Default::default() + }; + let tls_server_config = Arc::new( + ReloadableTlsServerConfig::try_new(opts.tls.clone()) + .expect("Failed to load certificates and keys"), + ); + + let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + )); + + pg_server + .start(fe_pg_addr.parse::().unwrap()) + .await + .unwrap(); + + (instance.guard, app, Arc::new(pg_server)) +} + pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) { let catalog_manager = instance .catalog_manager() diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 6c58ab0246c..e3b66fc4447 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -31,7 +31,8 @@ use sqlx::{Connection, Executor, Row}; use tests_integration::test_util::{ StorageType, setup_mysql_server, setup_mysql_server_with_slow_query_threshold, setup_mysql_server_with_user_provider, setup_pg_server, - setup_pg_server_with_slow_query_threshold, setup_pg_server_with_user_provider, + setup_pg_server_with_prom_native_histogram, setup_pg_server_with_slow_query_threshold, + setup_pg_server_with_user_provider, }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; @@ -94,6 +95,8 @@ macro_rules! sql_tests { test_mysql_federated_prepare_stmt, test_mysql_prepare_tql_and_show, test_postgres_extended_query_row_returning_statements, + test_postgres_native_histogram, + test_postgres_struct_types, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -879,6 +882,234 @@ pub async fn test_postgres_bytea(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_postgres_native_histogram(store_type: StorageType) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use axum::http::StatusCode; + use prost::Message; + use servers::http::test_helpers::TestClient; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_store; + + let (mut guard, app, fe_pg_server) = + setup_pg_server_with_prom_native_histogram(store_type, "test_postgres_native_histogram") + .await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + let http_client = TestClient::new(app).await; + + // Ingest one integer-count and one float-count native histogram via + // prometheus remote-write v2. + let write_request = remote_write_v2::request_with_labels_and_histograms( + vec![ + (prom_store::METRIC_NAME_LABEL, "pg_native_histogram_seconds"), + ("job", "api"), + ("instance", "localhost:9090"), + ], + vec![ + Histogram { + count: Some(Count::CountInt(8)), + sum: 10.0, + schema: 1, + zero_threshold: 0.001, + zero_count: Some(ZeroCount::ZeroCountInt(1)), + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + reset_hint: 2, + timestamp: 3000, + start_timestamp: 1500, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + sum: 20.0, + schema: 2, + zero_threshold: 0.002, + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + reset_hint: 3, + timestamp: 4000, + start_timestamp: 2500, + ..Default::default() + }, + ], + ); + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = http_client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + let rows = client + .simple_query( + "select greptime_native_histogram from pg_native_histogram_seconds order by greptime_timestamp", + ) + .await + .unwrap(); + let jsons: Vec<&str> = rows + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + assert_eq!(jsons.len(), 2); + + // The struct column is encoded as a JSON object, matching the HTTP query + // output for the same histograms. Before the fix every row was encoded as + // an empty struct (`{}`). + let histogram: serde_json::Value = serde_json::from_str(jsons[0]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let histogram: serde_json::Value = serde_json::from_str(jsons[1]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + +pub async fn test_postgres_struct_types(store_type: StorageType) { + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_struct_types").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + async fn query_one(client: &Client, sql: &str) -> String { + let messages = client.simple_query(sql).await.unwrap(); + messages + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0).map(|v| v.to_string()), + _ => None, + }) + .next() + .unwrap() + } + + // Struct and list-of-struct columns are serialized as JSON. View-typed + // struct fields must not fail row extraction, and a null struct inside a + // list stays null instead of becoming a struct of null fields. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + // '0102' casts to its UTF-8 bytes when interpreted as BinaryView. + let row = query_one(&client, "SELECT struct(arrow_cast('0102', 'BinaryView'))").await; + assert_eq!(row, "{\"c0\":[48,49,48,50]}"); + + let row = query_one(&client, "SELECT struct([struct(1), NULL])").await; + assert_eq!(row, "{\"c0\":[{\"c0\":1},null]}"); + + // Top-level control: the same null element in a bare list stays null too. + let row = query_one(&client, "SELECT [struct(1), NULL]").await; + assert_eq!(row, "[{\"c0\":1},null]"); + + // Unsupported arrow field types surface as query errors instead of + // dropping the connection. + let error = client + .simple_query("SELECT struct(arrow_cast('1', 'Decimal256(38, 10)'))") + .await + .unwrap_err(); + let message = match error.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => error.to_string(), + }; + assert!( + message.contains("Unsupported arrow data type"), + "unexpected error message: {message}" + ); + + // The connection stays usable after the error. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_slow_query(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server_with_slow_query_threshold( store_type, From 9281965bacf6c7ef867a508199643d186f6005b5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 23 Sep 2026 04:20:02 +0000 Subject: [PATCH 014/133] feat(flow): allow execution hook to rewrite completed plans (#9290) * refactor(flow): expose incremental aggregate plan analysis Signed-off-by: discord9 * feat(flow): allow execution hook to rewrite completed plans Signed-off-by: discord9 * test(flow): cover the execution hook receiving the completed plan The hook must see the plan that is dispatched after the incremental delta-sink merge, so the test records the plan a collaborator is handed and asserts the merge is already part of it. Signed-off-by: discord9 --------- Signed-off-by: discord9 --- .../src/batching_mode/batching_execution.rs | 6 ++ src/flow/src/batching_mode/task.rs | 4 + src/flow/src/batching_mode/task/test.rs | 93 +++++++++++++++++++ src/flow/src/lib.rs | 5 +- 4 files changed, 107 insertions(+), 1 deletion(-) diff --git a/src/flow/src/batching_mode/batching_execution.rs b/src/flow/src/batching_mode/batching_execution.rs index 0a699a227be..298522b8137 100644 --- a/src/flow/src/batching_mode/batching_execution.rs +++ b/src/flow/src/batching_mode/batching_execution.rs @@ -16,6 +16,7 @@ use std::sync::Arc; +use datafusion_expr::LogicalPlan; use query::QueryEngineRef; use table::TableRef; @@ -37,6 +38,11 @@ pub trait BatchingExecution: Send + Sync + 'static { max_window_cnt: Option, ) -> ExecuteOnceOutcome; + /// Rewrite the completed query plan after incremental merging and before execution. + fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> Result { + Ok(plan) + } + /// Retire this execution instance, rejecting new work and requesting that any /// retained local work stop. This is not an acknowledgement of remote quiescence. fn stop(&self) {} diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index c0ee6d1c1b5..346d39d1d46 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -828,6 +828,10 @@ impl BatchingTask { return Ok(None); } let plan = incremental_plan.unwrap_or_else(|| plan.clone()); + let plan = match &self.execution { + Some(execution) => execution.rewrite_plan(self, plan)?, + None => plan, + }; let extensions = self .build_flow_query_extensions(incremental_safe, coverage.is_incremental_delta()) diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 7a56c70823d..7add59664a5 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -143,6 +143,99 @@ impl crate::BatchingExecution for RetainingExecution { } } +/// Records the completed plan handed to it by the task. +#[derive(Default)] +struct PlanRecordingExecution { + plans: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl crate::BatchingExecution for PlanRecordingExecution { + async fn execute_once( + self: Arc, + _guard: BatchingExecutionGuard, + _task: &BatchingTask, + _engine: &QueryEngineRef, + _frontend: &Arc, + _max_window_cnt: Option, + ) -> ExecuteOnceOutcome { + ExecuteOnceOutcome { + new_query: None, + result: Ok(None), + } + } + + fn rewrite_plan(&self, _task: &BatchingTask, plan: LogicalPlan) -> crate::Result { + self.plans.lock().unwrap().push(plan.clone()); + Ok(plan) + } +} + +#[tokio::test] +async fn test_execution_hook_receives_the_completed_incremental_plan() { + let sink_table = "hook_sink"; + let query = "SELECT max(number) AS number, ts FROM numbers_with_ts GROUP BY ts"; + let TestTaskParts { + task, query_engine, .. + } = new_test_task_engine_and_plan_with_query(query, sink_table).await; + register_auto_created_aggregate_sink(&query_engine, sink_table); + + let ctx = task.state.read().unwrap().query_ctx.clone(); + let plan = sql_to_df_plan(ctx, query_engine.clone(), query, true) + .await + .unwrap(); + let (sink, _) = get_table_info_df_schema( + query_engine.engine_state().catalog_manager().clone(), + [ + "greptime".to_string(), + "public".to_string(), + sink_table.to_string(), + ], + ) + .await + .unwrap(); + let table_provider = Arc::new(DfTableProviderAdapter::new(sink)); + let table_source = Arc::new(DefaultTableSource::new(table_provider)); + let dml_plan = LogicalPlan::Dml(DmlStatement::new( + datafusion_common::TableReference::bare("test"), + table_source, + WriteOp::Insert(datafusion_expr::dml::InsertOp::Append), + Arc::new(plan), + )); + + task.state + .write() + .unwrap() + .advance_checkpoints(HashMap::from([(1_u64, 10_u64)])); + + let execution = Arc::new(PlanRecordingExecution::default()); + let task = task.with_execution(Some(execution.clone())); + let (frontend_client, _handler) = + FrontendClient::from_empty_grpc_handler(QueryOptions::default()); + let frontend = Arc::new(frontend_client); + let dirty_restore = DirtyRestore::Unscoped(dirty_range(10, 15)); + + let _ = task + .execute_plan_unlocked( + &query_engine, + &frontend, + &dml_plan, + &dirty_restore, + &QueryCoverage::IncrementalDelta, + ) + .await; + + let recorded = execution.plans.lock().unwrap(); + let plan = recorded + .first() + .expect("the execution hook must see the plan that is dispatched"); + let plan_text = plan.to_string(); + assert!( + plan_text.contains("Left Join"), + "the hook must receive the completed delta-sink merge plan, got:\n{plan_text}" + ); +} + #[tokio::test] async fn test_execution_delegate_dispatch_is_serialized() { let TestTaskParts { diff --git a/src/flow/src/lib.rs b/src/flow/src/lib.rs index e9ef90e09d6..36dc7c21818 100644 --- a/src/flow/src/lib.rs +++ b/src/flow/src/lib.rs @@ -48,7 +48,10 @@ pub use batching_mode::task::{ QueryCoverage, TaskArgs, }; pub use batching_mode::time_window::{TimeWindowExpr, find_time_window_expr}; -pub use batching_mode::utils::sql_to_df_plan; +pub use batching_mode::utils::{ + IncrementalAggregateAnalysis, IncrementalAggregateMergeColumn, IncrementalAggregateMergeOp, + analyze_incremental_aggregate_plan, sql_to_df_plan, +}; pub(crate) use engine::{CreateFlowArgs, FlowId, TableName}; pub use error::{Error, Result}; pub use server::{FlownodeBuilder, FlownodeInstance, FlownodeServer, FlownodeServiceBuilder}; From 953d01ac546cf13937b8df8b6770fe9395dd11de Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 23 Sep 2026 04:41:13 +0000 Subject: [PATCH 015/133] feat: support pending rows batching for MySQL and PostgreSQL (#9302) * feat: support pending rows batching for MySQL and PostgreSQL Signed-off-by: WenyXu * style: group batcher imports before item definitions Signed-off-by: WenyXu * fix: use 65536 as the default batcher worker channel capacity Signed-off-by: WenyXu * test: use a distinct custom worker channel capacity Signed-off-by: WenyXu * test: complete Prom config in worker capacity override case Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- config/config.md | 4 +- config/frontend.example.toml | 6 +- config/standalone.example.toml | 6 +- src/frontend/src/frontend.rs | 11 +- src/frontend/src/instance/builder.rs | 2 +- src/frontend/src/instance/logical_batcher.rs | 2 +- src/frontend/src/server.rs | 62 ++++-- .../service_config/pending_rows_batcher.rs | 44 ++++- src/frontend/src/service_config/prom_store.rs | 6 +- src/servers/src/batcher.rs | 39 ++++ src/servers/src/http.rs | 35 +--- src/servers/src/mysql/handler.rs | 31 ++- src/servers/src/mysql/server.rs | 11 +- src/servers/src/postgres.rs | 9 +- src/servers/src/postgres/handler.rs | 12 +- src/servers/src/postgres/server.rs | 12 +- src/servers/tests/http/prom_store_test.rs | 3 +- src/session/src/context.rs | 2 +- src/standalone/src/options.rs | 11 +- tests-integration/src/otlp.rs | 3 +- tests-integration/src/standalone.rs | 11 +- tests-integration/tests/sql.rs | 181 ++++++++++++++++++ 22 files changed, 406 insertions(+), 97 deletions(-) diff --git a/config/config.md b/config/config.md index d736c24d848..c33918a8573 100644 --- a/config/config.md +++ b/config/config.md @@ -75,7 +75,7 @@ | `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` | -- | -- | 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` | -- | -- | Ordinary-table batching for opted-in ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher.
MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged.
Single-connection writes and INSERT SELECT may incur additional flush waits. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | @@ -343,7 +343,7 @@ | `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` | -- | -- | 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` | -- | -- | Ordinary-table batching for opted-in ingestion protocols.
PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge
queue admission without waiting for storage; later failures cannot be returned to the client.
Omitted or empty protocols disables batching. Prom without metric engine uses this batcher.
OTLP logs, traces and ordinary metrics use this batcher.
MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged.
Single-connection writes and INSERT SELECT may incur additional flush waits. | | `pending_rows_batcher.pending_rows_flush_interval` | String | `0s` | Flush interval measured from the first pending submission. Zero disables batching. | | `pending_rows_batcher.max_batch_rows` | Integer | `100000` | Flush after a complete submission reaches this row threshold. | | `pending_rows_batcher.max_concurrent_flushes` | Integer | `256` | Maximum concurrent flushes shared by the frontend batcher. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 275db3105c5..055acbb93bd 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -231,11 +231,13 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Ordinary-table batching for opted-in HTTP ingestion protocols. +## Ordinary-table batching for opted-in ingestion protocols. ## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge ## queue admission without waiting for storage; later failures cannot be returned to the client. ## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. ## OTLP logs, traces and ordinary metrics use this batcher. +## MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged. +## Single-connection writes and INSERT SELECT may incur additional flush waits. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -246,6 +248,8 @@ default_merge_mode = "last_non_null" # "splunk", # "elasticsearch", # "http_sql", +# "mysql", +# "postgres", # "prom", # ] ## Flush interval measured from the first pending submission. Zero disables batching. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 257f48eda2c..357ec5904f7 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -210,11 +210,13 @@ enable = true ## Available values: "last_non_null", "last_row". default_merge_mode = "last_non_null" -## Ordinary-table batching for opted-in HTTP ingestion protocols. +## Ordinary-table batching for opted-in ingestion protocols. ## PENDING_ROWS_BATCH_SYNC defaults to true for both batchers. Set it to false to acknowledge ## queue admission without waiting for storage; later failures cannot be returned to the client. ## Omitted or empty protocols disables batching. Prom without metric engine uses this batcher. ## OTLP logs, traces and ordinary metrics use this batcher. +## MySQL and PostgreSQL use the same acknowledgement policy. Protocol timeouts are unchanged. +## Single-connection writes and INSERT SELECT may incur additional flush waits. [pending_rows_batcher] # protocols = [ # "influxdb", @@ -225,6 +227,8 @@ default_merge_mode = "last_non_null" # "splunk", # "elasticsearch", # "http_sql", +# "mysql", +# "postgres", # "prom", # ] ## Flush interval measured from the first pending submission. Zero disables batching. diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 72bf322b1f1..51e291c7906 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -206,6 +206,7 @@ mod tests { use futures::Stream; use meta_client::MetaClientRef; use meta_client::client::MetaClientBuilder; + use servers::batcher::BatchingProtocol; use servers::grpc::{FlightCompression, GRPC_SERVER}; use servers::http::HTTP_SERVER; use servers::http::result::greptime_result_v1::GreptimedbV1Response; @@ -249,10 +250,7 @@ mod tests { .logical_table .unwrap() .protocols, - vec![ - servers::http::BatchingProtocol::Prom, - servers::http::BatchingProtocol::Otlp - ] + vec![BatchingProtocol::Prom, BatchingProtocol::Otlp] ); }, ); @@ -270,10 +268,7 @@ mod tests { FrontendOptions::load_layered_options(None, "FRONTEND_BATCHER_TEST").unwrap(); assert_eq!( options.pending_rows_batcher.table.protocols, - vec![ - servers::http::BatchingProtocol::Influxdb, - servers::http::BatchingProtocol::HttpSql - ] + vec![BatchingProtocol::Influxdb, BatchingProtocol::HttpSql] ); }, ); diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index d9b86ac5a05..d261daed48e 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -50,8 +50,8 @@ use partition::manager::PartitionRuleManager; use pipeline::pipeline_operator::PipelineOperator; use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; +use servers::batcher::BatchingProtocol; use servers::batcher::table::TablePendingRowsBatcher; -use servers::http::BatchingProtocol; use snafu::{OptionExt, ResultExt}; use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs index 28ed19b3e14..e68dcf644d4 100644 --- a/src/frontend/src/instance/logical_batcher.rs +++ b/src/frontend/src/instance/logical_batcher.rs @@ -16,9 +16,9 @@ use std::sync::{Arc, Weak}; use api::v1::ColumnSchema; use async_trait::async_trait; +use servers::batcher::BatchingProtocol; use servers::batcher::logical_table::{LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer}; use servers::error::{BatcherChannelClosedSnafu, Result}; -use servers::http::BatchingProtocol; use session::context::QueryContextRef; use snafu::OptionExt; diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 0a6612ba31c..fb1ab2f2603 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -24,7 +24,7 @@ use common_base::Plugins; use common_config::Configurable; use common_telemetry::{info, warn}; use meta_client::MetaClientOptions; -use servers::batcher::pending_rows_batch_sync_enabled; +use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use servers::error::Error as ServerError; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::FlightCraftRef; @@ -34,7 +34,7 @@ use servers::grpc::{GrpcOptions, GrpcServer}; use servers::http::event::LogValidatorRef; use servers::http::result::error_result::ErrorResponse; use servers::http::utils::router::RouterConfigurator; -use servers::http::{BatchingProtocol, HttpOptions, HttpServer, HttpServerBuilder}; +use servers::http::{HttpOptions, HttpServer, HttpServerBuilder}; use servers::interceptor::LogIngestInterceptorRef; use servers::metrics_handler::MetricsHandler; use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef}; @@ -369,6 +369,15 @@ where } } + let table_batcher = opts.table_batcher_options(); + let batching_enabled = table_batcher.pending_rows_batching_enabled(); + let mysql_batching = + batching_enabled && table_batcher.protocols.contains(&BatchingProtocol::Mysql); + let postgres_batching = batching_enabled + && table_batcher + .protocols + .contains(&BatchingProtocol::Postgres); + if opts.mysql.enable { // Init MySQL server let opts = &opts.mysql; @@ -384,13 +393,16 @@ where let mysql_server = MysqlServer::create_server( common_runtime::global_runtime(), Arc::new(MysqlSpawnRef::new(instance.clone(), user_provider.clone())), - Arc::new(MysqlSpawnConfig::new( - opts.tls.should_force_tls(), - tls_server_config, - opts.keep_alive.as_secs(), - opts.reject_no_database.unwrap_or(false), - opts.prepared_stmt_cache_size, - )), + Arc::new( + MysqlSpawnConfig::new( + opts.tls.should_force_tls(), + tls_server_config, + opts.keep_alive.as_secs(), + opts.reject_no_database.unwrap_or(false), + opts.prepared_stmt_cache_size, + ) + .with_batching_enabled(mysql_batching), + ), Some(instance.process_manager().clone()), ); handlers.insert((mysql_server, mysql_addr)); @@ -407,15 +419,18 @@ where maybe_watch_server_tls_config(tls_server_config.clone()).context(StartServerSnafu)?; - let pg_server = Box::new(PostgresServer::new( - instance.clone(), - opts.tls.should_force_tls(), - tls_server_config, - opts.keep_alive.as_secs(), - common_runtime::global_runtime(), - user_provider.clone(), - Some(self.instance.process_manager().clone()), - )) as Box; + let pg_server = Box::new( + PostgresServer::new( + instance.clone(), + opts.tls.should_force_tls(), + tls_server_config, + opts.keep_alive.as_secs(), + common_runtime::global_runtime(), + user_provider.clone(), + Some(self.instance.process_manager().clone()), + ) + .with_batching_enabled(postgres_batching), + ) as Box; handlers.insert((pg_server, pg_addr)); } @@ -452,8 +467,11 @@ fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> let common_enabled = batch_sync && shared.pending_rows_batching_enabled() && shared.protocols.iter().any(|protocol| { - *protocol != BatchingProtocol::Prom - || (prom_store.enable && !prom_store.with_metric_engine) + !matches!( + protocol, + BatchingProtocol::Mysql | BatchingProtocol::Postgres + ) && (*protocol != BatchingProtocol::Prom + || (prom_store.enable && !prom_store.with_metric_engine)) }); let common_interval = common_enabled.then_some(shared.pending_rows_flush_interval); let prom_interval = (prom_store.pending_rows_batching_enabled() && batch_sync) @@ -655,6 +673,10 @@ mod tests { expected_secs, ) in [ (vec![BatchingProtocol::Prom], false, false, 5, 2, 1, 1), + (vec![BatchingProtocol::Mysql], false, false, 5, 0, 1, 1), + (vec![BatchingProtocol::Postgres], false, false, 5, 0, 1, 1), + (vec![BatchingProtocol::Mysql], false, true, 5, 0, 1, 1), + (vec![BatchingProtocol::Postgres], false, true, 5, 0, 1, 1), (vec![BatchingProtocol::Prom], true, false, 5, 2, 1, 1), (vec![BatchingProtocol::Prom], true, true, 5, 2, 1, 6), (vec![BatchingProtocol::Influxdb], true, false, 5, 2, 1, 1), diff --git a/src/frontend/src/service_config/pending_rows_batcher.rs b/src/frontend/src/service_config/pending_rows_batcher.rs index 61572015d7c..b671ee9a272 100644 --- a/src/frontend/src/service_config/pending_rows_batcher.rs +++ b/src/frontend/src/service_config/pending_rows_batcher.rs @@ -17,7 +17,7 @@ use std::time::Duration; use common_batcher::flush_policy::timing::TimingFlushPolicy; use serde::{Deserialize, Serialize}; -use servers::http::BatchingProtocol; +use servers::batcher::BatchingProtocol; use tokio::sync::Semaphore; use crate::frontend::FrontendOptions; @@ -38,11 +38,11 @@ pub struct PendingRowsBatcherOptions { pub logical_table: Option, } -/// Write batching controls shared by HTTP ingestion protocols. +/// Write batching controls shared by ingestion protocols. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct BatcherOptions { - /// HTTP write protocols sharing this batcher; empty disables all entrances. + /// Write protocols sharing this batcher; empty disables all entrances. pub protocols: Vec, /// Time from the first pending submission to a timed flush. Zero disables batching. #[serde(with = "humantime_serde")] @@ -238,6 +238,8 @@ mod tests { "opentsdb", "elasticsearch", "splunk", + "mysql", + "postgres", ] { assert!( toml::from_str::(&format!( @@ -251,9 +253,11 @@ mod tests { #[test] fn test_protocols() { let options: BatcherOptions = toml::from_str( - "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom']", + "protocols = ['influxdb', 'opentsdb', 'otlp', 'logs', 'loki', 'splunk', 'elasticsearch', 'http_sql', 'prom', 'mysql', 'postgres']", ).unwrap(); - assert_eq!(options.protocols.len(), 9); + assert_eq!(options.protocols.len(), 11); + assert!(options.protocols.contains(&BatchingProtocol::Mysql)); + assert!(options.protocols.contains(&BatchingProtocol::Postgres)); assert!(options.protocols.contains(&BatchingProtocol::HttpSql)); assert!(BatcherOptions::default().protocols.is_empty()); for invalid in ["sql", "jaeger", "unknown"] { @@ -289,6 +293,36 @@ mod tests { ); } + #[test] + fn test_worker_capacity_override() { + let options: FrontendOptions = toml::from_str( + r#" +[pending_rows_batcher] +worker_channel_capacity = 12345 +[pending_rows_batcher.logical_table] +worker_channel_capacity = 12345 +[prom_store] +enable = true +with_metric_engine = true +worker_channel_capacity = 12345 +"#, + ) + .unwrap(); + assert_eq!( + options.pending_rows_batcher.table.worker_channel_capacity, + 12_345 + ); + assert_eq!( + options + .pending_rows_batcher + .logical_table + .unwrap() + .worker_channel_capacity, + 12_345 + ); + assert_eq!(options.prom_store.worker_channel_capacity, 12_345); + } + #[test] fn test_partial_options_and_zero_controls() { let options: BatcherOptions = diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs index 2fdd05043db..3d00f48cdda 100644 --- a/src/frontend/src/service_config/prom_store.rs +++ b/src/frontend/src/service_config/prom_store.rs @@ -108,7 +108,6 @@ mod tests { use crate::service_config::prom_store::{ default_flow_notification_queue_capacity, default_max_batch_rows, default_max_concurrent_flushes, default_max_inflight_requests, - default_worker_channel_capacity, }; #[test] @@ -160,10 +159,7 @@ mod tests { default.max_concurrent_flushes, default_max_concurrent_flushes() ); - assert_eq!( - default.worker_channel_capacity, - default_worker_channel_capacity() - ); + assert_eq!(default.worker_channel_capacity, 65_536); assert_eq!( default.max_inflight_requests, default_max_inflight_requests() diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs index 6b6ca8f7ee7..3cd9a831aaa 100644 --- a/src/servers/src/batcher.rs +++ b/src/servers/src/batcher.rs @@ -22,6 +22,8 @@ pub mod table; #[cfg(test)] mod test_util; +use serde::{Deserialize, Serialize}; + /// Controls whether batching waits for storage before replying to the client. const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC"; @@ -39,3 +41,40 @@ pub fn pending_rows_batch_sync_enabled() -> bool { .and_then(|v| v.parse::().ok()) .unwrap_or(true) } + +/// Ingestion protocols that can opt into pending-row batching. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BatchingProtocol { + Prom, + Influxdb, + Opentsdb, + Otlp, + Logs, + Loki, + Splunk, + Elasticsearch, + HttpSql, + Mysql, + Postgres, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_protocol_names_reject_unknown_values() { + assert_eq!( + serde_json::from_str::("\"prom\"").unwrap(), + BatchingProtocol::Prom + ); + for name in ["sql", "unknown"] { + assert!(serde_json::from_str::(&format!("\"{name}\"")).is_err()); + } + assert_eq!( + serde_json::from_str::("\"http_sql\"").unwrap(), + BatchingProtocol::HttpSql + ); + } +} diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 4d168111ade..f947ec1621b 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -53,6 +53,7 @@ use tower_http::trace::TraceLayer; use self::authorize::AuthState; use self::result::table_result::TableResponse; +use crate::batcher::BatchingProtocol; use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; use crate::elasticsearch; use crate::error::{ @@ -189,22 +190,6 @@ pub(crate) enum HttpServerKind { Api, } -/// HTTP write protocols eligible for the shared pending-row batcher. -/// Prometheus uses this selector only when metric-engine storage is disabled. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum BatchingProtocol { - Prom, - Influxdb, - Opentsdb, - Otlp, - Logs, - Loki, - Splunk, - Elasticsearch, - HttpSql, -} - #[derive(Default)] pub struct HttpServer { router: StdMutex, @@ -2478,9 +2463,10 @@ mod batching_tests { use common_query::Output; use session::context::QueryContextRef; + use crate::batcher::BatchingProtocol; use crate::error::Result as ServerResult; use crate::http::test_helpers::TestClient; - use crate::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; + use crate::http::{HttpOptions, HttpServerBuilder}; use crate::influxdb::InfluxdbRequest; use crate::opentsdb::codec::DataPoint; use crate::query_handler::{InfluxdbLineProtocolHandler, OpentsdbProtocolHandler}; @@ -2523,21 +2509,6 @@ mod batching_tests { } } - #[test] - fn test_protocol_names_reject_unknown_values() { - assert_eq!( - serde_json::from_str::("\"prom\"").unwrap(), - BatchingProtocol::Prom - ); - for name in ["sql", "unknown"] { - assert!(serde_json::from_str::(&format!("\"{name}\"")).is_err()); - } - assert_eq!( - serde_json::from_str::("\"http_sql\"").unwrap(), - BatchingProtocol::HttpSql - ); - } - #[derive(Default)] struct RecordingWriteHandler { selections: Mutex>, diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index 3764591beaa..d0dd1f25279 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -85,6 +85,7 @@ pub struct MysqlInstanceShim { prepared_stmts_counter: AtomicU32, process_id: u32, prepared_stmt_cache_size: usize, + batching_enabled: bool, } impl MysqlInstanceShim { @@ -122,9 +123,22 @@ impl MysqlInstanceShim { prepared_stmts_counter: AtomicU32::new(1), process_id, prepared_stmt_cache_size, + batching_enabled: false, } } + /// Enables ordinary-table batching for this connection. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + + fn new_query_context(&self) -> QueryContextRef { + let mut ctx = self.session.new_query_context(); + Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); + ctx + } + #[tracing::instrument(skip_all, name = "mysql::do_query")] async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec> { if let Some(output) = @@ -495,7 +509,7 @@ impl AsyncMysqlShim for MysqlInstanceShi raw_query: &'a str, w: StatementMetaWriter<'a, W>, ) -> Result<()> { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let stmt_id = self.prepared_stmts_counter.fetch_add(1, Ordering::Relaxed); let stmt_key = uuid::Uuid::from_u128(stmt_id as u128).to_string(); let (params, columns) = match self @@ -525,7 +539,7 @@ impl AsyncMysqlShim for MysqlInstanceShi ) -> Result<()> { self.session.clear_warnings(); - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_MYSQL_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_MYSQL_BINQUERY, db.as_str()]) @@ -569,7 +583,7 @@ impl AsyncMysqlShim for MysqlInstanceShi query: &'a str, writer: QueryResultWriter<'a, W>, ) -> Result<()> { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_MYSQL_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_MYSQL_TEXTQUERY, db.as_str()]) @@ -1039,6 +1053,17 @@ mod tests { ) } + #[test] + fn test_batching_context() { + for enabled in [false, true] { + let shim = create_shim().with_batching_enabled(enabled); + let ctx = shim.new_query_context(); + assert_eq!(ctx.batching_enabled(), enabled); + assert!(!ctx.logical_batching_enabled()); + assert_eq!(ctx.channel(), Channel::Mysql); + } + } + fn statement_with_transformed_placeholders(query: &str) -> Statement { let mut statements = ParserContext::create_with_dialect(query, &MySqlDialect {}, ParseOptions::default()) diff --git a/src/servers/src/mysql/server.rs b/src/servers/src/mysql/server.rs index 0f48cb0d598..811fc9827ec 100644 --- a/src/servers/src/mysql/server.rs +++ b/src/servers/src/mysql/server.rs @@ -86,6 +86,7 @@ pub struct MysqlSpawnConfig { reject_no_database: bool, // prepared statement cache capacity prepared_stmt_cache_size: usize, + batching_enabled: bool, } impl MysqlSpawnConfig { @@ -102,9 +103,16 @@ impl MysqlSpawnConfig { keep_alive_secs, reject_no_database, prepared_stmt_cache_size, + batching_enabled: false, } } + /// Enables ordinary-table batching for connections accepted by this server. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + fn tls(&self) -> Option> { self.tls.get_config() } @@ -213,7 +221,8 @@ impl MysqlServer { stream.peer_addr()?, process_id, spawn_config.prepared_stmt_cache_size, - ); + ) + .with_batching_enabled(spawn_config.batching_enabled); let (mut r, w) = stream.into_split(); let mut w = BufWriter::with_capacity(DEFAULT_RESULT_SET_WRITE_BUFFER_SIZE, w); diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 58ef4fdd7b7..1cd5907157b 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -80,6 +80,7 @@ pub struct PostgresServerHandlerInner { force_tls: bool, param_provider: Arc, + batching_enabled: bool, session: Arc, query_parser: Arc, } @@ -114,7 +115,12 @@ impl PgWireServerHandlers for PostgresServerHandler { } impl MakePostgresServerHandler { - fn make(&self, addr: Option, process_id: u32) -> PostgresServerHandler { + fn make( + &self, + addr: Option, + process_id: u32, + batching_enabled: bool, + ) -> PostgresServerHandler { let session = Arc::new(Session::new( addr, Channel::Postgres, @@ -127,6 +133,7 @@ impl MakePostgresServerHandler { force_tls: self.force_tls, param_provider: self.param_provider.clone(), + batching_enabled, session: session.clone(), query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)), }; diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index e618c55614b..7bf1edca705 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -59,6 +59,14 @@ use crate::postgres::utils::convert_err; use crate::postgres::{PostgresServerHandlerInner, fixtures}; use crate::query_handler::sql::ServerSqlQueryHandlerRef; +impl PostgresServerHandlerInner { + fn new_query_context(&self) -> QueryContextRef { + let mut ctx = self.session.new_query_context(); + Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); + ctx + } +} + #[async_trait] impl SimpleQueryHandler for PostgresServerHandlerInner { #[tracing::instrument(skip_all, fields(protocol = "postgres"))] @@ -68,7 +76,7 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { C::Error: Debug, PgWireError: From<>::Error>, { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_POSTGRES_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_POSTGRES_SIMPLE_QUERY, db.as_str()]) @@ -425,7 +433,7 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner { C::Error: Debug, PgWireError: From<>::Error>, { - let query_ctx = self.session.new_query_context(); + let query_ctx = self.new_query_context(); let db = query_ctx.get_db_string(); let _timer = crate::metrics::METRIC_POSTGRES_QUERY_TIMER .with_label_values(&[crate::metrics::METRIC_POSTGRES_EXTENDED_QUERY, db.as_str()]) diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 09bc5015a0b..6e4aa942f9d 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -38,6 +38,7 @@ pub struct PostgresServer { make_handler: Arc, tls_server_config: Arc, keep_alive_secs: u64, + batching_enabled: bool, bind_addr: Option, process_manager: Option, } @@ -66,17 +67,25 @@ impl PostgresServer { make_handler, tls_server_config, keep_alive_secs, + batching_enabled: false, bind_addr: None, process_manager, } } + /// Enables ordinary-table batching for connections accepted by this server. + pub fn with_batching_enabled(mut self, enabled: bool) -> Self { + self.batching_enabled = enabled; + self + } + fn accept( &self, io_runtime: Runtime, accepting_stream: AbortableStream, ) -> impl Future + use<> { let handler_maker = self.make_handler.clone(); + let batching_enabled = self.batching_enabled; let tls_server_config = self.tls_server_config.clone(); let process_manager = self.process_manager.clone(); accepting_stream.for_each(move |tcp_stream| { @@ -102,7 +111,8 @@ impl PostgresServer { let _handle = io_runtime.spawn(async move { crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc(); - let pg_handler = Arc::new(handler_maker.make(addr, process_id)); + let pg_handler = + Arc::new(handler_maker.make(addr, process_id, batching_enabled)); let r = process_socket(io_stream, tls_acceptor.clone(), pg_handler).await; crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec(); diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index 7c07e3b0344..ffd2c9dcf7f 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -36,11 +36,12 @@ use datafusion_expr::LogicalPlan; use prost::Message; use query::parser::PromQuery; use query::query_engine::DescribeResult; +use servers::batcher::BatchingProtocol; use servers::error::{self, Result}; use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::http::test_helpers::{TestClient, TestResponse}; -use servers::http::{BatchingProtocol, HttpOptions, HttpServerBuilder}; +use servers::http::{HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store; diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 5ebf8b36d21..ec10dc6cf3a 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -471,7 +471,7 @@ impl QueryContext { self.logical_batching_enabled = enabled; } - /// Whether the local HTTP entry point selected ordinary-table batching. + /// Whether the local protocol entry point selected ordinary-table batching. pub fn batching_enabled(&self) -> bool { self.batching_enabled } diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index d0de7408789..fbfcc2ccc1d 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -217,6 +217,7 @@ mod tests { use std::sync::Arc; use common_event_recorder::EventTypeFilter; + use servers::batcher::BatchingProtocol; use crate::options::*; @@ -266,10 +267,7 @@ mod tests { .logical_table .unwrap() .protocols, - vec![ - servers::http::BatchingProtocol::Prom, - servers::http::BatchingProtocol::Otlp - ] + vec![BatchingProtocol::Prom, BatchingProtocol::Otlp] ); }, ); @@ -288,10 +286,7 @@ mod tests { .unwrap(); assert_eq!( options.pending_rows_batcher.table.protocols, - vec![ - servers::http::BatchingProtocol::Influxdb, - servers::http::BatchingProtocol::HttpSql - ] + vec![BatchingProtocol::Influxdb, BatchingProtocol::HttpSql] ); }, ); diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 338b03ff68f..9fbf80fbeea 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -514,8 +514,7 @@ WITH( ExponentialHistogram, ExponentialHistogramDataPoint, exponential_histogram_data_point, }; use prost::Message; - use servers::batcher::pending_rows_batch_sync_enabled; - use servers::http::BatchingProtocol; + use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use servers::http::test_helpers::TestClient; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 38f546e9ccb..6fc8ce231af 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -90,6 +90,7 @@ pub struct GreptimeDbStandaloneBuilder { auto_create_table: bool, experimental_metric_export: bool, logical_batcher: Option, + table_batcher: BatcherOptions, } impl GreptimeDbStandaloneBuilder { @@ -111,6 +112,7 @@ impl GreptimeDbStandaloneBuilder { auto_create_table: true, experimental_metric_export: false, logical_batcher: None, + table_batcher: BatcherOptions::default(), } } @@ -121,6 +123,13 @@ impl GreptimeDbStandaloneBuilder { self } + /// Configures ordinary-table batching for protocol integration tests. + #[must_use] + pub fn with_table_batcher(mut self, options: BatcherOptions) -> Self { + self.table_batcher = options; + self + } + /// Configures logical-table batching for integration tests. #[must_use] pub fn with_logical_batcher(mut self, options: BatcherOptions) -> Self { @@ -389,7 +398,7 @@ impl GreptimeDbStandaloneBuilder { experimental_metric_export: self.experimental_metric_export, pending_rows_batcher: PendingRowsBatcherOptions { logical_table: self.logical_batcher.clone(), - ..Default::default() + table: self.table_batcher.clone(), }, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index e3b66fc4447..cb3d151fa85 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -2672,3 +2672,184 @@ pub async fn test_declare_fetch_close_cursor(store_type: StorageType) { let _ = fe_pg_server.shutdown().await; guard.remove_all().await; } + +/// Keeps SQL and input rows fixed while changing only the selected batching protocols. +#[tokio::test(flavor = "multi_thread")] +async fn test_sql_batcher_alignment() { + use common_telemetry::{dump_metrics, info, init_default_ut_logging}; + use frontend::server::Services; + use frontend::service_config::BatcherOptions; + use servers::batcher::BatchingProtocol; + use servers::mysql::server::MYSQL_SERVER; + use tests_integration::standalone::GreptimeDbStandaloneBuilder; + + fn flushes() -> u64 { + dump_metrics() + .unwrap() + .lines() + .find_map(|line| { + line.strip_prefix("greptime_table_batcher_flush_total ") + .map(|value| value.parse().unwrap()) + }) + .unwrap_or(0) + } + + init_default_ut_logging(); + for protocols in [ + vec![], + vec![BatchingProtocol::Mysql], + vec![BatchingProtocol::Postgres], + vec![BatchingProtocol::Mysql, BatchingProtocol::Postgres], + ] { + let mysql_enabled = protocols.contains(&BatchingProtocol::Mysql); + let pg_enabled = protocols.contains(&BatchingProtocol::Postgres); + let mut instance = GreptimeDbStandaloneBuilder::new("sql_batcher_alignment") + .with_table_batcher(BatcherOptions { + protocols, + pending_rows_flush_interval: Duration::from_millis(10), + ..Default::default() + }) + .build() + .await; + let mut opts = instance.opts.clone(); + opts.http.addr = "127.0.0.1:0".into(); + opts.grpc.bind_addr = "127.0.0.1:0".into(); + opts.mysql.addr = "127.0.0.1:0".into(); + opts.postgres.addr = "127.0.0.1:0".into(); + let mut servers = Services::new(opts, instance.fe_instance().clone(), Default::default()) + .build() + .unwrap(); + servers.start_all().await.unwrap(); + let mysql_url = format!("mysql://{}/public", servers.addr(MYSQL_SERVER).unwrap()); + let mut mysql = MySqlConnection::connect(&mysql_url).await.unwrap(); + let (pg, connection) = tokio_postgres::connect( + &format!( + "postgres://{}/public", + servers.addr("POSTGRES_SERVER").unwrap() + ), + NoTls, + ) + .await + .unwrap(); + let pg_task = tokio::spawn(async move { connection.await.unwrap() }); + + for round in 0..2 { + mysql.execute("CREATE TABLE batch_sql (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val INT DEFAULT 7)").await.unwrap(); + let started = std::time::Instant::now(); + let before = flushes(); + // COM_QUERY and COM_STMT_EXECUTE must both opt in. + assert_eq!( + mysql + .execute("INSERT INTO batch_sql (ts, host) VALUES (0, 'a')") + .await + .unwrap() + .rows_affected(), + 1 + ); + assert_eq!( + sqlx::query("INSERT INTO batch_sql VALUES (0, ?, ?)") + .persistent(false) + .bind("b") + .bind(8_i32) + .execute(&mut mysql) + .await + .unwrap() + .rows_affected(), + 1 + ); + assert_eq!(flushes() - before, if mysql_enabled { 2 } else { 0 }); + + let before = flushes(); + let result = pg + .simple_query("INSERT INTO batch_sql (ts, host) VALUES (0, 'c')") + .await + .unwrap(); + assert!(matches!( + result.as_slice(), + [SimpleQueryMessage::CommandComplete(1)] + )); + let statement = pg + .prepare("INSERT INTO batch_sql VALUES (0, $1, $2)") + .await + .unwrap(); + assert_eq!(pg.execute(&statement, &[&"d", &9_i32]).await.unwrap(), 1); + assert_eq!(flushes() - before, if pg_enabled { 2 } else { 0 }); + + // Independent connections can submit to the same worker concurrently. + let (mysql_result, pg_result) = tokio::join!( + mysql.execute("INSERT INTO batch_sql VALUES (0, 'e', 10)"), + pg.simple_query("INSERT INTO batch_sql VALUES (0, 'f', 11)") + ); + assert_eq!(mysql_result.unwrap().rows_affected(), 1); + assert!(matches!( + pg_result.unwrap().as_slice(), + [SimpleQueryMessage::CommandComplete(1)] + )); + + // Errors must reach the client without affecting later requests. + assert!( + mysql + .execute("INSERT INTO batch_sql (missing) VALUES (1)") + .await + .is_err() + ); + assert!( + pg.simple_query("INSERT INTO batch_sql (missing) VALUES (1)") + .await + .is_err() + ); + let expected = vec![ + ("a".to_string(), 7_i32), + ("b".to_string(), 8), + ("c".to_string(), 7), + ("d".to_string(), 9), + ("e".to_string(), 10), + ("f".to_string(), 11), + ]; + // No polling: successful protocol responses guarantee visibility. + let rows: Vec<(String, i32)> = + sqlx::query_as("SELECT host, val FROM batch_sql ORDER BY host") + .persistent(false) + .fetch_all(&mut mysql) + .await + .unwrap(); + assert_eq!(rows, expected); + let rows: Vec<(String, i32)> = pg + .query("SELECT host, val FROM batch_sql ORDER BY host", &[]) + .await + .unwrap() + .iter() + .map(|row| (row.get(0), row.get(1))) + .collect(); + assert_eq!(rows, expected); + + mysql.execute("CREATE TABLE batch_copy (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val INT)").await.unwrap(); + assert_eq!( + mysql + .execute("INSERT INTO batch_copy SELECT * FROM batch_sql") + .await + .unwrap() + .rows_affected(), + 6 + ); + let rows: Vec<(String, i32)> = + sqlx::query_as("SELECT host, val FROM batch_copy ORDER BY host") + .persistent(false) + .fetch_all(&mut mysql) + .await + .unwrap(); + assert_eq!(rows, expected); + info!( + "sql_batcher_alignment mysql={mysql_enabled} postgres={pg_enabled} round={round} elapsed={:?}", + started.elapsed() + ); + mysql.execute("DROP TABLE batch_sql").await.unwrap(); + mysql.execute("DROP TABLE batch_copy").await.unwrap(); + } + mysql.close().await.unwrap(); + drop(pg); + pg_task.await.unwrap(); + servers.shutdown_all().await.unwrap(); + instance.guard.remove_all().await; + } +} From 1d8d95c12a61f33bf4116c1f3de940ad4eab2b87 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 04:56:44 +0000 Subject: [PATCH 016/133] chore(deps): replace cargo-udeps with cargo-shear for unused dependency checks (#9294) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(deps): replace cargo-udeps with cargo-shear for unused dependency checks cargo-udeps requires a nightly toolchain and its pinned version (0.1.61) no longer detects unused dependencies against current cargo internals — unused deps have landed on main undetected (e.g. humantime in common-frontend since #6689). cargo-shear is a standalone static analyzer that runs on any toolchain. - Swap 'make check-udeps' / 'make fix-udeps' recipes to 'cargo shear' / 'cargo shear --fix' and retire scripts/fix-udeps.py - CI: install cargo-shear in the check-udeps job; drop the build cache and protoc steps (cargo-shear never compiles) - Remove ~150 unused dependency declarations found by cargo-shear, move misplaced deps to the correct sections, drop orphaned [workspace.dependencies] entries (arrow-cast, rustc-hash) - Add [package.metadata.cargo-shear] ignored entries with explanations for dependencies that are structurally required despite no textual reference: sqlparser (required by sqlparser_derive expansions in datatypes, common-query), common-error (required by common-macro's stack_trace_debug expansions in session, tests-fuzz), k8s-openapi (feature-pinning for the transitive kube dependency in tests-fuzz), tikv-jemalloc-sys (link-only, enables jemalloc profiling features in common-mem-prof), protobuf (required by build.rs-generated bindings in log-store) - Drop the obsolete [package.metadata.cargo-udeps.ignore] sections Part of #9289 Signed-off-by: Ning Sun * fix(meta): populate physical metric table column ids (#9286) * fix(meta): populate physical metric table column ids Signed-off-by: dhruvxvaishnav * test(meta): verify physical metric column ids Signed-off-by: dhruvxvaishnav --------- Signed-off-by: dhruvxvaishnav * fix(postgres): return empty responses for comment-only SQL (#9295) fix(postgres): handle parsed empty queries in both protocols Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com> * ci: create docs follow-up issue on PR merge instead of on label (#9237) * ci: create docs follow-up issue on PR merge instead of on label The docbot workflow previously created a docs-repo issue as soon as the 'docs-required' condition was detected (PR opened/edited with the docs checkbox ticked), even if the PR was never merged. Now the workflow also triggers on PR 'closed': - opened/edited: only manage the docs-required/docs-not-required labels - closed: create the docs issue only when the PR was actually merged and carries the docs-required label This also lets maintainers control issue creation by manually adding or removing the docs-required label before merging. Signed-off-by: Ning Sun * fix: address review comments on docs issue creation timing - Only touch docs labels when the docs checkbox state actually changed in an edit. Previously, editing any other part of the PR body while the checkbox stayed checked removed the docs-required label, silently dropping the docs follow-up now that issue creation happens at merge. Unchanged checkbox now leaves labels untouched, which also preserves manual label overrides. - Do not trust the closed event's stale label snapshot at merge time: re-read the live PR via the API and create the docs issue if the docs-required label is present OR the checkbox is ticked in the current body. - Make the workflow concurrency group action-aware so a merge run does not cancel an in-flight label update from an edit run. Signed-off-by: Ning Sun * fix: make docs-required label the single source of truth at merge The label-OR-checkbox merge condition could not distinguish an intentional opt-out from an unfinished label update: removing docs-required while the checkbox stayed checked still produced an issue, and unchecking the box could still produce one if the merge read the stale label before the edit run removed it. At merge time, wait for any pending docbot runs on the PR head SHA to finish their label updates (bounded to 5 minutes), then decide solely by the live docs-required label. Adds actions: read permission for listing workflow runs. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun * perf(promql): push label filters into grouped join inputs (#9280) * perf(promql): propagate matching filters through grouped joins Signed-off-by: Dennis Zhuang * perf(promql): check matcher safety on the receiving operand Signed-off-by: Dennis Zhuang * refactor(promql): spell out the shapes a filter may cross `preserves_filter` ended in `_ => true`, which was only sound because `selector_matchers` independently rejects label rewriting, `count_values`, subqueries and non-rollup calls on the same operand. Loosening the latter alone would have silently pushed a matcher below a label rewrite. List the shapes that carry a scan filter instead and default to `false`. Cite #9207 for the result labels the grouped cases record: the join projects the right operand's tag set, so `zone` is missing wherever the right side aggregates it away. No behavior change. Signed-off-by: Dennis Zhuang * test(promql): assert the new pushdowns reach the scan The grouped-join unit tests feed tag columns by hand and the SQLness case only checks results, which are identical whether or not the rewrite fires. Nothing would have failed if scalar arithmetic, ranking or grouped matching stopped propagating. Assert through the planner that the matcher reaches both scans, with a global topk one-side as the counter-example. Also state that the duplicate-one-side cases record a cross product Prometheus rejects (#9209), so the baseline is not read as intended semantics. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang * fix(ci): build tests-integration lib with meta-srv/mock (#9299) * fix(ci): build tests-integration lib with meta-srv/mock tests-integration's lib code (src/cluster.rs) uses meta_srv::mocks, but the dependency carrying the mock feature sits in [dev-dependencies]. Builds that only touch the lib, such as the apidoc job's cargo doc --workspace, resolve meta-srv without mock and fail with E0432. --all-targets builds unify dev-dependency features, which is why check, clippy and nextest stayed green. Move the mock-enabled meta-srv entry back to [dependencies]. The other testing features moved out in #9072 are not needed by the lib and stay in [dev-dependencies]. Signed-off-by: Dennis Zhuang * test(repartition): split per-case repartition tests test_repartition_metric ran four format/primary-key-encoding cases in a single test function, and test_repartition_mito ran two format cases. Each case builds its own 3-datanode cluster and runs a full repartition plus GC cycle, so on S3 the metric test took 165-178s against the 180s nextest terminate-after. Merge queue runs failed on it at random. Split each case into its own test. Cases were already independent, so they now run in parallel and each stays far inside the timeout, and a failure points at one encoding instead of four. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang * feat(json2): support altering JSON2 settings (#9029) * feat(sql): support alter syntax for JSON2 columns Signed-off-by: fys * fix(json2): preserve rows on type hint mismatch during compaction * refactor(json2): simplify alter settings handling * fix(json2): preserve coerced values during compaction * chore: remove unnecessary clone * chor: reduce memory allocations * fix: cargo clippy * chore: update greptime-proto to main branch * refactor(datatypes): unify string handling with other JSON type hints * fix: cr --------- Signed-off-by: fys * fix: keep compaction pruning, metadata, and index work on compact runtime (#9304) * fix: run compaction pruner tasks on compact runtime Signed-off-by: Lei, HUANG * fix: keep compaction metadata and index work on compact runtime Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG * feat: add AI matching, classification, and scoring functions (#9300) * feat: return matching scores from jev Replace the experimental three-argument Boolean function with jev(text, prompt) returning a Float64 probability in [0, 1]. Move threshold comparisons into SQL and update tests and migration examples. Signed-off-by: Lei, HUANG * feat: add Jev choice and score functions Share asynchronous execution across Noul, Choice, and Score. Validate JSON criteria before requests and return typed scalar answers. Add SQL and HTTP mock coverage with usage examples. Signed-off-by: Lei, HUANG * refactor: use generic AI SQL function names Expose ai_match, ai_choose, and ai_score and move their implementation, tests, and usage guide under generic AI names. Document the current unreleased interface without migration history. Signed-off-by: Lei, HUANG * fix: share constant AI criteria within each batch Borrow scalar string arguments and lazily parse constant criteria once per batch. Share the parsed allocation across requests while preserving NULL propagation and batch validation before HTTP calls. Signed-off-by: Lei, HUANG * feat: preserve AI score uncertainty in JSONB results Return score, confidence, and probabilities in criteria-level order from one evaluation. Validate the distribution and preserve provider precision. Add JSON extraction, uncertainty, and single-request regressions, and document confidence-aware ranking. Signed-off-by: Lei, HUANG * docs: explain reuse of volatile AI evaluations Document repeated SELECT and WHERE evaluation costs as N + M requests, and show subquery aliases for reusing scalar or structured AI results without additional model calls. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG * feat: share logical table batching with OTLP metrics (#9288) * feat: share logical table batching with OTLP metrics Signed-off-by: WenyXu * fix: unify pending rows batch acknowledgement policy Signed-off-by: WenyXu * fix: align logical batcher example configuration expectations Signed-off-by: WenyXu * fix: align batcher worker channel defaults to 65536 Signed-off-by: WenyXu --------- Signed-off-by: WenyXu * perf(mito2): lazily decode dense primary key columns (#9226) * perf(mito2): lazily decode dense primary key columns Signed-off-by: Lei, HUANG * perf(mito2): bypass lazy decoding for full primary keys Signed-off-by: Lei, HUANG * fix(mito-codec): preserve prefix decoding errors Signed-off-by: Lei, HUANG * refactor(mito-codec): align encoded length helper naming Rename encoded_length to encoded_len and update all callers to match the other length helpers in the module. Signed-off-by: Lei, HUANG * refactor(mito2): clarify conditional dense key decoding Rename decode_dense_pk to ensure_dense_pk_decoded so callers can see that existing decoded values are preserved and only missing caches are populated. Signed-off-by: Lei, HUANG * refactor(mito-codec): share string framing in row converter Move encoded_string_len to the parent module so Dense and Sparse use the same framing helper without depending on each other. Preserve its implementation and visibility. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG * chore: sync lock * fix: shear and check issues --------- Signed-off-by: Ning Sun Signed-off-by: dhruvxvaishnav Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com> Signed-off-by: Dennis Zhuang Signed-off-by: fys Signed-off-by: Lei, HUANG Signed-off-by: WenyXu Co-authored-by: Dhruv Vaishnav Co-authored-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com> Co-authored-by: dennis zhuang Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com> Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com> Co-authored-by: Weny Xu --- .github/workflows/rust.yml | 15 +- .gitignore | 2 - Cargo.lock | 130 +------------ Cargo.toml | 2 - Makefile | 7 +- flake.nix | 2 +- scripts/fix-udeps.py | 265 --------------------------- src/catalog/Cargo.toml | 2 - src/cli/Cargo.toml | 8 - src/client/Cargo.toml | 4 - src/cmd/Cargo.toml | 10 +- src/common/config/Cargo.toml | 1 - src/common/frontend/Cargo.toml | 5 - src/common/function/Cargo.toml | 5 - src/common/grpc-expr/Cargo.toml | 4 - src/common/mem-prof/Cargo.toml | 6 + src/common/meta/Cargo.toml | 2 - src/common/procedure-test/Cargo.toml | 1 - src/common/query/Cargo.toml | 5 + src/common/runtime/Cargo.toml | 2 - src/common/sql/Cargo.toml | 1 - src/common/stat/Cargo.toml | 1 - src/common/substrait/Cargo.toml | 5 - src/datanode/Cargo.toml | 3 - src/datatypes/Cargo.toml | 5 + src/file-engine/Cargo.toml | 1 - src/flow/Cargo.toml | 17 -- src/frontend/Cargo.toml | 6 - src/log-store/Cargo.toml | 8 +- src/meta-client/Cargo.toml | 4 - src/meta-srv/Cargo.toml | 10 - src/metric-engine/Cargo.toml | 12 +- src/mito2/Cargo.toml | 20 +- src/object-store/Cargo.toml | 13 +- src/object-store/src/lib.rs | 1 + src/operator/Cargo.toml | 2 - src/partition/Cargo.toml | 4 - src/pipeline/Cargo.toml | 9 - src/query/Cargo.toml | 7 - src/servers/Cargo.toml | 4 +- src/session/Cargo.toml | 6 + src/sql/Cargo.toml | 7 - src/store-api/Cargo.toml | 5 - src/table/Cargo.toml | 3 - tests-fuzz/Cargo.toml | 9 + tests-integration/Cargo.toml | 39 ++-- tests/runner/Cargo.toml | 1 - 47 files changed, 79 insertions(+), 602 deletions(-) delete mode 100755 scripts/fix-udeps.py diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e4483b4a42b..1326ab5fd3d 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -24,7 +24,7 @@ on: name: Rust CI env: - CARGO_UDEPS_VERSION: "0.1.61" + CARGO_SHEAR_VERSION: "1.13.4" concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -84,20 +84,11 @@ jobs: - uses: actions/checkout@v4 with: persist-credentials: false - - uses: arduino/setup-protoc@v3 - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} - uses: actions-rust-lang/setup-rust-toolchain@v1 - - name: Rust Cache - uses: Swatinem/rust-cache@v2 - with: - shared-key: "check-udeps" - cache-all-crates: "true" - save-if: ${{ github.ref == 'refs/heads/main' }} - - name: Install cargo-udeps + - name: Install cargo-shear uses: taiki-e/install-action@v2 with: - tool: cargo-udeps@${{ env.CARGO_UDEPS_VERSION }} + tool: cargo-shear@${{ env.CARGO_SHEAR_VERSION }} fallback: none - name: Check unused dependencies run: make check-udeps diff --git a/.gitignore b/.gitignore index 90893f8aac1..dd60d96224e 100644 --- a/.gitignore +++ b/.gitignore @@ -52,8 +52,6 @@ venv/ tests-fuzz/artifacts/ tests-fuzz/corpus/ -# cargo-udeps reports -udeps-report.json # Nix .direnv diff --git a/Cargo.lock b/Cargo.lock index ea55f1b297d..6c7916d8194 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1673,7 +1673,6 @@ dependencies = [ "async-trait", "bytes", "cache", - "catalog", "chrono", "common-base", "common-catalog", @@ -1701,7 +1700,6 @@ dependencies = [ "lazy_static", "meta-client", "moka", - "object-store", "partition", "paste", "prometheus 0.14.0", @@ -2006,10 +2004,8 @@ version = "1.3.0-alpha.1" dependencies = [ "async-stream", "async-trait", - "auth", "base64 0.22.1", "cache", - "catalog", "chrono", "clap", "client", @@ -2022,13 +2018,10 @@ dependencies = [ "common-macro", "common-meta", "common-procedure", - "common-query", - "common-recordbatch", "common-runtime", "common-telemetry", "common-test-util", "common-time", - "common-version", "common-wal", "datatypes", "etcd-client", @@ -2048,7 +2041,6 @@ dependencies = [ "serde", "serde_json", "servers", - "session", "snafu 0.8.6", "sql", "standalone", @@ -2056,7 +2048,6 @@ dependencies = [ "table", "tempfile", "tokio", - "tracing-appender", "url", "uuid", ] @@ -2074,7 +2065,6 @@ dependencies = [ "common-catalog", "common-error", "common-grpc", - "common-grpc-expr", "common-macro", "common-meta", "common-query", @@ -2100,7 +2090,6 @@ dependencies = [ "tokio", "tokio-stream", "tonic 0.14.2", - "tracing", ] [[package]] @@ -2150,10 +2139,8 @@ dependencies = [ "api", "async-trait", "auth", - "base64 0.22.1", "cache", "catalog", - "chrono", "clap", "cli", "client", @@ -2170,7 +2157,6 @@ dependencies = [ "common-options", "common-procedure", "common-query", - "common-recordbatch", "common-runtime", "common-stat", "common-telemetry", @@ -2184,7 +2170,6 @@ dependencies = [ "datafusion-physical-plan", "datanode", "datatypes", - "either", "etcd-client", "file-engine", "flow", @@ -2192,7 +2177,6 @@ dependencies = [ "futures", "hex", "human-panic", - "humantime", "lazy_static", "log-store", "meta-client", @@ -2216,20 +2200,17 @@ dependencies = [ "serde", "serde_json", "servers", - "session", "similar-asserts", "smallvec", "snafu 0.8.6", "sqlparser", "standalone", "store-api", - "table", "temp-env", "tempfile", "tikv-jemallocator", "tokio", "toml 0.8.23", - "tonic 0.14.2", "tracing-appender", ] @@ -2355,7 +2336,6 @@ dependencies = [ "serde_json", "snafu 0.8.6", "temp-env", - "tempfile", "toml 0.8.23", ] @@ -2456,14 +2436,11 @@ dependencies = [ "common-grpc", "common-macro", "common-meta", - "common-time", "greptime-proto", - "humantime", "meta-client", "serde", "serde_json", "snafu 0.8.6", - "tokio", "tonic 0.14.2", ] @@ -2474,22 +2451,18 @@ dependencies = [ "ahash 0.8.12", "api", "approx 0.5.1", - "arc-swap", "arrow 59.2.0", - "arrow-cast 59.2.0", "arrow-schema 59.2.0", "async-trait", "axum 0.8.4", "bincode", "catalog", - "chrono", "common-base", "common-catalog", "common-error", "common-macro", "common-meta", "common-query", - "common-runtime", "common-telemetry", "common-time", "common-version", @@ -2514,7 +2487,6 @@ dependencies = [ "memchr", "mito-codec", "nalgebra", - "num", "num-traits", "paste", "pretty_assertions", @@ -2591,7 +2563,6 @@ name = "common-grpc-expr" version = "1.3.0-alpha.1" dependencies = [ "api", - "common-base", "common-catalog", "common-error", "common-macro", @@ -2599,7 +2570,6 @@ dependencies = [ "common-sql", "common-time", "datatypes", - "paste", "prost 0.14.1", "snafu 0.8.6", "store-api", @@ -2675,7 +2645,6 @@ dependencies = [ "common-wal", "common-workload", "datafusion-common", - "datafusion-expr", "datatypes", "deadpool", "deadpool-postgres", @@ -2684,7 +2653,6 @@ dependencies = [ "flexbuffers", "futures", "futures-util", - "greptime-proto", "hex", "humantime-serde", "itertools 0.14.0", @@ -2779,7 +2747,6 @@ version = "1.3.0-alpha.1" dependencies = [ "async-trait", "common-procedure", - "snafu 0.8.6", "tokio", ] @@ -2852,13 +2819,11 @@ dependencies = [ "lazy_static", "num_cpus", "once_cell", - "parking_lot 0.12.4", "paste", "pin-project", "prometheus 0.14.0", "ratelimit", "serde", - "serde_json", "snafu 0.8.6", "tempfile", "tokio", @@ -2885,7 +2850,6 @@ dependencies = [ "common-decimal", "common-error", "common-macro", - "common-telemetry", "common-time", "datafusion-sql", "datatypes", @@ -2903,7 +2867,6 @@ dependencies = [ "common-base", "common-runtime", "common-telemetry", - "lazy_static", "nix 0.30.1", "num_cpus", "prometheus 0.14.0", @@ -4533,7 +4496,6 @@ dependencies = [ "async-trait", "bytes", "cache", - "client", "common-base", "common-catalog", "common-config", @@ -4541,7 +4503,6 @@ dependencies = [ "common-error", "common-function", "common-greptimedb-telemetry", - "common-grpc", "common-macro", "common-meta", "common-options", @@ -5332,7 +5293,6 @@ dependencies = [ "common-datasource", "common-error", "common-macro", - "common-query", "common-recordbatch", "common-telemetry", "common-test-util", @@ -5472,19 +5432,14 @@ dependencies = [ "arrow 59.2.0", "arrow-flight", "arrow-schema 59.2.0", - "async-recursion", "async-trait", - "bytes", "catalog", "chrono", "client", "common-base", "common-catalog", "common-config", - "common-datasource", - "common-decimal", "common-error", - "common-frontend", "common-function", "common-grpc", "common-macro", @@ -5497,42 +5452,31 @@ dependencies = [ "common-telemetry", "common-time", "common-version", - "config", "datafusion", "datafusion-common", "datafusion-expr", "datafusion-physical-expr", - "datafusion-substrait", "datatypes", - "enum_dispatch", "futures", "get-size2", "greptime-proto", "hostname 0.4.1", - "http 1.5.0", "humantime-serde", "itertools 0.14.0", "lazy_static", "meta-client", - "nom 7.1.3", - "num-traits", - "num_cpus", "operator", - "partition", "pretty_assertions", "prometheus 0.14.0", - "prost 0.14.1", "query", "rand 0.9.4", "serde", "serde_json", "servers", "session", - "smallvec", "snafu 0.8.6", "sql", "store-api", - "strum 0.27.1", "substrait 1.3.0-alpha.1", "table", "tokio", @@ -5633,7 +5577,6 @@ name = "frontend" version = "1.3.0-alpha.1" dependencies = [ "api", - "arc-swap", "arrow-flight", "async-stream", "async-trait", @@ -5657,7 +5600,6 @@ dependencies = [ "common-memory-manager", "common-meta", "common-options", - "common-procedure", "common-query", "common-recordbatch", "common-runtime", @@ -5672,13 +5614,11 @@ dependencies = [ "datatypes", "futures", "hostname 0.4.1", - "humantime", "humantime-serde", "hyper-util", "lazy_static", "log-query", "meta-client", - "num_cpus", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", @@ -5689,7 +5629,6 @@ dependencies = [ "promql-parser", "prost 0.14.1", "query", - "rand 0.9.4", "reqwest 0.13.4", "serde", "serde_json", @@ -5708,7 +5647,6 @@ dependencies = [ "tonic 0.14.2", "tower 0.5.2", "tracing", - "uuid", ] [[package]] @@ -8121,7 +8059,6 @@ dependencies = [ "derive_builder 0.20.2", "futures", "futures-util", - "itertools 0.14.0", "lazy_static", "object-store", "pin-project", @@ -8465,20 +8402,16 @@ dependencies = [ "common-options", "common-telemetry", "common-time", - "datatypes", "futures", - "futures-util", "humantime-serde", "meta-srv", "rand 0.9.4", "serde", "serde_json", - "session", "snafu 0.8.6", "tokio", "tokio-stream", "tonic 0.14.2", - "tower 0.5.2", "tracing", "tracing-subscriber", ] @@ -8490,9 +8423,7 @@ dependencies = [ "api", "async-trait", "axum 0.8.4", - "axum-extra", "axum-macros", - "bytes", "chrono", "clap", "client", @@ -8525,19 +8456,14 @@ dependencies = [ "futures", "h2 0.3.26", "hostname 0.4.1", - "http-body-util", - "humantime", "humantime-serde", "hyper-util", "itertools 0.14.0", "lazy_static", - "local-ip-address", "once_cell", "ordered-float 4.6.0", - "parking_lot 0.12.4", "partition", "prometheus 0.14.0", - "prost 0.14.1", "rand 0.9.4", "regex", "rskafka", @@ -8555,7 +8481,6 @@ dependencies = [ "toml 0.8.23", "tonic 0.14.2", "tower 0.5.2", - "tower-http 0.6.11", "tracing", "tracing-subscriber", "typetag", @@ -8603,7 +8528,6 @@ dependencies = [ "common-recordbatch", "common-runtime", "common-telemetry", - "common-test-util", "common-time", "common-wal", "criterion 0.7.0", @@ -8723,9 +8647,7 @@ dependencies = [ "bytes", "chrono", "common-base", - "common-config", "common-datasource", - "common-decimal", "common-error", "common-function", "common-grpc", @@ -8764,7 +8686,6 @@ dependencies = [ "object-store", "parquet", "partition", - "paste", "pin-project", "prometheus 0.14.0", "prost 0.14.1", @@ -8776,11 +8697,9 @@ dependencies = [ "rskafka", "rstest", "rstest_reuse", - "scopeguard", "serde", "serde_json", "serde_with", - "session", "smallvec", "snafu 0.8.6", "store-api", @@ -9562,10 +9481,8 @@ name = "object-store" version = "1.3.0-alpha.1" dependencies = [ "anyhow", - "async-trait", "bytes", "cap-std", - "chrono", "common-base", "common-error", "common-macro", @@ -9575,13 +9492,12 @@ dependencies = [ "derive_builder 0.20.2", "futures", "humantime-serde", - "lazy_static", + "object-store", "object_store", "object_store_opendal", "opendal", "opendal-http-transport-reqwest", "prometheus 0.14.0", - "rand 0.9.4", "reqwest 0.13.4", "serde", "snafu 0.8.6", @@ -10096,7 +10012,6 @@ dependencies = [ "ahash 0.8.12", "api", "arrow 59.2.0", - "arrow-ipc 59.2.0", "async-stream", "async-trait", "axum 0.8.4", @@ -10143,7 +10058,6 @@ dependencies = [ "path-slash", "percent-encoding", "prometheus 0.14.0", - "prost 0.14.1", "prost-types 0.14.1", "query", "regex", @@ -10492,19 +10406,15 @@ name = "partition" version = "1.3.0-alpha.1" dependencies = [ "api", - "async-trait", "common-base", "common-error", "common-macro", "common-meta", - "common-query", - "common-telemetry", "criterion 0.7.0", "datafusion-common", "datafusion-expr", "datafusion-physical-expr", "datatypes", - "itertools 0.14.0", "moka", "rand 0.8.5", "serde", @@ -10893,29 +10803,22 @@ dependencies = [ "api", "arrow 59.2.0", "arrow-schema 59.2.0", - "async-trait", "catalog", "chrono", "chrono-tz", "common-catalog", "common-error", - "common-function", "common-macro", "common-meta", "common-query", "common-recordbatch", - "common-runtime", "common-telemetry", "common-time", "criterion 0.7.0", - "crossbeam-utils", "csv", - "dashmap", "datafusion", "datafusion-common", "datafusion-expr", - "datafusion-functions", - "datafusion-physical-expr", "datatypes", "dyn-fmt", "enum_dispatch", @@ -10930,10 +10833,8 @@ dependencies = [ "once_cell", "operator", "ordered-float 4.6.0", - "paste", "prometheus 0.14.0", "query", - "rayon", "regex", "serde", "serde_json", @@ -11867,7 +11768,6 @@ version = "1.3.0-alpha.1" dependencies = [ "ahash 0.8.12", "api", - "arc-swap", "arrow 59.2.0", "arrow-schema 59.2.0", "async-recursion", @@ -11901,12 +11801,10 @@ dependencies = [ "datafusion-proto", "datafusion-sql", "datatypes", - "either", "fastrand", "futures", "futures-util", "greptime-proto", - "humantime", "itertools 0.14.0", "jsonb", "lazy_static", @@ -11915,11 +11813,8 @@ dependencies = [ "meter-macros", "moka", "nalgebra", - "num", - "num-traits", "object-store", "once_cell", - "parking_lot 0.12.4", "partition", "paste", "pretty_assertions", @@ -11939,7 +11834,6 @@ dependencies = [ "substrait 1.3.0-alpha.1", "table", "tokio", - "tokio-stream", "toml 0.8.23", "tracing", "uuid", @@ -13582,7 +13476,6 @@ dependencies = [ "common-catalog", "common-decimal", "common-error", - "common-frontend", "common-grpc", "common-macro", "common-mem-prof", @@ -13602,7 +13495,6 @@ dependencies = [ "common-version", "criterion 0.5.1", "csv", - "dashmap", "datafusion", "datafusion-common", "datafusion-expr", @@ -14070,25 +13962,19 @@ dependencies = [ "common-base", "common-catalog", "common-datasource", - "common-decimal", "common-error", "common-macro", "common-query", "common-sql", - "common-telemetry", "common-time", "datafusion", "datafusion-common", "datafusion-expr", - "datafusion-physical-expr", "datafusion-sql", "datatypes", "either", - "hex", - "humantime", "iso8601", "itertools 0.14.0", - "jsonb", "lazy_static", "pretty_assertions", "promql-parser", @@ -14098,7 +13984,6 @@ dependencies = [ "snafu 0.8.6", "sqlparser", "sqlparser_derive 0.1.1", - "store-api", "table", "uuid", ] @@ -14130,7 +14015,6 @@ dependencies = [ "client", "common-error", "common-query", - "common-recordbatch", "common-time", "datatypes", "flate2", @@ -14458,7 +14342,6 @@ dependencies = [ "api", "aquamarine", "arrow-schema 59.2.0", - "async-stream", "async-trait", "bytes", "common-base", @@ -14485,7 +14368,6 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.6", - "sqlparser", "strum 0.27.1", "tokio", "uuid", @@ -14654,16 +14536,13 @@ dependencies = [ "common-error", "common-function", "common-macro", - "common-telemetry", "datafusion", "datafusion-common", "datafusion-expr", "datafusion-substrait", - "datatypes", "promql", "prost 0.14.1", "snafu 0.8.6", - "tokio", ] [[package]] @@ -14812,7 +14691,6 @@ dependencies = [ "futures", "greptime-proto", "humantime", - "humantime-serde", "lazy_static", "once_cell", "parquet", @@ -14820,10 +14698,8 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.6", - "sqlparser", "store-api", "tokio", - "tokio-util", ] [[package]] @@ -15121,7 +14997,6 @@ version = "1.3.0-alpha.1" dependencies = [ "api", "arrow-flight", - "async-stream", "async-trait", "auth", "axum 0.8.4", @@ -15149,7 +15024,6 @@ dependencies = [ "common-query", "common-recordbatch", "common-runtime", - "common-stat", "common-telemetry", "common-test-util", "common-time", @@ -15180,7 +15054,6 @@ dependencies = [ "opentelemetry-proto", "operator", "otel-arrow-rust", - "partition", "paste", "pipeline", "plugins", @@ -15195,7 +15068,6 @@ dependencies = [ "servers", "session", "similar-asserts", - "snafu 0.8.6", "sql", "sqlx", "standalone", diff --git a/Cargo.toml b/Cargo.toml index 7c5c53d91c8..5cdc3cf63a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -104,7 +104,6 @@ aquamarine = "0.6" arrow = { version = "=59.2.0", features = ["prettyprint"] } arrow-array = { version = "=59.2.0", default-features = false, features = ["chrono-tz"] } arrow-buffer = "=59.2.0" -arrow-cast = "=59.2.0" arrow-flight = "=59.2.0" arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] } arrow-schema = { version = "=59.2.0", features = ["serde"] } @@ -236,7 +235,6 @@ rskafka = { git = "https://github.com/influxdata/rskafka.git", rev = "80f42eb15e rstest = "0.25" rstest_reuse = "0.7" rust_decimal = "1.33" -rustc-hash = "2.0" rustls = { version = "0.23.25", default-features = false } ryu = "1.0" sea-query = "0.32" diff --git a/Makefile b/Makefile index 7176eafd0bb..1379d86683c 100644 --- a/Makefile +++ b/Makefile @@ -233,14 +233,11 @@ fix-clippy: ## Fix clippy violations. .PHONY: check-udeps check-udeps: ## Check unused dependencies. - cargo udeps --workspace --all-targets + cargo shear .PHONY: fix-udeps fix-udeps: ## Remove unused dependencies automatically. - @echo "Running cargo-udeps to find unused dependencies..." - @cargo udeps --workspace --all-targets --output json > udeps-report.json || true - @echo "Removing unused dependencies..." - @python3 scripts/fix-udeps.py udeps-report.json + cargo shear --fix .PHONY: fmt-check fmt-check: ## Check code format. diff --git a/flake.nix b/flake.nix index a5886b8add1..2d95d8dcfd6 100644 --- a/flake.nix +++ b/flake.nix @@ -46,7 +46,7 @@ ]) cargo-nextest cargo-llvm-cov - cargo-udeps + cargo-shear taplo curl gnuplot ## for cargo bench diff --git a/scripts/fix-udeps.py b/scripts/fix-udeps.py deleted file mode 100755 index 7547b897361..00000000000 --- a/scripts/fix-udeps.py +++ /dev/null @@ -1,265 +0,0 @@ -# Copyright 2023 Greptime Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import os -import re -import sys - - -def load_udeps_report(report_path): - try: - with open(report_path, "r") as f: - return json.load(f) - except FileNotFoundError: - print(f"Error: Report file '{report_path}' not found.") - return None - except json.JSONDecodeError as e: - print(f"Error: Invalid JSON in report file: {e}") - return None - - -def extract_unused_dependencies(report): - """ - Extract and organize unused dependencies from the cargo-udeps JSON report. - - The cargo-udeps report has this structure: - { - "unused_deps": { - "package_name v0.1.0 (/path/to/package)": { - "normal": ["dep1", "dep2"], - "development": ["dev_dep1"], - "build": ["build_dep1"], - "manifest_path": "/path/to/Cargo.toml" - } - } - } - - Args: - report (dict): The parsed JSON report from cargo-udeps - - Returns: - dict: Organized unused dependencies by package name: - { - "package_name": { - "dependencies": [("dep1", "normal"), ("dev_dep1", "dev")], - "manifest_path": "/path/to/Cargo.toml" - } - } - """ - if not report or "unused_deps" not in report: - return {} - - unused_deps = {} - for package_full_name, deps_info in report["unused_deps"].items(): - package_name = package_full_name.split(" ")[0] - - all_unused = [] - if deps_info.get("normal"): - all_unused.extend([(dep, "normal") for dep in deps_info["normal"]]) - if deps_info.get("development"): - all_unused.extend([(dep, "dev") for dep in deps_info["development"]]) - if deps_info.get("build"): - all_unused.extend([(dep, "build") for dep in deps_info["build"]]) - - if all_unused: - unused_deps[package_name] = { - "dependencies": all_unused, - "manifest_path": deps_info.get("manifest_path", "unknown"), - } - - return unused_deps - - -def get_section_pattern(dep_type): - """ - Get regex patterns to identify different dependency sections in Cargo.toml. - - Args: - dep_type (str): Type of dependency ("normal", "dev", or "build") - - Returns: - list: List of regex patterns to match the appropriate section headers - - """ - patterns = { - "normal": [r"\[dependencies\]", r"\[dependencies\..*?\]"], - "dev": [r"\[dev-dependencies\]", r"\[dev-dependencies\..*?\]"], - "build": [r"\[build-dependencies\]", r"\[build-dependencies\..*?\]"], - } - return patterns.get(dep_type, []) - - -def remove_dependency_line(content, dep_name, section_start, section_end): - """ - Remove a dependency line from a specific section of a Cargo.toml file. - - Args: - content (str): The entire content of the Cargo.toml file - dep_name (str): Name of the dependency to remove (e.g., "serde", "tokio") - section_start (int): Starting position of the section in the content - section_end (int): Ending position of the section in the content - - Returns: - tuple: (new_content, removed) where: - - new_content (str): The modified content with dependency removed - - removed (bool): True if dependency was found and removed, False otherwise - - Example input content format: - content = ''' - [package] - name = "my-crate" - version = "0.1.0" - - [dependencies] - serde = "1.0" - tokio = { version = "1.0", features = ["full"] } - serde_json.workspace = true - - [dev-dependencies] - tempfile = "3.0" - ''' - - # If dep_name = "serde", section_start = start of [dependencies], - # section_end = start of [dev-dependencies], this function will: - # 1. Extract the section: "serde = "1.0"\ntokio = { version = "1.0", features = ["full"] }\nserde_json.workspace = true\n" - # 2. Find and remove the line: "serde = "1.0"" - # 3. Return the modified content with that line removed - """ - section_content = content[section_start:section_end] - - dep_patterns = [ - rf"^{re.escape(dep_name)}\s*=.*$", # e.g., "serde = "1.0"" - rf"^{re.escape(dep_name)}\.workspace\s*=.*$", # e.g., "serde_json.workspace = true" - ] - - for pattern in dep_patterns: - match = re.search(pattern, section_content, re.MULTILINE) - if match: - line_start = section_start + match.start() # Start of the matched line - line_end = section_start + match.end() # End of the matched line - - if line_end < len(content) and content[line_end] == "\n": - line_end += 1 - - return content[:line_start] + content[line_end:], True - - return content, False - - -def remove_dependency_from_toml(file_path, dep_name, dep_type): - """ - Remove a specific dependency from a Cargo.toml file. - - Args: - file_path (str): Path to the Cargo.toml file - dep_name (str): Name of the dependency to remove - dep_type (str): Type of dependency ("normal", "dev", or "build") - - Returns: - bool: True if dependency was successfully removed, False otherwise - """ - try: - with open(file_path, "r") as f: - content = f.read() - - section_patterns = get_section_pattern(dep_type) - if not section_patterns: - return False - - for pattern in section_patterns: - section_match = re.search(pattern, content, re.IGNORECASE) - if not section_match: - continue - - section_start = section_match.end() - next_section = re.search(r"\n\s*\[", content[section_start:]) - section_end = ( - section_start + next_section.start() if next_section else len(content) - ) - - new_content, removed = remove_dependency_line( - content, dep_name, section_start, section_end - ) - if removed: - with open(file_path, "w") as f: - f.write(new_content) - return True - - return False - - except Exception as e: - print(f"Error processing {file_path}: {e}") - return False - - -def process_unused_dependencies(unused_deps): - """ - Process and remove all unused dependencies from their respective Cargo.toml files. - - Args: - unused_deps (dict): Dictionary of unused dependencies organized by package: - { - "package_name": { - "dependencies": [("dep1", "normal"), ("dev_dep1", "dev")], - "manifest_path": "/path/to/Cargo.toml" - } - } - - """ - if not unused_deps: - print("No unused dependencies found.") - return - - total_removed = 0 - total_failed = 0 - - for package, info in unused_deps.items(): - deps = info["dependencies"] - manifest_path = info["manifest_path"] - - if not os.path.exists(manifest_path): - print(f"Manifest file not found: {manifest_path}") - total_failed += len(deps) - continue - - for dep, dep_type in deps: - if remove_dependency_from_toml(manifest_path, dep, dep_type): - print(f"Removed {dep} from {package}") - total_removed += 1 - else: - print(f"Failed to remove {dep} from {package}") - total_failed += 1 - - print(f"Removed {total_removed} dependencies") - if total_failed > 0: - print(f"Failed to remove {total_failed} dependencies") - - -def main(): - if len(sys.argv) > 1: - report_path = sys.argv[1] - else: - report_path = "udeps-report.json" - - report = load_udeps_report(report_path) - if report is None: - sys.exit(1) - - unused_deps = extract_unused_dependencies(report) - process_unused_dependencies(unused_deps) - - -if __name__ == "__main__": - main() diff --git a/src/catalog/Cargo.toml b/src/catalog/Cargo.toml index f46887e148a..fe42ddc2d13 100644 --- a/src/catalog/Cargo.toml +++ b/src/catalog/Cargo.toml @@ -61,9 +61,7 @@ tokio-stream.workspace = true [dev-dependencies] cache.workspace = true -catalog = { workspace = true, features = ["testing"] } chrono.workspace = true common-meta = { workspace = true, features = ["testing"] } common-query = { workspace = true, features = ["testing"] } -object-store.workspace = true tokio.workspace = true diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 587e5b8cffa..8178fa1b0b5 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -18,10 +18,8 @@ workspace = true [dependencies] async-stream.workspace = true async-trait.workspace = true -auth.workspace = true base64.workspace = true cache.workspace = true -catalog.workspace = true chrono.workspace = true clap.workspace = true client = { workspace = true, features = ["testing"] } @@ -34,14 +32,11 @@ common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true common-procedure.workspace = true -common-query.workspace = true -common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry = { workspace = true, features = [ "deadlock_detection", ] } common-time.workspace = true -common-version.workspace = true common-wal.workspace = true datatypes.workspace = true etcd-client.workspace = true @@ -61,21 +56,18 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true -session.workspace = true snafu.workspace = true sql.workspace = true standalone.workspace = true store-api.workspace = true table.workspace = true tokio.workspace = true -tracing-appender.workspace = true url.workspace = true uuid.workspace = true [dev-dependencies] common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true -common-version.workspace = true object-store = { workspace = true, features = ["testing"] } serde.workspace = true tempfile.workspace = true diff --git a/src/client/Cargo.toml b/src/client/Cargo.toml index 5b479b1fbf2..36c3f60a9a4 100644 --- a/src/client/Cargo.toml +++ b/src/client/Cargo.toml @@ -45,7 +45,3 @@ substrait.workspace = true tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true - -[dev-dependencies] -common-grpc-expr.workspace = true -tracing = "0.1" diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index fff268dea45..367b86809d8 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -43,10 +43,8 @@ workspace = true api.workspace = true async-trait.workspace = true auth.workspace = true -base64.workspace = true cache.workspace = true catalog.workspace = true -chrono.workspace = true clap.workspace = true cli.workspace = true client.workspace = true @@ -63,7 +61,6 @@ common-meta.workspace = true common-options.workspace = true common-procedure.workspace = true common-query.workspace = true -common-recordbatch.workspace = true common-runtime.workspace = true common-stat.workspace = true common-telemetry = { workspace = true, features = [ @@ -78,19 +75,16 @@ datafusion-physical-plan.workspace = true datafusion_object_store.workspace = true datanode.workspace = true datatypes.workspace = true -either = "1.15" etcd-client.workspace = true flow.workspace = true frontend = { workspace = true, default-features = false } futures.workspace = true hex.workspace = true human-panic = "2.0" -humantime.workspace = true lazy_static.workspace = true log-store.workspace = true meta-client.workspace = true meta-srv.workspace = true -metric-engine.workspace = true mito-codec.workspace = true mito2.workspace = true moka = { workspace = true, features = ["future"] } @@ -107,16 +101,13 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true -session.workspace = true smallvec.workspace = true snafu.workspace = true sqlparser.workspace = true standalone.workspace = true store-api.workspace = true -table.workspace = true tokio.workspace = true toml.workspace = true -tonic.workspace = true tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] @@ -134,6 +125,7 @@ client = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-version.workspace = true file-engine.workspace = true +metric-engine.workspace = true mito2.workspace = true serde.workspace = true similar-asserts.workspace = true diff --git a/src/common/config/Cargo.toml b/src/common/config/Cargo.toml index 27b238add73..ad356716867 100644 --- a/src/common/config/Cargo.toml +++ b/src/common/config/Cargo.toml @@ -29,4 +29,3 @@ datanode.workspace = true meta-client.workspace = true serde.workspace = true temp-env = "0.3" -tempfile.workspace = true diff --git a/src/common/frontend/Cargo.toml b/src/common/frontend/Cargo.toml index 32a5146b566..26092ae13f7 100644 --- a/src/common/frontend/Cargo.toml +++ b/src/common/frontend/Cargo.toml @@ -12,14 +12,9 @@ common-event-recorder.workspace = true common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true -common-time.workspace = true greptime-proto.workspace = true -humantime.workspace = true meta-client.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true tonic.workspace = true - -[dev-dependencies] -tokio.workspace = true diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index ee6cfac9dfb..3a5da30c136 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -17,21 +17,17 @@ ai_functions = ["dep:futures", "dep:reqwest"] [dependencies] ahash.workspace = true api.workspace = true -arc-swap = "1.0" arrow.workspace = true -arrow-cast.workspace = true arrow-schema.workspace = true async-trait.workspace = true bincode = "=1.3.3" catalog.workspace = true -chrono.workspace = true common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true -common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true common-version.workspace = true @@ -55,7 +51,6 @@ jsonb.workspace = true memchr = "2.7" mito-codec.workspace = true nalgebra.workspace = true -num = "0.4" num-traits = "0.2" paste.workspace = true regex.workspace = true diff --git a/src/common/grpc-expr/Cargo.toml b/src/common/grpc-expr/Cargo.toml index a689de7c018..b735e8f9ba8 100644 --- a/src/common/grpc-expr/Cargo.toml +++ b/src/common/grpc-expr/Cargo.toml @@ -9,7 +9,6 @@ workspace = true [dependencies] api.workspace = true -common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true @@ -21,6 +20,3 @@ prost.workspace = true snafu.workspace = true store-api.workspace = true table.workspace = true - -[dev-dependencies] -paste.workspace = true diff --git a/src/common/mem-prof/Cargo.toml b/src/common/mem-prof/Cargo.toml index 10b67cf7fa7..46d959b02cc 100644 --- a/src/common/mem-prof/Cargo.toml +++ b/src/common/mem-prof/Cargo.toml @@ -26,3 +26,9 @@ jemalloc-pprof-mappings = { version = "0.7", package = "mappings" } # for get th [target.'cfg(not(windows))'.dependencies.tikv-jemalloc-sys] features = ["stats", "profiling", "unprefixed_malloc_on_supported_platforms"] version = "0.7" + +[package.metadata.cargo-shear] +# `tikv-jemalloc-sys` is a link-only dependency: it enables jemalloc +# profiling/stats features and pins the linked jemalloc instance, but is +# never referenced by name in source code. +ignored = ["tikv-jemalloc-sys"] diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8deb3d00814..bd974deb4cf 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -46,7 +46,6 @@ common-time.workspace = true common-wal.workspace = true common-workload.workspace = true datafusion-common.workspace = true -datafusion-expr.workspace = true datatypes.workspace = true deadpool = { workspace = true, optional = true } deadpool-postgres = { workspace = true, optional = true } @@ -55,7 +54,6 @@ etcd-client.workspace = true flexbuffers = "25.2" futures.workspace = true futures-util.workspace = true -greptime-proto.workspace = true hex.workspace = true humantime-serde.workspace = true itertools.workspace = true diff --git a/src/common/procedure-test/Cargo.toml b/src/common/procedure-test/Cargo.toml index e445a235bc9..907035dd052 100644 --- a/src/common/procedure-test/Cargo.toml +++ b/src/common/procedure-test/Cargo.toml @@ -10,5 +10,4 @@ workspace = true [dependencies] async-trait.workspace = true common-procedure = { workspace = true, features = ["testing"] } -snafu.workspace = true tokio.workspace = true diff --git a/src/common/query/Cargo.toml b/src/common/query/Cargo.toml index 9706098826f..4acdfafcb6c 100644 --- a/src/common/query/Cargo.toml +++ b/src/common/query/Cargo.toml @@ -39,3 +39,8 @@ common-base.workspace = true futures-util.workspace = true serde_json.workspace = true tokio.workspace = true + +[package.metadata.cargo-shear] +# `sqlparser` is required by `sqlparser_derive`: the generated code references +# `::sqlparser` paths, so it must stay a direct dependency. +ignored = ["sqlparser"] diff --git a/src/common/runtime/Cargo.toml b/src/common/runtime/Cargo.toml index 927aa17b9c2..410173f6445 100644 --- a/src/common/runtime/Cargo.toml +++ b/src/common/runtime/Cargo.toml @@ -25,13 +25,11 @@ futures.workspace = true lazy_static.workspace = true num_cpus.workspace = true once_cell.workspace = true -parking_lot.workspace = true paste.workspace = true pin-project.workspace = true prometheus.workspace = true ratelimit.workspace = true serde.workspace = true -serde_json.workspace = true snafu.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/src/common/sql/Cargo.toml b/src/common/sql/Cargo.toml index 3a66ec2d04e..318685996d5 100644 --- a/src/common/sql/Cargo.toml +++ b/src/common/sql/Cargo.toml @@ -10,7 +10,6 @@ common-base.workspace = true common-decimal.workspace = true common-error.workspace = true common-macro.workspace = true -common-telemetry.workspace = true common-time.workspace = true datafusion-sql.workspace = true datatypes.workspace = true diff --git a/src/common/stat/Cargo.toml b/src/common/stat/Cargo.toml index 2a99107c909..fcf927e9cca 100644 --- a/src/common/stat/Cargo.toml +++ b/src/common/stat/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true common-base.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -lazy_static.workspace = true num_cpus.workspace = true prometheus.workspace = true sysinfo.workspace = true diff --git a/src/common/substrait/Cargo.toml b/src/common/substrait/Cargo.toml index a4e78475051..a56b56b8301 100644 --- a/src/common/substrait/Cargo.toml +++ b/src/common/substrait/Cargo.toml @@ -13,7 +13,6 @@ bytes.workspace = true common-error.workspace = true common-function.workspace = true common-macro.workspace = true -common-telemetry.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true @@ -21,7 +20,3 @@ datafusion-substrait.workspace = true promql.workspace = true prost.workspace = true snafu.workspace = true - -[dev-dependencies] -datatypes.workspace = true -tokio.workspace = true diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 6b78858bd2f..7c126764723 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -16,7 +16,6 @@ api.workspace = true arrow-flight.workspace = true async-trait.workspace = true bytes.workspace = true -client.workspace = true common-base.workspace = true common-catalog.workspace = true common-config.workspace = true @@ -24,7 +23,6 @@ common-datasource.workspace = true common-error.workspace = true common-function.workspace = true common-greptimedb-telemetry.workspace = true -common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true common-options.workspace = true @@ -75,7 +73,6 @@ tracing.workspace = true [dev-dependencies] cache.workspace = true -client.workspace = true common-meta = { workspace = true, features = ["testing"] } common-query.workspace = true common-test-util.workspace = true diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index 124b1e0ea64..d1dd70acef1 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -37,3 +37,8 @@ serde_json.workspace = true snafu.workspace = true sqlparser.workspace = true sqlparser_derive = "0.1" + +[package.metadata.cargo-shear] +# `sqlparser` is required by `sqlparser_derive`: the generated code references +# `::sqlparser` paths, so it must stay a direct dependency. +ignored = ["sqlparser"] diff --git a/src/file-engine/Cargo.toml b/src/file-engine/Cargo.toml index ec904addf1e..dda5fb4221c 100644 --- a/src/file-engine/Cargo.toml +++ b/src/file-engine/Cargo.toml @@ -18,7 +18,6 @@ common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true common-macro.workspace = true -common-query.workspace = true common-recordbatch.workspace = true common-telemetry.workspace = true common-test-util = { workspace = true, optional = true } diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index 4586e612693..d0ac5eee9ac 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -11,18 +11,13 @@ workspace = true api.workspace = true arrow.workspace = true arrow-schema.workspace = true -async-recursion = "1.0" async-trait.workspace = true -bytes.workspace = true catalog.workspace = true chrono.workspace = true client.workspace = true common-base.workspace = true common-config.workspace = true -common-datasource.workspace = true -common-decimal.workspace = true common-error.workspace = true -common-frontend.workspace = true common-function.workspace = true common-grpc.workspace = true common-macro.workspace = true @@ -35,41 +30,30 @@ common-stat.workspace = true common-telemetry.workspace = true common-time.workspace = true common-version.workspace = true -config.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datafusion-physical-expr.workspace = true -datafusion-substrait.workspace = true datatypes.workspace = true -enum_dispatch = "0.3" futures.workspace = true get-size2 = "0.1.2" greptime-proto.workspace = true hostname.workspace = true -http.workspace = true humantime-serde.workspace = true itertools.workspace = true lazy_static.workspace = true meta-client.workspace = true -nom = "7.1.3" -num-traits = "0.2" -num_cpus.workspace = true operator.workspace = true -partition.workspace = true prometheus.workspace = true -prost.workspace = true query.workspace = true rand.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true session.workspace = true -smallvec.workspace = true snafu.workspace = true sql.workspace = true store-api.workspace = true -strum.workspace = true substrait.workspace = true table.workspace = true tokio.workspace = true @@ -81,7 +65,6 @@ catalog = { workspace = true, features = ["testing"] } common-catalog.workspace = true common-meta = { workspace = true, features = ["testing"] } pretty_assertions.workspace = true -prost.workspace = true query.workspace = true session.workspace = true table.workspace = true diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index b59e1e83a4a..1aa057d0fcd 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -19,7 +19,6 @@ workspace = true [dependencies] api.workspace = true -arc-swap = "1.0" async-stream.workspace = true async-trait.workspace = true auth.workspace = true @@ -42,7 +41,6 @@ common-macro.workspace = true common-memory-manager.workspace = true common-meta.workspace = true common-options.workspace = true -common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true @@ -56,12 +54,10 @@ datafusion-expr.workspace = true datatypes.workspace = true futures.workspace = true hostname.workspace = true -humantime.workspace = true humantime-serde.workspace = true lazy_static.workspace = true log-query.workspace = true meta-client.workspace = true -num_cpus.workspace = true opentelemetry-proto.workspace = true opentelemetry-semantic-conventions = { version = "0.31", features = ["semconv_experimental"] } operator.workspace = true @@ -72,7 +68,6 @@ prometheus.workspace = true promql-parser.workspace = true prost.workspace = true query.workspace = true -rand.workspace = true serde.workspace = true serde_json.workspace = true servers.workspace = true @@ -100,4 +95,3 @@ serde_json.workspace = true strfmt = "0.2" temp-env.workspace = true tower.workspace = true -uuid.workspace = true diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 15560b06b2c..5d068ca344d 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -31,7 +31,6 @@ delta-encoding = "0.4" derive_builder.workspace = true futures.workspace = true futures-util.workspace = true -itertools.workspace = true lazy_static.workspace = true object-store.workspace = true pin-project.workspace = true @@ -49,8 +48,13 @@ tokio.workspace = true common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-wal = { workspace = true, features = ["testing"] } -itertools.workspace = true object-store = { workspace = true, features = ["testing"] } rand.workspace = true rskafka = { workspace = true, features = ["unstable-fuzzing"] } uuid.workspace = true + +[package.metadata.cargo-shear] +# `protobuf` is required by the code generated by `build.rs`: the generated +# protobuf bindings (included via `include!`) reference `::protobuf` paths, +# but no source file does. +ignored = ["protobuf"] diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 984f434cc51..9863b32bbbf 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -19,7 +19,6 @@ common-options.workspace = true common-telemetry.workspace = true common-time.workspace = true futures.workspace = true -futures-util.workspace = true humantime-serde.workspace = true rand.workspace = true serde.workspace = true @@ -31,10 +30,7 @@ tonic.workspace = true [dev-dependencies] common-event-recorder.workspace = true -datatypes.workspace = true futures.workspace = true meta-srv = { workspace = true, features = ["mock"] } -session.workspace = true -tower.workspace = true tracing = "0.1" tracing-subscriber.workspace = true diff --git a/src/meta-srv/Cargo.toml b/src/meta-srv/Cargo.toml index e6585505db4..6baca0934e1 100644 --- a/src/meta-srv/Cargo.toml +++ b/src/meta-srv/Cargo.toml @@ -14,16 +14,11 @@ enterprise = ["common-meta/enterprise"] [lints] workspace = true -[target.'cfg(not(target_os = "android"))'.dependencies] -local-ip-address.workspace = true - [dependencies] api.workspace = true async-trait.workspace = true axum.workspace = true -axum-extra.workspace = true axum-macros.workspace = true -bytes.workspace = true chrono.workspace = true clap.workspace = true client.workspace = true @@ -55,18 +50,14 @@ etcd-client.workspace = true futures.workspace = true h2 = "0.3" hostname.workspace = true -http-body-util = "0.1" -humantime.workspace = true humantime-serde.workspace = true hyper-util = { workspace = true, features = ["tokio"] } itertools.workspace = true lazy_static.workspace = true once_cell.workspace = true ordered-float.workspace = true -parking_lot.workspace = true partition.workspace = true prometheus.workspace = true -prost.workspace = true rand.workspace = true regex.workspace = true rskafka.workspace = true @@ -87,7 +78,6 @@ tokio-stream = { workspace = true, features = ["net"] } toml.workspace = true tonic.workspace = true tower.workspace = true -tower-http.workspace = true tracing.workspace = true typetag.workspace = true url = "2.3" diff --git a/src/metric-engine/Cargo.toml b/src/metric-engine/Cargo.toml index 0005a4dc543..92000549832 100644 --- a/src/metric-engine/Cargo.toml +++ b/src/metric-engine/Cargo.toml @@ -14,7 +14,7 @@ async-stream.workspace = true async-trait.workspace = true base64.workspace = true bytes.workspace = true -fxhash = "0.2" +chrono.workspace = true common-base.workspace = true common-error.workspace = true common-grpc.workspace = true @@ -24,10 +24,10 @@ common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true -chrono.workspace = true datafusion.workspace = true datatypes.workspace = true futures-util.workspace = true +fxhash = "0.2" humantime-serde.workspace = true itertools.workspace = true lazy_static = "1.4" @@ -45,18 +45,14 @@ tokio.workspace = true tracing.workspace = true [dev-dependencies] +common-function = { workspace = true } common-meta = { workspace = true, features = ["testing"] } -common-test-util.workspace = true -mito2 = { workspace = true, features = ["test"] } common-wal = { workspace = true } criterion = { workspace = true, features = ["async", "async_tokio"] } +mito2 = { workspace = true, features = ["test"] } mur3 = "0.1" -common-function = { workspace = true } partition = { workspace = true } [[bench]] name = "bench_tsid_generator" harness = false - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 2be72af57cf..82a2884549a 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -20,7 +20,6 @@ api.workspace = true aquamarine.workspace = true arrow-schema.workspace = true async-channel = "1.9" -common-stat.workspace = true async-stream.workspace = true async-trait.workspace = true base64.workspace = true @@ -28,58 +27,55 @@ bytemuck.workspace = true bytes.workspace = true chrono.workspace = true common-base.workspace = true -common-config.workspace = true common-datasource.workspace = true -common-decimal.workspace = true common-error.workspace = true -common-grpc.workspace = true common-function.workspace = true +common-grpc.workspace = true common-macro.workspace = true common-memory-manager.workspace = true common-meta.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true +common-stat.workspace = true common-telemetry.workspace = true common-test-util = { workspace = true, optional = true } common-time.workspace = true common-wal.workspace = true crc32fast = "1" crossbeam-utils.workspace = true +dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datatypes.workspace = true -dashmap.workspace = true derive_more.workspace = true dotenv.workspace = true either.workspace = true futures.workspace = true +greptime-proto.workspace = true humantime.workspace = true humantime-serde.workspace = true index.workspace = true itertools.workspace = true -greptime-proto.workspace = true lazy_static = "1.4" log-store = { workspace = true } mito-codec.workspace = true moka = { workspace = true, features = ["sync", "future"] } object-store = { workspace = true, features = ["testing"] } parquet = { workspace = true, features = ["async", "variant_experimental"] } -paste.workspace = true +partition.workspace = true pin-project.workspace = true prometheus.workspace = true prost.workspace = true -partition.workspace = true puffin.workspace = true rand.workspace = true rayon = "1.10" -roaring = { version = "0.10", optional = true } regex.workspace = true +roaring = { version = "0.10", optional = true } rskafka = { workspace = true, optional = true } rstest = { workspace = true, optional = true } rstest_reuse = { workspace = true, optional = true } -scopeguard = "1.2" serde.workspace = true serde_json.workspace = true serde_with.workspace = true @@ -108,7 +104,6 @@ object-store = { workspace = true, features = ["services-memory"] } rskafka.workspace = true rstest.workspace = true rstest_reuse.workspace = true -session.workspace = true toml.workspace = true [[bench]] @@ -159,6 +154,3 @@ required-features = ["testing"] name = "bench_dense_index_update" harness = false required-features = ["testing"] - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index 412f3563298..4083271a967 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -10,23 +10,18 @@ workspace = true [features] mysql-object-store = ["opendal/services-mysql"] services-memory = ["opendal/services-memory"] -testing = ["derive_builder"] +testing = ["derive_builder", "uuid"] [dependencies] -async-trait.workspace = true bytes.workspace = true cap-std.workspace = true -chrono.workspace = true common-base.workspace = true common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -datafusion_object_store.workspace = true derive_builder = { workspace = true, optional = true } -futures.workspace = true humantime-serde.workspace = true -lazy_static.workspace = true opendal = { version = "0.59.2", features = [ "layers-tracing", "layers-prometheus", @@ -43,14 +38,16 @@ reqwest.workspace = true serde.workspace = true snafu.workspace = true tokio.workspace = true -uuid.workspace = true +uuid = { workspace = true, optional = true } [dev-dependencies] anyhow = "1.0" common-telemetry.workspace = true common-test-util.workspace = true +datafusion_object_store.workspace = true +futures.workspace = true +object-store = { workspace = true, features = ["testing"] } object_store_opendal.workspace = true -rand.workspace = true tempfile.workspace = true tokio.workspace = true toml.workspace = true diff --git a/src/object-store/src/lib.rs b/src/object-store/src/lib.rs index 05c344485c6..2de2b570082 100644 --- a/src/object-store/src/lib.rs +++ b/src/object-store/src/lib.rs @@ -26,6 +26,7 @@ pub mod layers; pub mod manager; mod metrics; pub mod secure_fs; +#[cfg(feature = "testing")] pub mod test_util; pub mod util; diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index ea1bc990f74..72f953e0f89 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -22,7 +22,6 @@ workspace = true ahash.workspace = true api.workspace = true arrow.workspace = true -arrow-ipc.workspace = true async-stream.workspace = true async-trait.workspace = true bytes.workspace = true @@ -66,7 +65,6 @@ object-store.workspace = true partition.workspace = true percent-encoding.workspace = true prometheus.workspace = true -prost.workspace = true prost-types = { workspace = true, optional = true } query.workspace = true regex.workspace = true diff --git a/src/partition/Cargo.toml b/src/partition/Cargo.toml index a8e3a8ae11a..6773cc5878b 100644 --- a/src/partition/Cargo.toml +++ b/src/partition/Cargo.toml @@ -9,18 +9,14 @@ workspace = true [dependencies] api.workspace = true -async-trait.workspace = true common-base.workspace = true common-error.workspace = true common-macro.workspace = true common-meta.workspace = true -common-query.workspace = true -common-telemetry.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true datafusion-physical-expr.workspace = true datatypes.workspace = true -itertools.workspace = true moka = { workspace = true, features = ["future"] } serde.workspace = true serde_json.workspace = true diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 620ac36ac77..9a355f4c8dc 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -14,28 +14,21 @@ ahash.workspace = true api.workspace = true arrow.workspace = true arrow-schema.workspace = true -async-trait.workspace = true catalog.workspace = true chrono.workspace = true chrono-tz.workspace = true common-catalog.workspace = true common-error.workspace = true -common-function.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true common-recordbatch.workspace = true -common-runtime.workspace = true common-telemetry.workspace = true common-time.workspace = true -crossbeam-utils.workspace = true csv = "1.3" -dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true -datafusion-functions.workspace = true -datafusion-physical-expr.workspace = true datatypes.workspace = true dyn-fmt = "0.4" enum_dispatch = "0.3" @@ -50,7 +43,6 @@ moka = { workspace = true, features = ["future"] } once_cell.workspace = true operator.workspace = true ordered-float.workspace = true -paste.workspace = true prometheus.workspace = true query.workspace = true regex.workspace = true @@ -68,7 +60,6 @@ yaml-rust = "0.4" [dev-dependencies] catalog = { workspace = true, features = ["testing"] } criterion = { workspace = true, features = ["html_reports"] } -rayon = "1.0" session = { workspace = true, features = ["testing"] } [[bench]] diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index 480c95a5adb..e81d93b7781 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -16,7 +16,6 @@ vector_index = [] [dependencies] ahash.workspace = true api.workspace = true -arc-swap = "1.0" arrow.workspace = true arrow-schema.workspace = true async-recursion = "1.0" @@ -50,11 +49,9 @@ datafusion-physical-expr.workspace = true datafusion-proto.workspace = true datafusion-sql.workspace = true datatypes.workspace = true -either.workspace = true futures.workspace = true futures-util.workspace = true greptime-proto.workspace = true -humantime.workspace = true itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true @@ -63,7 +60,6 @@ meter-core.workspace = true meter-macros.workspace = true object-store.workspace = true once_cell.workspace = true -parking_lot.workspace = true partition.workspace = true prometheus.workspace = true promql.workspace = true @@ -93,13 +89,10 @@ common-query = { workspace = true, features = ["testing"] } fastrand = "2.0" moka = { workspace = true, features = ["future"] } nalgebra.workspace = true -num = "0.4" -num-traits = "0.2" paste.workspace = true pretty_assertions.workspace = true rand.workspace = true session = { workspace = true, features = ["testing"] } store-api.workspace = true table = { workspace = true, features = ["testing"] } -tokio-stream.workspace = true toml.workspace = true diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 0e2d896db8c..06648254a06 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -40,7 +40,6 @@ common-batcher.workspace = true common-catalog.workspace = true common-decimal.workspace = true common-error.workspace = true -common-frontend.workspace = true common-grpc.workspace = true common-macro.workspace = true common-mem-prof = { workspace = true, optional = true } @@ -58,7 +57,6 @@ common-telemetry.workspace = true common-time.workspace = true common-version = { workspace = true, features = ["codec"] } csv = "1.3" -dashmap.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true @@ -123,7 +121,6 @@ serde.workspace = true serde_json.workspace = true session.workspace = true simd-json.workspace = true -simdutf8 = "0.1" smallvec.workspace = true snafu.workspace = true snap = "1" @@ -169,6 +166,7 @@ rand.workspace = true serde_json.workspace = true servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } +simdutf8 = "0.1" table.workspace = true tempfile = "3.0.0" tokio = { workspace = true, features = ["test-util"] } diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 889b911c3cc..01ef368f1f0 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -29,3 +29,9 @@ derive_more.workspace = true snafu.workspace = true sql.workspace = true uuid.workspace = true + +[package.metadata.cargo-shear] +# `common-error` is required by `common-macro`'s `stack_trace_debug`: the +# generated code references `::common_error` paths, so it must stay a direct +# dependency. +ignored = ["common-error"] diff --git a/src/sql/Cargo.toml b/src/sql/Cargo.toml index 1ab966fc7e3..685d96336b1 100644 --- a/src/sql/Cargo.toml +++ b/src/sql/Cargo.toml @@ -16,25 +16,19 @@ arrow-buffer.workspace = true chrono.workspace = true common-base.workspace = true common-catalog.workspace = true -common-decimal.workspace = true common-error.workspace = true common-macro.workspace = true common-query.workspace = true common-sql.workspace = true -common-telemetry.workspace = true common-time.workspace = true datafusion.workspace = true datafusion-common.workspace = true datafusion-expr.workspace = true -datafusion-physical-expr.workspace = true datafusion-sql.workspace = true datatypes.workspace = true either.workspace = true -hex = "0.4" -humantime.workspace = true iso8601 = "0.6.1" itertools.workspace = true -jsonb.workspace = true lazy_static.workspace = true promql-parser.workspace = true regex.workspace = true @@ -43,7 +37,6 @@ serde_json.workspace = true snafu.workspace = true sqlparser.workspace = true sqlparser_derive = "0.1" -store-api.workspace = true table.workspace = true uuid.workspace = true diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml index 2e5a203912e..4a3616afb8a 100644 --- a/src/store-api/Cargo.toml +++ b/src/store-api/Cargo.toml @@ -36,15 +36,10 @@ prost.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true -sqlparser.workspace = true strum.workspace = true tokio.workspace = true uuid.workspace = true [dev-dependencies] -async-stream.workspace = true common-meta.workspace = true serde_json.workspace = true - -[package.metadata.cargo-udeps.ignore] -normal = ["aquamarine"] diff --git a/src/table/Cargo.toml b/src/table/Cargo.toml index 8ecc51a320c..357d005e612 100644 --- a/src/table/Cargo.toml +++ b/src/table/Cargo.toml @@ -34,13 +34,11 @@ derive_builder.workspace = true futures.workspace = true greptime-proto.workspace = true humantime.workspace = true -humantime-serde.workspace = true lazy_static.workspace = true once_cell.workspace = true paste.workspace = true serde.workspace = true snafu.workspace = true -sqlparser.workspace = true store-api.workspace = true tokio.workspace = true @@ -48,4 +46,3 @@ tokio.workspace = true common-test-util.workspace = true parquet = { workspace = true, features = ["async"] } serde_json.workspace = true -tokio-util = { version = "0.7", features = ["compat"] } diff --git a/tests-fuzz/Cargo.toml b/tests-fuzz/Cargo.toml index 2cd66b8afe5..1189f53c42f 100644 --- a/tests-fuzz/Cargo.toml +++ b/tests-fuzz/Cargo.toml @@ -188,3 +188,12 @@ path = "utils/kafka_wal_helper.rs" test = false bench = false doc = false + +[package.metadata.cargo-shear] +# `common-error` is required by `common-macro`'s `stack_trace_debug`: the +# generated code references `::common_error` paths, so it must stay a direct +# dependency. +# `k8s-openapi` is a feature-pinning dependency: it enables the `v1_30` +# version feature required by the transitive `kube` -> `k8s-openapi` chain, +# but is never referenced by name in source code. +ignored = ["common-error", "k8s-openapi"] diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index a5ac94924e4..a3d694f758b 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -20,16 +20,11 @@ workspace = true [dependencies] api.workspace = true arrow-flight.workspace = true -async-stream.workspace = true async-trait.workspace = true auth.workspace = true axum.workspace = true -base64.workspace = true cache.workspace = true catalog.workspace = true -chrono.workspace = true -clap.workspace = true -cli.workspace = true client.workspace = true cmd.workspace = true common-base.workspace = true @@ -38,16 +33,12 @@ common-config.workspace = true common-datasource.workspace = true common-error.workspace = true common-event-recorder.workspace = true -common-frontend.workspace = true common-grpc.workspace = true -common-memory-manager.workspace = true common-meta.workspace = true -common-options.workspace = true common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true -common-stat.workspace = true common-telemetry.workspace = true common-test-util.workspace = true common-time.workspace = true @@ -55,62 +46,59 @@ common-wal.workspace = true datanode = { workspace = true } datatypes.workspace = true dotenv.workspace = true -flate2 = "1.0" flow.workspace = true frontend.workspace = true futures.workspace = true futures-util.workspace = true http.workspace = true hyper-util = { workspace = true, features = ["tokio"] } -log-query = { workspace = true } -loki-proto.workspace = true meta-client.workspace = true meta-srv = { workspace = true, features = ["mock"] } mito2.workspace = true object-store.workspace = true operator.workspace = true -plugins.workspace = true prost.workspace = true query.workspace = true rand.workspace = true serde_json.workspace = true servers.workspace = true session.workspace = true -snafu.workspace = true sql.workspace = true -sqlx = { workspace = true, features = [ - "mysql", - "postgres", - "chrono", -] } standalone.workspace = true store-api.workspace = true substrait.workspace = true table.workspace = true tempfile.workspace = true -time = "0.3" tokio.workspace = true tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true tower.workspace = true uuid.workspace = true -zstd.workspace = true [dev-dependencies] +base64.workspace = true catalog = { workspace = true, features = ["testing"] } +chrono.workspace = true +clap.workspace = true cli.workspace = true client = { workspace = true, features = ["testing"] } +common-frontend.workspace = true +common-memory-manager.workspace = true common-meta = { workspace = true, features = ["testing"] } +common-options.workspace = true common-procedure-test.workspace = true common-wal = { workspace = true, features = ["testing"] } datafusion.workspace = true datafusion-expr.workspace = true either.workspace = true +flate2 = "1.0" frontend = { workspace = true, features = ["testing"] } hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true +log-query.workspace = true +loki-proto.workspace = true mito2 = { workspace = true, features = ["test", "test-shared-fs-region-migration"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", @@ -119,9 +107,9 @@ mysql_async = { version = "0.37", default-features = false, features = [ opentelemetry-proto.workspace = true operator = { workspace = true, features = ["testing"] } otel-arrow-rust.workspace = true -partition.workspace = true paste.workspace = true pipeline.workspace = true +plugins.workspace = true prost.workspace = true rand.workspace = true rskafka.workspace = true @@ -131,9 +119,16 @@ sea-query.workspace = true servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } similar-asserts.workspace = true +sqlx = { workspace = true, features = [ + "mysql", + "postgres", + "chrono", +] } store-api.workspace = true +time = "0.3" tokio-postgres = { workspace = true } tokio-util.workspace = true url = "2.3" urlencoding = "2.1" yaml-rust = "0.4" +zstd.workspace = true diff --git a/tests/runner/Cargo.toml b/tests/runner/Cargo.toml index 2941e85fa5a..35de1443aa8 100644 --- a/tests/runner/Cargo.toml +++ b/tests/runner/Cargo.toml @@ -13,7 +13,6 @@ clap.workspace = true client.workspace = true common-error.workspace = true common-query.workspace = true -common-recordbatch.workspace = true common-time.workspace = true datatypes = { workspace = true } flate2 = "1.0" From a35f9d5c782ddccfe7c94fe9c7f291a6b7910518 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 23 Sep 2026 05:16:41 +0000 Subject: [PATCH 017/133] fix: address Windows test failures and run full Windows CI (#9305) * fix: use relative object keys for Windows filesystem access Signed-off-by: WenyXu * fix: cover Windows path and time limits in full test CI Signed-off-by: WenyXu * fix: use relative keys in metadata snapshot tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/workflows/rust.yml | 42 ++++ src/cli/src/metadata/snapshot.rs | 218 ++++++++++++++---- src/common/datasource/src/file_format/csv.rs | 19 +- src/common/datasource/src/file_format/json.rs | 11 +- .../datasource/src/file_format/tests.rs | 28 +-- src/common/datasource/src/test_util.rs | 31 +-- src/common/datasource/src/tests.rs | 25 +- src/common/meta/src/snapshot.rs | 40 +--- src/mito2/src/cache/write_cache.rs | 3 +- src/query/src/promql/label_values.rs | 29 ++- 10 files changed, 303 insertions(+), 143 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 1326ab5fd3d..5778d3856bb 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -148,6 +148,48 @@ jobs: env: RUSTFLAGS: "-D warnings" + test-windows: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + name: Test Windows + runs-on: windows-2022 + timeout-minutes: 90 + env: + CARGO_INCREMENTAL: 0 + CARGO_PROFILE_DEV_DEBUG: 0 + CARGO_PROFILE_TEST_DEBUG: 0 + RUSTFLAGS: "-D warnings -C linker=lld-link" + RUST_BACKTRACE: 1 + RUST_MIN_STACK: 8388608 # 8MB + steps: + - run: git config --global core.autocrlf false + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: ./.github/actions/setup-cyborg + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - uses: KyleMayes/install-llvm-action@v1 + with: + version: "14.0" + - uses: actions-rust-lang/setup-rust-toolchain@v1 + with: + cache: false + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "test-windows" + cache-all-crates: "true" + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Cargo Nextest + uses: taiki-e/install-action@nextest + - name: Install WSL distribution + uses: Vampire/setup-wsl@v2 + with: + distribution: Ubuntu-22.04 + - name: Run Windows tests + run: cargo nextest run --locked --workspace -F dashboard --no-fail-fast + conflict-check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} name: Check for conflict diff --git a/src/cli/src/metadata/snapshot.rs b/src/cli/src/metadata/snapshot.rs index 59b2599ad9c..9c010c2235b 100644 --- a/src/cli/src/metadata/snapshot.rs +++ b/src/cli/src/metadata/snapshot.rs @@ -12,11 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::path::{Component, Path, PathBuf}; + use async_trait::async_trait; use clap::{Parser, Subcommand}; use common_error::ext::BoxedError; use common_meta::snapshot::MetadataSnapshotManager; -use object_store::{ObjectStore, services}; +use object_store::ObjectStore; use crate::Tool; use crate::common::{ObjectStoreConfig, StoreConfig, new_fs_object_store}; @@ -270,25 +272,44 @@ fn build_object_store_and_resolve_file_path( fs_root: &str, file_path: &str, ) -> Result<(ObjectStore, String), BoxedError> { - let object_store = object_store.build().map_err(BoxedError::new)?; - let object_store = match object_store { - Some(object_store) => object_store, - None => new_fs_object_store(fs_root)?, - }; + if let Some(object_store) = object_store.build().map_err(BoxedError::new)? { + return Ok((object_store, file_path.to_string())); + } - let file_path = if object_store.info().scheme() == services::FS_SCHEME { - resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)? + let absolute_path = + resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)?; + let (root, key) = split_fs_snapshot_path(fs_root, &absolute_path); + Ok((new_fs_object_store(&root)?, key)) +} + +/// Separates a local absolute path into an operator root and a relative key. +fn split_fs_snapshot_path(fs_root: &str, absolute_path: &str) -> (String, String) { + let path = Path::new(absolute_path); + // OpenDAL keys cannot contain a Windows drive/UNC prefix or a root directory. + // Keep the filesystem anchor in the operator root, not in the object key. + let anchor: PathBuf = path + .components() + .take_while(|c| matches!(c, Component::Prefix(_) | Component::RootDir)) + .collect(); + let key = path + .components() + .filter(|c| !matches!(c, Component::Prefix(_) | Component::RootDir)) + .map(|c| c.as_os_str().to_string_lossy()) + .collect::>() + .join("/"); + // Preserve the existing custom-root semantics: cwd-resolved keys are rooted + // under --dir. The default root follows the input path's drive on Windows. + let root = if fs_root == "/" { + anchor.to_string_lossy() } else { - file_path.to_string() + fs_root.into() }; - - Ok((object_store, file_path)) + (root.into_owned(), key) } #[cfg(test)] mod tests { use std::env; - use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -311,44 +332,159 @@ mod tests { #[tokio::test] async fn test_cmd_resolve_file_path() { - common_telemetry::init_default_ut_logging(); - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); let current_dir = env::current_dir().unwrap(); - let file_path = current_dir.join("metadata_snapshot.metadata.fb"); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + for file_path in [ + "metadata_snapshot.metadata.fb".to_string(), + current_dir + .join("metadata_snapshot.metadata.fb") + .display() + .to_string(), + ] { + let cmd = RestoreCommand::parse_from([ + "", + "--file_name", + &file_path, + "--backend", + "memory-store", + "--store-addrs", + "memory://", + ]); + let tool = cmd.build().await.unwrap(); + assert!(!Path::new(&tool.file_path).is_absolute()); + assert!( + Path::new(&tool.file_path) + .components() + .all(|c| matches!(c, Component::Normal(_))) + ); + let expected = current_dir + .components() + .filter_map(|c| match c { + Component::Normal(part) => Some(part.to_str().unwrap()), + _ => None, + }) + .chain(["metadata_snapshot.metadata.fb"]) + .collect::>() + .join("/"); + assert_eq!(tool.file_path, expected); + } + } + #[cfg(windows)] + #[test] + fn test_split_windows_snapshot_path() { + for (path, root, key) in [ + ( + r"C:\backup\snapshot.metadata.fb", + r"C:\", + "backup/snapshot.metadata.fb", + ), + ( + r"D:\备份 dir\snapshot.metadata.fb", + r"D:\", + "备份 dir/snapshot.metadata.fb", + ), + ( + r"\\server\share\backup\snapshot.metadata.fb", + r"\\server\share\", + "backup/snapshot.metadata.fb", + ), + ( + r"\\?\C:\backup\snapshot.metadata.fb", + r"\\?\C:\", + "backup/snapshot.metadata.fb", + ), + ] { + assert_eq!( + split_fs_snapshot_path("/", path), + (root.to_string(), key.to_string()) + ); + assert_eq!( + split_fs_snapshot_path(r"D:\root", path), + (r"D:\root".to_string(), key.to_string()) + ); + } + } + + #[tokio::test] + async fn test_fs_snapshot_paths() { + let temp_dir = tempfile::tempdir().unwrap(); + for name in ["snapshot.metadata.fb", "nested dir/å¿«ç…§.metadata.fb"] { + let path = temp_dir.path().join(name); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + "/", + path.to_str().unwrap(), + ) + .unwrap(); + assert!(!path.exists()); + store.write(&key, "snapshot").await.unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), b"snapshot"); + assert_eq!(store.read(&key).await.unwrap().to_bytes(), &b"snapshot"[..]); + } + } + + #[tokio::test] + async fn test_fs_snapshot_custom_root() { + let temp_dir = tempfile::tempdir().unwrap(); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + temp_dir.path().to_str().unwrap(), + "snapshot.metadata.fb", + ) + .unwrap(); + let current_dir = env::current_dir().unwrap(); + let expected = current_dir + .components() + .filter_map(|c| match c { + Component::Normal(part) => Some(part.to_str().unwrap()), + _ => None, + }) + .chain(["snapshot.metadata.fb"]) + .collect::>() + .join("/"); + assert_eq!(key, expected); + store.write(&key, "snapshot").await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join(expected)).unwrap(), + b"snapshot" + ); + } + + #[tokio::test] + async fn test_fs_snapshot_rejects_parent_traversal() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("../snapshot.metadata.fb"); + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + "/", + path.to_str().unwrap(), + ) + .unwrap(); + assert!(key.split('/').any(|part| part == "..")); + let err = store.write(&key, "snapshot").await.unwrap_err(); + assert_eq!(err.kind(), object_store::ErrorKind::NotFound); + assert!(err.to_string().contains("path escapes the configured root")); + } + + #[test] + fn test_remote_snapshot_key_unchanged() { let cmd = RestoreCommand::parse_from([ "", - "--file_name", - "metadata_snapshot.metadata.fb", "--backend", "memory-store", "--store-addrs", "memory://", + "--s3", + "--s3-bucket", + "snapshots", + "--s3-region", + "us-east-1", ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); - - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + let key = "nested/snapshot.metadata.fb"; + let (store, resolved) = + build_object_store_and_resolve_file_path(cmd.object_store, "/ignored", key).unwrap(); + assert_eq!(store.info().scheme(), object_store::services::S3_SCHEME); + assert_eq!(resolved, key); } async fn setup_backup_file(object_store: ObjectStore, file_path: &str) { diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index d016d4eabe9..d4cfb45088c 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -647,16 +647,15 @@ mod tests { let compressed_file_name = format!("test_compressed_csv.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export CSV with compression let rows = stream_to_csv( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -708,13 +707,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( CsvFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -722,7 +721,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, csv_source.clone(), None, compression_type, @@ -761,13 +760,13 @@ mod tests { ) .unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) @@ -793,13 +792,13 @@ mod tests { let csv_file_path = temp_dir.path().join("input.csv"); std::fs::write(&csv_file_path, "id,name,value\n1,Alice,10.5\n2,Bob\n").unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) diff --git a/src/common/datasource/src/file_format/json.rs b/src/common/datasource/src/file_format/json.rs index 3c24a4b7e27..23ce461fcb2 100644 --- a/src/common/datasource/src/file_format/json.rs +++ b/src/common/datasource/src/file_format/json.rs @@ -313,16 +313,15 @@ mod tests { let compressed_file_name = format!("test_compressed_json.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export JSON with compression let rows = stream_to_json( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -374,13 +373,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( JsonFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -388,7 +387,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, json_source.clone(), None, compression_type, diff --git a/src/common/datasource/src/file_format/tests.rs b/src/common/datasource/src/file_format/tests.rs index 27b86e1b114..14be2b1777f 100644 --- a/src/common/datasource/src/file_format/tests.rs +++ b/src/common/datasource/src/file_format/tests.rs @@ -66,13 +66,12 @@ impl Test<'_> { #[tokio::test] async fn test_json_opener() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = basic_schema_with_time_format(); let file_source = Arc::new(JsonSource::new(schema)).with_batch_size(test_util::TEST_BATCH_SIZE); - let path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let path = "json/basic.json"; let tests = [ Test { @@ -108,11 +107,10 @@ async fn test_json_opener() { #[tokio::test] async fn test_csv_opener() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = basic_schema_with_time_format(); - let path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); + let path = "csv/basic.csv"; let file_source = CsvSource::new(schema).with_batch_size(test_util::TEST_BATCH_SIZE); @@ -150,13 +148,12 @@ async fn test_csv_opener() { #[tokio::test(flavor = "multi_thread")] async fn test_parquet_exec() { - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = test_basic_schema(); - let path = &find_workspace_path("/src/common/datasource/tests/parquet/basic.parquet") - .display() - .to_string(); + let path = "parquet/basic.parquet"; let parquet_source = ParquetSource::new(schema) .with_parquet_file_reader_factory(Arc::new(DefaultParquetFileReaderFactory::new(store))); @@ -191,11 +188,10 @@ async fn test_parquet_exec() { #[tokio::test] async fn test_orc_opener() { - let path = &find_workspace_path("/src/common/datasource/tests/orc/test.orc") - .display() - .to_string(); + let path = "orc/test.orc"; - let store = test_store("/"); + let root = find_workspace_path("/src/common/datasource/tests"); + let store = test_store(root.to_str().unwrap()); let schema = Arc::new(OrcFormat.infer_schema(&store, path).await.unwrap()); let file_source = Arc::new(OrcSource::new(schema.into())); diff --git a/src/common/datasource/src/test_util.rs b/src/common/datasource/src/test_util.rs index 4cefda88890..f95d24930b2 100644 --- a/src/common/datasource/src/test_util.rs +++ b/src/common/datasource/src/test_util.rs @@ -56,8 +56,8 @@ pub fn test_store(root: &str) -> ObjectStore { pub fn test_tmp_store(root: &str) -> (ObjectStore, TempDir) { let dir = create_temp_dir(root); - let builder = Fs::default(); - (ObjectStore::new(builder.root("/")).unwrap(), dir) + let store = test_store(dir.path().to_str().unwrap()); + (store, dir) } pub fn test_basic_schema() -> SchemaRef { @@ -94,9 +94,11 @@ pub(crate) fn scan_config( .build() } -pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usize) -> usize) { - let store = test_store("/"); - +pub async fn setup_stream_to_json_test( + store: &ObjectStore, + origin_path: &str, + threshold: impl Fn(usize) -> usize, +) { let schema = basic_schema_with_time_format(); let json_opener = JsonOpener::new( @@ -117,9 +119,9 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_json"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_json"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let json_format = JsonFormat::default(); @@ -127,7 +129,7 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi stream_to_json( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &json_format, @@ -136,18 +138,17 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let origin = store.read(origin_path).await.unwrap(); assert_eq_lines(written.to_vec(), origin.to_vec()); } pub async fn setup_stream_to_csv_test( + store: &ObjectStore, origin_path: &str, format_path: &str, threshold: impl Fn(usize) -> usize, ) { - let store = test_store("/"); - let schema = basic_schema_with_time_format(); let csv_source = CsvSource::new(schema).with_batch_size(TEST_BATCH_SIZE); @@ -168,9 +169,9 @@ pub async fn setup_stream_to_csv_test( .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_csv"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_csv"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let csv_format = CsvFormat { timestamp_format: Some("%m-%d-%Y".to_string()), @@ -183,7 +184,7 @@ pub async fn setup_stream_to_csv_test( stream_to_csv( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &csv_format, @@ -192,7 +193,7 @@ pub async fn setup_stream_to_csv_test( .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let format_expect = store.read(format_path).await.unwrap(); assert_eq_lines(written.to_vec(), format_expect.to_vec()); } diff --git a/src/common/datasource/src/tests.rs b/src/common/datasource/src/tests.rs index 8be9add7c88..b983ef86627 100644 --- a/src/common/datasource/src/tests.rs +++ b/src/common/datasource/src/tests.rs @@ -18,34 +18,31 @@ use crate::test_util; #[tokio::test] async fn test_stream_to_json() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let root = find_workspace_path("/src/common/datasource/tests/json"); + let store = test_util::test_store(root.to_str().unwrap()); + let origin_path = "basic.json"; // A small threshold // Triggers the flush each writes - test_util::setup_stream_to_json_test(origin_path, |size| size / 2).await; + test_util::setup_stream_to_json_test(&store, origin_path, |size| size / 2).await; // A large threshold // Only triggers the flush at last - test_util::setup_stream_to_json_test(origin_path, |size| size * 2).await; + test_util::setup_stream_to_json_test(&store, origin_path, |size| size * 2).await; } #[tokio::test] async fn test_stream_to_csv() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); - - let format_path = &find_workspace_path("/src/common/datasource/tests/csv/basic_format.csv") - .display() - .to_string(); + let root = find_workspace_path("/src/common/datasource/tests/csv"); + let store = test_util::test_store(root.to_str().unwrap()); + let origin_path = "basic.csv"; + let format_path = "basic_format.csv"; // A small threshold // Triggers the flush each writes - test_util::setup_stream_to_csv_test(origin_path, format_path, |size| size / 2).await; + test_util::setup_stream_to_csv_test(&store, origin_path, format_path, |size| size / 2).await; // A large threshold // Only triggers the flush at last - test_util::setup_stream_to_csv_test(origin_path, format_path, |size| size * 2).await; + test_util::setup_stream_to_csv_test(&store, origin_path, format_path, |size| size * 2).await; } diff --git a/src/common/meta/src/snapshot.rs b/src/common/meta/src/snapshot.rs index f35c86e2e7b..3582b0acecb 100644 --- a/src/common/meta/src/snapshot.rs +++ b/src/common/meta/src/snapshot.rs @@ -408,7 +408,6 @@ mod tests { async fn test_dump_and_restore() { common_telemetry::init_default_ut_logging(); let (temp_dir, kv_backend, manager) = test_env("test_dump_and_restore"); - let temp_path = temp_dir.path(); for i in 0..10 { kv_backend @@ -420,32 +419,16 @@ mod tests { .await .unwrap(); } - let dump_path = temp_path.join("snapshot"); - manager - .dump(&format!( - "{}/metadata_snapshot", - &dump_path.as_path().display().to_string() - )) - .await - .unwrap(); + manager.dump("snapshot/metadata_snapshot").await.unwrap(); + let snapshot_key = "snapshot/metadata_snapshot.metadata.fb"; + assert!(temp_dir.path().join("data").join(snapshot_key).is_file()); // Clean up the kv backend kv_backend.clear(); - let err = manager - .dump(&format!( - "{}/metadata_snapshot.metadata.fb", - &dump_path.as_path().display().to_string() - )) - .await - .unwrap_err(); + let err = manager.dump(snapshot_key).await.unwrap_err(); assert_matches!(err, Error::Unexpected { .. }); assert!(err.to_string().contains("already exists")); - let restore_path = dump_path - .join("metadata_snapshot.metadata.fb") - .as_path() - .display() - .to_string(); - manager.restore(&restore_path).await.unwrap(); + manager.restore(snapshot_key).await.unwrap(); for i in 0..10 { let key = format!("test_{}", i); @@ -456,14 +439,11 @@ mod tests { #[tokio::test] async fn test_restore_from_nonexistent_file() { - let (temp_dir, _kv_backend, manager) = test_env("test_restore_from_nonexistent_file"); - let restore_path = temp_dir - .path() - .join("nonexistent.metadata.fb") - .as_path() - .display() - .to_string(); - let err = manager.restore(&restore_path).await.unwrap_err(); + let (_temp_dir, _kv_backend, manager) = test_env("test_restore_from_nonexistent_file"); + let err = manager + .restore("nonexistent.metadata.fb") + .await + .unwrap_err(); assert_matches!(err, Error::ReadObject { .. }) } } diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 0e96c2f16a2..87ab4f0081a 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -790,7 +790,8 @@ mod tests { async fn test_read_metadata_from_write_cache() { common_telemetry::init_default_ut_logging(); let mut env = TestEnv::new().await; - let data_home = env.data_home().display().to_string(); + // Object keys are relative to the mock store root, not native paths. + let data_home = "write-cache-test".to_string(); let mock_store = env.init_object_store_manager(); let local_dir = create_temp_dir(""); diff --git a/src/query/src/promql/label_values.rs b/src/query/src/promql/label_values.rs index 44c173477bb..841666758e4 100644 --- a/src/query/src/promql/label_values.rs +++ b/src/query/src/promql/label_values.rs @@ -37,15 +37,15 @@ fn millis_since_epoch(time: SystemTime) -> Result { Ok(duration) => (duration.as_millis(), false), Err(earlier) => (earlier.duration().as_millis(), true), }; - let millis = i64::try_from(millis) - .ok() - .with_context(|| SystemTimeOutOfRangeSnafu { time })?; + let millis = + signed_millis(millis, before_epoch).with_context(|| SystemTimeOutOfRangeSnafu { time })?; + Ok(Timestamp::new_millisecond(millis)) +} - Ok(Timestamp::new_millisecond(if before_epoch { - -millis - } else { - millis - })) +/// Converts a millisecond magnitude and epoch direction without platform time limits. +fn signed_millis(millis: u128, before_epoch: bool) -> Option { + let millis = i64::try_from(millis).ok()?; + Some(if before_epoch { -millis } else { millis }) } fn build_time_filter(time_index_expr: Expr, start: Timestamp, end: Timestamp) -> Expr { @@ -139,7 +139,16 @@ mod tests { #[test] fn millis_beyond_i64_are_rejected() { - let time = UNIX_EPOCH + Duration::from_secs(1 << 60); - assert!(millis_since_epoch(time).is_err()); + let limit = i64::MAX as u128; + assert_eq!(signed_millis(limit, false), Some(i64::MAX)); + assert_eq!(signed_millis(limit, true), Some(-i64::MAX)); + assert_eq!(signed_millis(limit + 1, false), None); + assert_eq!(signed_millis(limit + 1, true), None); + + // Windows cannot represent this SystemTime. The conversion boundary is + // tested above on every platform; exercise the wrapper where possible. + if let Some(time) = UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64 + 1)) { + assert!(millis_since_epoch(time).is_err()); + } } } From bc800714093e705e2722dc1212f6b0fc9d58dee3 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 23 Sep 2026 06:23:52 +0000 Subject: [PATCH 018/133] feat: expose region open failure metrics (#9283) Signed-off-by: evenyag --- src/datanode/src/metrics.rs | 10 +++++++ src/datanode/src/region_server.rs | 43 +++++++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/datanode/src/metrics.rs b/src/datanode/src/metrics.rs index 7573490c747..e9f947ef9da 100644 --- a/src/datanode/src/metrics.rs +++ b/src/datanode/src/metrics.rs @@ -85,6 +85,16 @@ lazy_static! { ) .unwrap(); + /// Failed region open attempts, including individual and batch requests. + /// Counts every request in an engine group when no individual results are available. + /// Failed retries count separately; explicitly ignored missing regions do not count. + pub static ref REGION_OPEN_FAILURES_TOTAL: IntCounterVec = register_int_counter_vec!( + "greptime_datanode_region_open_failures_total", + "Failed region open attempts, including retries and all requests in failed engine groups, excluding ignored missing regions", + &["engine", "status_code"] + ) + .unwrap(); + /// Total count of failed insert requests to region server, labeled with request /// type and status code. pub static ref REGION_SERVER_INSERT_FAIL_COUNT: IntCounterVec = register_int_counter_vec!( diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 527ab0491c0..aba09de8797 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1285,13 +1285,17 @@ impl RegionServerInner { requests: Vec<(RegionId, RegionOpenRequest)>, ignore_nonexistent_region: bool, ) -> Result> { + let request_count = requests.len(); let region_changes = requests .iter() .map(|(region_id, open)| { let attribute = parse_region_attribute(&open.engine, &open.options)?; Ok((*region_id, RegionChange::Register(attribute))) }) - .collect::>>()?; + .collect::>>() + .inspect_err(|err| { + record_region_open_failures(engine.name(), err.status_code(), request_count); + })?; for (®ion_id, region_change) in ®ion_changes { self.set_region_status_not_ready(region_id, &engine, region_change) @@ -1314,6 +1318,7 @@ impl RegionServerInner { .await { error!(e; "Failed to set region to ready: {}", region_id); + record_region_open_failures(engine.name(), e.status_code(), 1); errors.push(BoxedError::new(e)); } else { open_regions.push(region_id) @@ -1327,6 +1332,7 @@ impl RegionServerInner { warn!("Region {} not found, ignore it, source: {:?}", region_id, e); } else { error!(e; "Failed to open region: {}", region_id); + record_region_open_failures(engine.name(), e.status_code(), 1); errors.push(e); } } @@ -1338,6 +1344,7 @@ impl RegionServerInner { self.unset_region_status(region_id, &engine, *region_change); } error!(e; "Failed to open batch regions"); + record_region_open_failures(engine.name(), e.status_code(), request_count); errors.push(BoxedError::new(e)); } } @@ -1373,7 +1380,10 @@ impl RegionServerInner { .read() .unwrap() .get(&engine) - .with_context(|| RegionEngineNotFoundSnafu { name: &engine })? + .with_context(|| RegionEngineNotFoundSnafu { name: &engine }) + .inspect_err(|err| { + record_region_open_failures(&engine, err.status_code(), requests.len()); + })? .clone(); results.push( self.handle_batch_open_requests_inner( @@ -1574,6 +1584,24 @@ impl RegionServerInner { &self, region_id: RegionId, request: RegionRequest, + ) -> Result { + let open_engine = match &request { + RegionRequest::Open(open) => Some(open.engine.clone()), + _ => None, + }; + self.handle_request_inner(region_id, request) + .await + .inspect_err(|err| { + if let Some(engine) = open_engine { + record_region_open_failures(&engine, err.status_code(), 1); + } + }) + } + + async fn handle_request_inner( + &self, + region_id: RegionId, + request: RegionRequest, ) -> Result { let request_type = request.request_type(); let _timer = crate::metrics::HANDLE_REGION_REQUEST_ELAPSED @@ -1946,6 +1974,17 @@ fn is_metric_engine(engine: &str) -> bool { engine == METRIC_ENGINE_NAME } +/// Records failed open attempts without using arbitrary engine names as metric labels. +fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usize) { + let engine = match engine { + MITO_ENGINE_NAME | METRIC_ENGINE_NAME | FILE_ENGINE_NAME => engine, + _ => "unknown", + }; + crate::metrics::REGION_OPEN_FAILURES_TOTAL + .with_label_values(&[engine, status_code.as_ref()]) + .inc_by(count as u64); +} + fn parse_region_attribute( engine: &str, options: &HashMap, From 5a78e07be6269baa993dc5e25978d68aecd9886d Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:19:18 +0000 Subject: [PATCH 019/133] fix(json2): restrict JSON2 type hints (#9316) * fix(sql): restrict JSON2 type hints Signed-off-by: fys * fix(sql): allow equivalent 64-bit aliases in JSON2 type hints Signed-off-by: fys * fix(sql): support UInt64 conversion and preserve JSON2 hint aliases Signed-off-by: fys * refactor(sql): remove redundant JSON2 type hint normalization Signed-off-by: fys * fix(datatypes): restrict JSON2 hint types in JsonSettings::try_new Signed-off-by: fys --------- Signed-off-by: fys --- src/datatypes/src/json.rs | 11 +- src/sql/src/parsers/create_parser/json.rs | 162 ++++++++++-------- src/sql/src/statements.rs | 5 +- src/sql/src/statements/create.rs | 62 ++----- .../common/types/json/json2_alter.result | 2 +- .../common/types/json/json2_alter.sql | 2 +- .../json2_alter_type_hints_compaction.result | 2 +- .../json2_alter_type_hints_compaction.sql | 2 +- ...er_type_hints_compaction_conversion.result | 2 +- ...alter_type_hints_compaction_conversion.sql | 2 +- ..._compaction_null_type_hint_mismatch.result | 2 +- ...on2_compaction_null_type_hint_mismatch.sql | 2 +- .../common/types/json/json2_type_hints.result | 131 +++++++++++++- .../common/types/json/json2_type_hints.sql | 45 +++++ 14 files changed, 300 insertions(+), 132 deletions(-) diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index cd73e202bf3..e3c408698eb 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -237,19 +237,10 @@ fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { } let data_type = match &hint.data_type { ConcreteDataType::Boolean(_) - | ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) | ConcreteDataType::UInt64(_) - | ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) | ConcreteDataType::Int64(_) - | ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) - | ConcreteDataType::String(_) - | ConcreteDataType::List(_) - | ConcreteDataType::Struct(_) => (&hint.data_type).into(), + | ConcreteDataType::String(_) => (&hint.data_type).into(), data_type => { return InvalidJson2SettingsSnafu { reason: format!("unsupported JSON2 type hint data type: {data_type}"), diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 2bf3c9407a6..0adadc1d193 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -25,7 +25,6 @@ use crate::dialect::GreptimeDbDialect; use crate::error::{InvalidSqlSnafu, Result, SyntaxSnafu}; use crate::parsers::create_parser::{INVERTED, SKIPPING}; use crate::statements::create::{Json2Options, JsonTypeHint}; -use crate::statements::transform::type_alias::get_type_by_alias; const JSON2_TYPE_NAME: &str = "JSON2"; const MAX_AUTO_EXPANDED_PATHS: &str = "max_auto_expanded_paths"; @@ -155,7 +154,7 @@ fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { } ); let data_type = parser.parse_data_type().context(SyntaxSnafu)?; - let data_type = normalize_json2_type_hint_type(data_type)?; + let data_type = validate_json2_type_hint_type(data_type)?; let mut inverted_index = false; @@ -238,35 +237,23 @@ fn parse_json2_path(parser: &mut Parser<'_>) -> Result> { Ok(path) } -fn normalize_json2_type_hint_type(data_type: DataType) -> Result { - let data_type = get_type_by_alias(&data_type).unwrap_or(data_type); - let normalized = match data_type { - DataType::String(_) | DataType::Text | DataType::Varchar(_) | DataType::Char(_) => { - DataType::String(None) +fn validate_json2_type_hint_type(data_type: DataType) -> Result { + match data_type { + DataType::Int64 + | DataType::UInt64 + | DataType::Float64 + | DataType::String(None) + | DataType::BigInt(None) + | DataType::BigIntUnsigned(None) + | DataType::Double(ExactNumberInfo::None) + | DataType::Boolean => Ok(data_type), + _ => InvalidSqlSnafu { + msg: format!( + "unsupported JSON2 type hint data type: {data_type}; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64" + ), } - DataType::TinyInt(_) - | DataType::SmallInt(_) - | DataType::Int(_) - | DataType::Integer(_) - | DataType::BigInt(_) => DataType::BigInt(None), - DataType::TinyIntUnsigned(_) - | DataType::SmallIntUnsigned(_) - | DataType::IntUnsigned(_) - | DataType::UnsignedInteger - | DataType::BigIntUnsigned(_) => DataType::BigIntUnsigned(None), - DataType::Float(_) | DataType::Real | DataType::Double(_) => { - DataType::Double(ExactNumberInfo::None) - } - DataType::Boolean => DataType::Boolean, - _ => { - return InvalidSqlSnafu { - msg: format!("unsupported JSON2 type hint data type: {data_type}"), - } - .fail(); - } - }; - - Ok(normalized) + .fail(), + } } fn ensure_no_path_conflict(hints: &[JsonTypeHint], path: &[String]) -> Result<()> { @@ -330,7 +317,7 @@ CREATE TABLE traces ( log_json_data JSON2 ( "service.name" STRING INVERTED INDEX, http.method STRING, - status_code INT64, + status_code BIGINT, comment STRING, ), ts TIMESTAMP TIME INDEX, @@ -463,40 +450,81 @@ CREATE TABLE traces ( } #[test] - fn test_parse_json2_type_hint_normalizes_numeric_types() { - let column = parse_json2_column( - r#" -CREATE TABLE traces ( - log_json_data JSON2 ( - tinyint_value TINYINT, - smallint_value SMALLINT, - int_value INT, - integer_value INTEGER, - bigint_value BIGINT, - int64_value INT64, - tinyuint_value TINYINT UNSIGNED, - smalluint_value SMALLINT UNSIGNED, - uint_value INT UNSIGNED, - uint64_value UINT64, - float_value FLOAT, - real_value REAL, - double_value DOUBLE, - float64_value FLOAT64 - ), - ts TIMESTAMP TIME INDEX, -)"#, - ); + fn test_parse_json2_type_hint_supported_types() { + for (sql_type, expected) in [ + ("STRING", DataType::String(None)), + ("BIGINT", DataType::BigInt(None)), + ("BIGINT UNSIGNED", DataType::BigIntUnsigned(None)), + ("DOUBLE", DataType::Double(ExactNumberInfo::None)), + ("BOOLEAN", DataType::Boolean), + ("INT64", DataType::Int64), + ("UINT64", DataType::UInt64), + ("FLOAT64", DataType::Float64), + ] { + for sql_type in [sql_type.to_string(), sql_type.to_lowercase()] { + let column = parse_json2_column(&format!( + "CREATE TABLE traces (j JSON2 (value {sql_type}), ts TIMESTAMP TIME INDEX)" + )); + let options = column.extensions.json2_options.unwrap(); + let settings = options.build_json_settings().unwrap(); + assert_eq!(settings.type_hints().len(), 1); + let hints = options.type_hints; + assert_eq!(hints[0].data_type, expected); + } + } + } - let hints = column.extensions.json2_options.unwrap().type_hints; - assert_eq!(hints.len(), 14); - for hint in hints.iter().take(6) { - assert_eq!(hint.data_type, DataType::BigInt(None)); - } - for hint in hints.iter().skip(6).take(4) { - assert_eq!(hint.data_type, DataType::BigIntUnsigned(None)); - } - for hint in hints.iter().skip(10) { - assert_eq!(hint.data_type, DataType::Double(ExactNumberInfo::None)); + #[test] + fn test_parse_json2_type_hint_rejects_unsupported_types() { + for sql_type in [ + "INT2", + "INT4", + "INT8", + "INT16", + "INT32", + "TINYINT", + "SMALLINT", + "INT", + "INTEGER", + "UINT8", + "UINT16", + "UINT32", + "TINYINT UNSIGNED", + "SMALLINT UNSIGNED", + "INT UNSIGNED", + "FLOAT", + "REAL", + "FLOAT4", + "FLOAT8", + "FLOAT32", + "BOOL", + "TEXT", + "VARCHAR(10)", + "CHAR(10)", + "TIMESTAMP", + "DECIMAL(10, 2)", + "STRING(10)", + "BIGINT(10)", + ] { + for sql in [ + format!( + "CREATE TABLE traces (j JSON2 (value {sql_type}), ts TIMESTAMP TIME INDEX)" + ), + format!("ALTER TABLE traces MODIFY COLUMN j JSON2 (value {sql_type})"), + ] { + let err = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string().contains( + "supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64" + ), + "{sql}: {err}" + ); + } } } @@ -504,7 +532,7 @@ CREATE TABLE traces ( fn test_parse_json2_type_hint_rejects_default() { for default in ["-5", "abs(-1)"] { let sql = format!( - "CREATE TABLE traces (log_json_data JSON2 (status_code INT64 DEFAULT {default}), ts TIMESTAMP TIME INDEX)" + "CREATE TABLE traces (log_json_data JSON2 (status_code BIGINT DEFAULT {default}), ts TIMESTAMP TIME INDEX)" ); let err = ParserContext::create_with_dialect( &sql, @@ -521,7 +549,7 @@ CREATE TABLE traces ( let result = ParserContext::create_with_dialect( r#" CREATE TABLE traces ( - log_json_data JSON2 (a.b STRING, a.b INT64), + log_json_data JSON2 (a.b STRING, a.b BIGINT), ts TIMESTAMP TIME INDEX, )"#, &GreptimeDbDialect {}, @@ -537,7 +565,7 @@ CREATE TABLE traces ( let result = ParserContext::create_with_dialect( r#" CREATE TABLE traces ( - log_json_data JSON2 (a STRING, a.b INT64), + log_json_data JSON2 (a STRING, a.b BIGINT), ts TIMESTAMP TIME INDEX, )"#, &GreptimeDbDialect {}, diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 808e30f842c..e3f684bfcc2 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -225,7 +225,9 @@ pub fn sql_column_def_to_grpc_column_def( pub fn sql_data_type_to_concrete_data_type(data_type: &SqlDataType) -> Result { match data_type { SqlDataType::BigInt(_) | SqlDataType::Int64 => Ok(ConcreteDataType::int64_datatype()), - SqlDataType::BigIntUnsigned(_) => Ok(ConcreteDataType::uint64_datatype()), + SqlDataType::BigIntUnsigned(_) | SqlDataType::UInt64 => { + Ok(ConcreteDataType::uint64_datatype()) + } SqlDataType::Int(_) | SqlDataType::Integer(_) => Ok(ConcreteDataType::int32_datatype()), SqlDataType::IntUnsigned(_) | SqlDataType::UnsignedInteger => { Ok(ConcreteDataType::uint32_datatype()) @@ -422,6 +424,7 @@ mod tests { SqlDataType::BigIntUnsigned(None), ConcreteDataType::uint64_datatype(), ); + check_type(SqlDataType::UInt64, ConcreteDataType::uint64_datatype()); check_type( SqlDataType::IntUnsigned(None), ConcreteDataType::uint32_datatype(), diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 2310531841c..7c31790deaf 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -165,7 +165,7 @@ impl Json2Options { .map(|hint| { Ok(datatypes::json::JsonTypeHint { path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, + data_type: sql_data_type_to_concrete_data_type(&hint.data_type)?, inverted_index: hint.inverted_index, }) }) @@ -396,47 +396,19 @@ impl ColumnExtensions { } } -fn json_type_hint_concrete_data_type(data_type: &DataType) -> Result { - let data_type = sql_data_type_to_concrete_data_type(data_type)?; - normalize_json_type_hint_concrete_data_type(&data_type) -} - -fn normalize_json_type_hint_concrete_data_type( - data_type: &ConcreteDataType, -) -> Result { - let normalized = match data_type { - ConcreteDataType::String(_) => ConcreteDataType::string_datatype(), - ConcreteDataType::Int8(_) - | ConcreteDataType::Int16(_) - | ConcreteDataType::Int32(_) - | ConcreteDataType::Int64(_) => ConcreteDataType::int64_datatype(), - ConcreteDataType::UInt8(_) - | ConcreteDataType::UInt16(_) - | ConcreteDataType::UInt32(_) - | ConcreteDataType::UInt64(_) => ConcreteDataType::uint64_datatype(), - ConcreteDataType::Float32(_) | ConcreteDataType::Float64(_) => { - ConcreteDataType::float64_datatype() - } - ConcreteDataType::Boolean(_) => ConcreteDataType::boolean_datatype(), - _ => { - return InvalidSqlSnafu { - msg: format!("unsupported JSON2 type hint data type: {data_type}"), - } - .fail(); - } - }; - Ok(normalized) -} - fn json_type_hint_sql_data_type(data_type: &ConcreteDataType) -> Result { - let data_type = normalize_json_type_hint_concrete_data_type(data_type)?; let sql_type = match data_type { ConcreteDataType::String(_) => DataType::String(None), ConcreteDataType::Int64(_) => DataType::BigInt(None), ConcreteDataType::UInt64(_) => DataType::BigIntUnsigned(None), ConcreteDataType::Float64(_) => DataType::Double(sqlparser::ast::ExactNumberInfo::None), ConcreteDataType::Boolean(_) => DataType::Boolean, - _ => unreachable!("JSON2 type hint data type should have been normalized"), + _ => { + return InvalidSqlSnafu { + msg: format!("unsupported JSON2 type hint data type: {data_type}"), + } + .fail(); + } }; Ok(sql_type) } @@ -933,7 +905,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "service.name" STRING, - "a.b"."c" INT64, + "a.b"."c" BIGINT, a."b.c" STRING ), ts TIMESTAMP TIME INDEX @@ -977,7 +949,7 @@ ENGINE=mito fn test_parse_json2_max_auto_expanded_paths_option() -> Result<()> { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64, + status_code BIGINT, max_auto_expanded_paths = 1 ), ts TIMESTAMP TIME INDEX @@ -1003,7 +975,7 @@ ENGINE=mito let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( "1abc" STRING, - a."2b" INT64 + a."2b" BIGINT ), ts TIMESTAMP TIME INDEX )"#; @@ -1045,8 +1017,8 @@ ENGINE=mito fn test_json2_type_hint_rejects_default() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64 DEFAULT -5, - duration FLOAT64 DEFAULT +1.5, + status_code BIGINT DEFAULT -5, + duration DOUBLE DEFAULT +1.5, error BOOLEAN DEFAULT false, message STRING DEFAULT 'unknown' ), @@ -1062,7 +1034,7 @@ ENGINE=mito fn test_json2_type_hint_rejects_not_null() { let sql = r#"CREATE TABLE traces ( log_json_data JSON2 ( - status_code INT64 NOT NULL DEFAULT NULL + status_code BIGINT NOT NULL DEFAULT NULL ), ts TIMESTAMP TIME INDEX )"#; @@ -1073,23 +1045,23 @@ ENGINE=mito } #[test] - fn test_set_json_settings_normalizes_type_hint_sql_types() -> Result<()> { + fn test_set_json_settings_preserves_type_hint_sql_types() -> Result<()> { let mut extensions = super::ColumnExtensions::default(); let settings = JsonSettings::try_new( vec![ DatatypeJsonTypeHint { path: vec!["i".to_string()], - data_type: ConcreteDataType::int32_datatype(), + data_type: ConcreteDataType::int64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["f".to_string()], - data_type: ConcreteDataType::float32_datatype(), + data_type: ConcreteDataType::float64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { path: vec!["u".to_string()], - data_type: ConcreteDataType::uint32_datatype(), + data_type: ConcreteDataType::uint64_datatype(), inverted_index: false, }, DatatypeJsonTypeHint { diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index ab9e39dc777..153aa9e7902 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -18,7 +18,7 @@ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, service STRING, - duration INT64, + duration BIGINT, trace_id STRING ); diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql index 7fce02b76ee..457943e93f6 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -14,7 +14,7 @@ ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ( max_auto_expanded_paths = 2000, service STRING, - duration INT64, + duration BIGINT, trace_id STRING ); diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result index 284b7e6f850..aaefd5ff0ca 100644 --- a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.result @@ -24,7 +24,7 @@ ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); ALTER TABLE json2_alter_type_hints_compaction MODIFY COLUMN attrs JSON2 ( - kind INT64 + kind BIGINT ); Affected Rows: 0 diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql index 45d859d4ece..34a16dafbb7 100644 --- a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction.sql @@ -14,7 +14,7 @@ ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction'); ALTER TABLE json2_alter_type_hints_compaction MODIFY COLUMN attrs JSON2 ( - kind INT64 + kind BIGINT ); INSERT INTO json2_alter_type_hints_compaction VALUES diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result index 5ba2cf4700a..e3517c90e38 100644 --- a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.result @@ -24,7 +24,7 @@ ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); ALTER TABLE json2_alter_type_hints_compaction_conversion MODIFY COLUMN attrs JSON2 ( to_string STRING, - to_int INT64 + to_int BIGINT ); Affected Rows: 0 diff --git a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql index de9f211cfb5..cfadce89a93 100644 --- a/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql +++ b/tests/cases/standalone/common/types/json/json2_alter_type_hints_compaction_conversion.sql @@ -14,7 +14,7 @@ ADMIN FLUSH_TABLE('json2_alter_type_hints_compaction_conversion'); ALTER TABLE json2_alter_type_hints_compaction_conversion MODIFY COLUMN attrs JSON2 ( to_string STRING, - to_int INT64 + to_int BIGINT ); INSERT INTO json2_alter_type_hints_compaction_conversion VALUES diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result index 6e0f833ccd6..5ad7a46ce83 100644 --- a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.result @@ -24,7 +24,7 @@ ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); ALTER TABLE json2_compaction_null_type_hint_mismatch MODIFY COLUMN attrs JSON2 ( - kind INT64 + kind BIGINT ); Affected Rows: 0 diff --git a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql index 7ebff9cceb1..8025977c380 100644 --- a/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql +++ b/tests/cases/standalone/common/types/json/json2_compaction_null_type_hint_mismatch.sql @@ -14,7 +14,7 @@ ADMIN FLUSH_TABLE('json2_compaction_null_type_hint_mismatch'); ALTER TABLE json2_compaction_null_type_hint_mismatch MODIFY COLUMN attrs JSON2 ( - kind INT64 + kind BIGINT ); INSERT INTO json2_compaction_null_type_hint_mismatch VALUES diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.result b/tests/cases/standalone/common/types/json/json2_type_hints.result index 1a892b6e4cb..eda5aee7692 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.result +++ b/tests/cases/standalone/common/types/json/json2_type_hints.result @@ -71,7 +71,7 @@ CREATE TABLE json2_type_hints_timestamp ( ) ); -Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TIMESTAMP; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 -- A type hint at the maximum supported depth is accepted. CREATE TABLE json2_type_hint_depth_50 ( @@ -105,3 +105,132 @@ DROP TABLE json2_type_hint_depth_50; Affected Rows: 0 +-- Type hints reject unsupported types and aliases. +CREATE TABLE json2_hint_int8 (ts TIMESTAMP TIME INDEX, j JSON2 (value INT8)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: INT8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_tinyint (ts TIMESTAMP TIME INDEX, j JSON2 (value TINYINT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TINYINT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_uint8 (ts TIMESTAMP TIME INDEX, j JSON2 (value UINT8)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: UInt8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_float (ts TIMESTAMP TIME INDEX, j JSON2 (value FLOAT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: FLOAT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_text (ts TIMESTAMP TIME INDEX, j JSON2 (value TEXT)); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: TEXT; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +CREATE TABLE json2_hint_unsigned ( + ts TIMESTAMP TIME INDEX, + j JSON2 (value BIGINT UNSIGNED) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +INSERT INTO json2_hint_unsigned VALUES (1, '{"value":18446744073709551615}'); + +Affected Rows: 1 + +SELECT j.value FROM json2_hint_unsigned; + ++--------------------------------------------------------------+ +| json_get(json2_hint_unsigned.j,Utf8("$.value"),UInt64(NULL)) | ++--------------------------------------------------------------+ +| 18446744073709551615 | ++--------------------------------------------------------------+ + +ALTER TABLE json2_hint_unsigned MODIFY COLUMN j JSON2 (value INT8); + +Error: 2000(InvalidSyntax), Invalid SQL, error: unsupported JSON2 type hint data type: INT8; supported types: STRING, BIGINT, BIGINT UNSIGNED, DOUBLE, BOOLEAN; supported aliases: INT64, UINT64, FLOAT64 + +DROP TABLE json2_hint_unsigned; + +Affected Rows: 0 + +-- These aliases preserve the exact type and are displayed using canonical names. +CREATE TABLE json2_hint_aliases ( + ts TIMESTAMP TIME INDEX, + j JSON2 (i INT64, u UINT64, f FLOAT64) +) WITH ('append_mode' = 'true'); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_hint_aliases; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_hint_aliases | CREATE TABLE IF NOT EXISTS "json2_hint_aliases" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "i" BIGINT, | +| | "u" BIGINT UNSIGNED, | +| | "f" DOUBLE | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------+---------------------------------------------------+ + +INSERT INTO json2_hint_aliases VALUES + (1, '{"i":-9223372036854775808,"u":18446744073709551615,"f":1.5}'); + +Affected Rows: 1 + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + ++----------------------+----------------------+-----+ +| i | u | f | ++----------------------+----------------------+-----+ +| -9223372036854775808 | 18446744073709551615 | 1.5 | ++----------------------+----------------------+-----+ + +ALTER TABLE json2_hint_aliases MODIFY COLUMN j JSON2 (i INT64, u UINT64, f FLOAT64); + +Affected Rows: 0 + +SHOW CREATE TABLE json2_hint_aliases; + ++--------------------+---------------------------------------------------+ +| Table | Create Table | ++--------------------+---------------------------------------------------+ +| json2_hint_aliases | CREATE TABLE IF NOT EXISTS "json2_hint_aliases" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "i" BIGINT, | +| | "u" BIGINT UNSIGNED, | +| | "f" DOUBLE | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------+---------------------------------------------------+ + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + ++----------------------+----------------------+-----+ +| i | u | f | ++----------------------+----------------------+-----+ +| -9223372036854775808 | 18446744073709551615 | 1.5 | ++----------------------+----------------------+-----+ + +DROP TABLE json2_hint_aliases; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_type_hints.sql b/tests/cases/standalone/common/types/json/json2_type_hints.sql index aa88ae60bf0..3b301700333 100644 --- a/tests/cases/standalone/common/types/json/json2_type_hints.sql +++ b/tests/cases/standalone/common/types/json/json2_type_hints.sql @@ -59,3 +59,48 @@ CREATE TABLE json2_type_hint_depth_51 ( DROP TABLE json2_type_hints; DROP TABLE json2_type_hint_depth_50; + +-- Type hints reject unsupported types and aliases. +CREATE TABLE json2_hint_int8 (ts TIMESTAMP TIME INDEX, j JSON2 (value INT8)); + +CREATE TABLE json2_hint_tinyint (ts TIMESTAMP TIME INDEX, j JSON2 (value TINYINT)); + +CREATE TABLE json2_hint_uint8 (ts TIMESTAMP TIME INDEX, j JSON2 (value UINT8)); + +CREATE TABLE json2_hint_float (ts TIMESTAMP TIME INDEX, j JSON2 (value FLOAT)); + +CREATE TABLE json2_hint_text (ts TIMESTAMP TIME INDEX, j JSON2 (value TEXT)); + +CREATE TABLE json2_hint_unsigned ( + ts TIMESTAMP TIME INDEX, + j JSON2 (value BIGINT UNSIGNED) +) WITH ('append_mode' = 'true'); + +INSERT INTO json2_hint_unsigned VALUES (1, '{"value":18446744073709551615}'); + +SELECT j.value FROM json2_hint_unsigned; + +ALTER TABLE json2_hint_unsigned MODIFY COLUMN j JSON2 (value INT8); + +DROP TABLE json2_hint_unsigned; + +-- These aliases preserve the exact type and are displayed using canonical names. +CREATE TABLE json2_hint_aliases ( + ts TIMESTAMP TIME INDEX, + j JSON2 (i INT64, u UINT64, f FLOAT64) +) WITH ('append_mode' = 'true'); + +SHOW CREATE TABLE json2_hint_aliases; + +INSERT INTO json2_hint_aliases VALUES + (1, '{"i":-9223372036854775808,"u":18446744073709551615,"f":1.5}'); + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + +ALTER TABLE json2_hint_aliases MODIFY COLUMN j JSON2 (i INT64, u UINT64, f FLOAT64); + +SHOW CREATE TABLE json2_hint_aliases; + +SELECT j.i AS i, j.u AS u, j.f AS f FROM json2_hint_aliases; + +DROP TABLE json2_hint_aliases; From d910183cc46983074d8dfbb46166568ef819cb77 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 23 Sep 2026 09:21:40 +0000 Subject: [PATCH 020/133] feat(flow): freeze recovery windows and retention bounds for incremental flows (#9312) * feat(flow): discover recovery windows from exact sequence reads Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(flow): freeze explicit recovery windows and retention bounds Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): validate retained source windows before recovery Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): qualify recovery source tables before serialization Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/flow/Cargo.toml | 1 + src/flow/src/batching_mode/state.rs | 139 +++- src/flow/src/batching_mode/task.rs | 437 ++++++++++++- src/flow/src/batching_mode/task/test.rs | 820 +++++++++++++++++++++++- 5 files changed, 1376 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6c7916d8194..3ab47005352 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5468,6 +5468,7 @@ dependencies = [ "operator", "pretty_assertions", "prometheus 0.14.0", + "prost 0.14.1", "query", "rand 0.9.4", "serde", diff --git a/src/flow/Cargo.toml b/src/flow/Cargo.toml index d0ac5eee9ac..fea7b1eeda2 100644 --- a/src/flow/Cargo.toml +++ b/src/flow/Cargo.toml @@ -65,6 +65,7 @@ catalog = { workspace = true, features = ["testing"] } common-catalog.workspace = true common-meta = { workspace = true, features = ["testing"] } pretty_assertions.workspace = true +prost.workspace = true query.workspace = true session.workspace = true table.workspace = true diff --git a/src/flow/src/batching_mode/state.rs b/src/flow/src/batching_mode/state.rs index b9710a25a0d..55307dbcdec 100644 --- a/src/flow/src/batching_mode/state.rs +++ b/src/flow/src/batching_mode/state.rs @@ -208,6 +208,25 @@ impl TaskState { self.pending_fenced_repair.as_ref() } + /// Start repairing explicit bounded windows under a frozen high `H`. + /// + /// The caller supplies the complete repair scope. Live dirty windows are + /// left unchanged so signals received after `H` remain separate. + pub fn start_fenced_repair_windows( + &mut self, + high: BTreeMap, + windows: Vec<(Timestamp, Timestamp)>, + ) { + let mut pending_windows = self.dirty_time_windows.clone(); + pending_windows.clean(); + pending_windows.add_windows(windows); + self.pending_fenced_repair = Some(FencedRepair { + high, + pending_windows, + }); + self.checkpoint_mode = CheckpointMode::FullSnapshot; + } + /// Finish the fenced repair and promote the frozen high watermark to the /// checkpoint map. Incremental-disabled flows stay in FullSnapshot mode. pub fn finish_fenced_repair(&mut self) -> Option> { @@ -259,22 +278,18 @@ impl TaskState { task_ctx: Option<&BatchingTask>, ) -> Result, Error> { if let Some(repair) = self.pending_fenced_repair.as_mut() { - let expr = repair.pending_windows.gen_filter_exprs( + // Fenced windows are an explicit frozen repair scope. They must not + // be pruned by the moving live-data expiration boundary, and an + // empty repair must remain active until its high watermark is + // explicitly finished. + return repair.pending_windows.gen_filter_exprs( col_name, - expire_lower_bound, + None, window_size, window_cnt, flow_id, task_ctx, - )?; - if expr.is_some() || !repair.pending_windows.is_empty() { - return Ok(expr); - } - - // All pending repair windows may have expired during merge. Clear - // the empty repair so this call can fall back to live dirty windows - // instead of routing future executions to an empty queue forever. - self.pending_fenced_repair = None; + ); } self.dirty_time_windows.gen_filter_exprs( @@ -1393,6 +1408,108 @@ mod test { assert_eq!(state.dirty_time_windows.len(), 2); } + #[test] + fn test_explicit_fenced_repair_windows_keep_live_windows_separate() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows( + high.clone(), + vec![(Timestamp::new_second(10), Timestamp::new_second(15))], + ); + + assert_eq!(state.checkpoint_mode(), CheckpointMode::FullSnapshot); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 1); + + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(10), Timestamp::new_second(15))] + ); + state + .dirty_time_windows + .add_window(Timestamp::new_second(20), Some(Timestamp::new_second(25))); + state.restore_scoped_windows(&filter); + + assert_eq!( + state + .pending_fenced_repair() + .unwrap() + .pending_windows() + .len(), + 1 + ); + assert_eq!(state.dirty_time_windows.len(), 2); + } + + #[test] + fn test_explicit_fenced_repair_keeps_empty_high_and_old_windows() { + let mut state = state_with_past_update(Duration::from_secs(1)); + state + .dirty_time_windows + .add_window(Timestamp::new_second(0), Some(Timestamp::new_second(1_000))); + let high = BTreeMap::from([(1, 10)]); + state.start_fenced_repair_windows(high.clone(), Vec::new()); + + assert!( + state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .is_none() + ); + assert_eq!(state.pending_fenced_repair().unwrap().high(), &high); + assert_eq!(state.dirty_time_windows.len(), 1); + assert_eq!(state.finish_fenced_repair(), Some(high)); + + state.start_fenced_repair_windows( + BTreeMap::from([(1, 11)]), + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))], + ); + let filter = state + .gen_scoped_filter_exprs( + "ts", + Some(Timestamp::new_second(100)), + chrono::Duration::seconds(5), + 1, + 1, + None, + ) + .unwrap() + .unwrap(); + assert_eq!( + filter.time_ranges, + vec![(Timestamp::new_second(-20), Timestamp::new_second(-15))] + ); + } + #[test] fn test_disable_incremental_persists_full_snapshot_mode() { let query_ctx = QueryContext::arc(); diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 346d39d1d46..40020f639c1 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -18,21 +18,29 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use api::v1::{CreateTableExpr, TableName}; use catalog::CatalogManagerRef; +use catalog::kvbackend::KvBackendCatalogManager; use client::OutputWithMetrics; use common_error::ext::BoxedError; +use common_meta::key::schema_name::SchemaNameKey; +use common_query::OutputData; use common_query::logical_plan::breakup_insert_plan; use common_telemetry::tracing::warn; use common_telemetry::{debug, info}; -use common_time::Timestamp; +use common_time::{TimeToLive, Timestamp}; use datafusion::datasource::DefaultTableSource; use datafusion::sql::unparser::expr_to_sql; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::utils::quote_identifier; use datafusion_common::{DFSchemaRef, ScalarValue, TableReference}; -use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp, col, lit}; +use datafusion_expr::{DmlStatement, LogicalPlan, Projection, WriteOp, col, lit}; use datatypes::schema::Schema; +use datatypes::vectors::Helper; +use futures::TryStreamExt; use query::QueryEngineRef; -use query::options::FLOW_INCREMENTAL_MODE; +use query::options::{ + FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE, + FLOW_RETURN_REGION_SEQ, +}; use query::query_engine::DefaultSerializer; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; @@ -55,7 +63,7 @@ use crate::batching_mode::state::{ use crate::batching_mode::table_creator::{QueryType, create_table_with_expr}; use crate::batching_mode::time_window::TimeWindowExpr; use crate::batching_mode::utils::{ - AddFilterRewriter, ColumnMatcherRewriter, df_plan_to_sql, + AddFilterRewriter, ColumnMatcherRewriter, analyze_incremental_aggregate_plan, df_plan_to_sql, gen_plan_with_matching_schema_and_values, get_table_info_df_schema, sql_to_df_plan, }; use crate::df_optimizer::apply_df_optimizer; @@ -192,6 +200,56 @@ fn encode_insert_plan_request( }) } +fn recovery_aggregate_input(plan: &LogicalPlan) -> Result { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + _ => plan, + }; + let LogicalPlan::Aggregate(aggregate) = plan else { + return UnexpectedSnafu { + reason: "Recovery timestamp projection did not find an aggregate".to_string(), + } + .fail(); + }; + Ok(aggregate.input.as_ref().clone()) +} + +fn capture_recovery_batch_windows( + batch: &common_recordbatch::RecordBatch, + time_window_expr: &TimeWindowExpr, + windows: &mut BTreeSet<(Timestamp, Timestamp)>, +) -> Result<(), Error> { + if batch.num_columns() != 1 { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned {} columns instead of one", + batch.num_columns() + ), + } + .fail(); + } + let values = Helper::try_into_vector(batch.column(0).clone()) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + for index in 0..values.len() { + let timestamp = values.get(index).as_timestamp().context(UnexpectedSnafu { + reason: "Recovery timestamp projection returned a null or non-timestamp value" + .to_string(), + })?; + let (start, end) = time_window_expr.eval(timestamp)?; + let window = ( + start.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no lower bound".to_string(), + })?, + end.context(UnexpectedSnafu { + reason: "Recovery time-window expression returned no upper bound".to_string(), + })?, + ); + windows.insert(window); + } + Ok(()) +} + fn format_insert_target_columns(plan: &LogicalPlan) -> String { plan.schema() .fields() @@ -483,6 +541,377 @@ impl BatchingTask { Arc::new(query_ctx.fork()) } + /// Returns the retention lower bound aligned to this task's time window. + pub fn recovery_retention_lower_bound(&self) -> Result, Error> { + let Some(expire_after) = self.config.expire_after else { + return Ok(None); + }; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after}", + self.config.flow_id + ), + } + .build() + })?; + let now = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to read recovery wall clock: {err}"), + } + .build() + })? + .as_secs() as i64, + ); + let lower = now + .sub_duration(Duration::from_secs(expire_after)) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + self.config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: "Recovery expiry requires a time-window expression".to_string(), + })? + .eval(lower)? + .0 + .context(UnexpectedSnafu { + reason: "Recovery expiry time-window expression returned no lower bound" + .to_string(), + }) + .map(Some) + } + + /// Verifies that recovery can still read all source data required by its retained scope. + pub async fn validate_recovery_retention( + &self, + retention_lower: Option, + windows: &[(Timestamp, Timestamp)], + ) -> Result<(), Error> { + for name in &self.config.source_table_names { + let table = self + .config + .catalog_manager + .table(&name[0], &name[1], &name[2], None) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} is unavailable for recovery retention validation", + self.config.flow_id, + name.join(".") + ), + })?; + let ttl = if let Some(ttl) = table.table_info().meta.options.ttl { + ttl + } else { + let manager = self + .config + .catalog_manager + .as_any() + .downcast_ref::() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} cannot resolve inherited TTL for source table {} during recovery", + self.config.flow_id, + name.join(".") + ), + })?; + manager + .table_metadata_manager_ref() + .schema_manager() + .get(SchemaNameKey::new(&name[0], &name[1])) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + .context(UnexpectedSnafu { + reason: format!( + "Flow {} schema {}.{} is unavailable for recovery retention validation", + self.config.flow_id, name[0], name[1] + ), + })? + .ttl + .map(Into::into) + .unwrap_or(TimeToLive::Forever) + }; + + match ttl { + TimeToLive::Forever => {} + TimeToLive::Instant => { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has instant TTL and cannot be recovered", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + TimeToLive::Duration(ttl) => { + let expire_after = self.config.expire_after.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has TTL {ttl:?}, but recovery has no expire_after", + self.config.flow_id, + name.join(".") + ), + })?; + let expire_after = u64::try_from(expire_after).map_err(|_| { + UnexpectedSnafu { + reason: format!( + "Flow {} has negative expire_after {expire_after} during recovery retention validation", + self.config.flow_id + ), + } + .build() + })?; + if ttl <= Duration::from_secs(expire_after) { + return UnexpectedSnafu { + reason: format!( + "Flow {} source table {} TTL {ttl:?} must exceed expire_after {expire_after}s for recovery", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + let mut oldest = retention_lower.context(UnexpectedSnafu { + reason: format!( + "Flow {} source table {} has finite TTL {ttl:?}, but recovery has no retained lower bound", + self.config.flow_id, + name.join(".") + ), + })?; + for (start, _) in windows { + oldest = oldest.min(*start); + } + let cutoff = Timestamp::current_millis() + .sub_duration(ttl) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if oldest <= cutoff { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery requires source table {} data from {oldest:?}, older than TTL {ttl:?} cutoff {cutoff:?}", + self.config.flow_id, + name.join(".") + ), + } + .fail(); + } + } + } + } + Ok(()) + } + + /// Discovers the source time windows touched by the exact sequence range `(C, H]`. + /// + /// This compatibility wrapper captures the retention bound before discovery. + pub async fn capture_recovery_windows( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + let retention_lower = self.recovery_retention_lower_bound()?; + self.capture_recovery_windows_since(engine, frontend_client, lower, retention_lower) + .await + } + + /// Discovers windows under `lower` using a caller-frozen retention bound. + /// + /// The caller owns the execution guard and freezes this bound with its recovery scope. + /// This read-only method leaves task state untouched. Terminal proof must cover every + /// region in `lower`; subset or pruned proofs are rejected. + pub async fn capture_recovery_windows_since( + &self, + engine: &QueryEngineRef, + frontend_client: &FrontendClient, + lower: &BTreeMap, + retention_lower: Option, + ) -> Result<(BTreeMap, Vec<(Timestamp, Timestamp)>), Error> { + if lower.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires nonempty lower sequence bounds", + self.config.flow_id + ), + } + .fail(); + } + if !self.sequence_range_capable().await? { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires sequence_range-capable sources", + self.config.flow_id + ), + } + .fail(); + } + + let query_ctx = self.query_context_snapshot(); + let plan = sql_to_df_plan(query_ctx, engine.clone(), &self.config.query, false).await?; + let input = recovery_aggregate_input(&plan)?; + let Some(analysis) = analyze_incremental_aggregate_plan(&plan)? else { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a supported aggregate input", + self.config.flow_id + ), + } + .fail(); + }; + if !analysis.unsupported_exprs.is_empty() { + return UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture has unsupported aggregate expressions: {:?}", + self.config.flow_id, analysis.unsupported_exprs + ), + } + .fail(); + } + let time_window_expr = self + .config + .time_window_expr + .as_ref() + .context(UnexpectedSnafu { + reason: format!( + "Flow {} recovery window capture requires a time-window expression", + self.config.flow_id + ), + })?; + let input = if let Some(retention_lower) = retention_lower { + let mut add_filter = AddFilterRewriter::new( + col(&time_window_expr.column_name).gt_eq(lit(to_df_literal(retention_lower)?)), + ); + input + .rewrite(&mut add_filter) + .with_context(|_| DatafusionSnafu { + context: "Failed to apply recovery expire_after filter".to_string(), + })? + .data + } else { + input + }; + let timestamp_plan = LogicalPlan::Projection( + Projection::try_new(vec![col(&time_window_expr.column_name)], Arc::new(input)) + .context(DatafusionSnafu { + context: "Failed to project recovery source timestamps".to_string(), + })?, + ); + let catalog = &self.config.sink_table_name[0]; + let schema = &self.config.sink_table_name[1]; + let timestamp_plan = timestamp_plan + .clone() + .transform_down_with_subqueries(|p| { + if let LogicalPlan::TableScan(mut table_scan) = p { + let resolved = table_scan.table_name.resolve(catalog, schema); + table_scan.table_name = resolved.into(); + Ok(Transformed::yes(LogicalPlan::TableScan(table_scan))) + } else { + Ok(Transformed::no(p)) + } + }) + .with_context(|_| DatafusionSnafu { + context: format!( + "Failed to fix table ref in recovery timestamp plan, plan={:?}", + timestamp_plan + ), + })? + .data; + let message = DFLogicalSubstraitConvertor {} + .encode(×tamp_plan, DefaultSerializer) + .context(SubstraitEncodeLogicalPlanSnafu)?; + let request = api::v1::QueryRequest { + query: Some(api::v1::query_request::Query::LogicalPlan(message.to_vec())), + }; + let lower_json = serde_json::to_string(lower).map_err(|err| { + UnexpectedSnafu { + reason: format!("Failed to serialize recovery lower sequence bounds: {err}"), + } + .build() + })?; + let extensions = [ + (FLOW_RETURN_REGION_SEQ, "true"), + (FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_SEQUENCE_RANGE), + (FLOW_INCREMENTAL_AFTER_SEQS, lower_json.as_str()), + ]; + let mut peer_desc = None; + let result = frontend_client + .query_with_terminal_metrics( + catalog, + schema, + request, + &extensions, + &HashMap::new(), + &mut peer_desc, + ) + .await?; + let mut aligned_windows = BTreeSet::new(); + let metrics = result.metrics.clone(); + match result.output.data { + OutputData::AffectedRows(_) => { + return UnexpectedSnafu { + reason: "Recovery timestamp projection unexpectedly returned affected rows" + .to_string(), + } + .fail(); + } + OutputData::RecordBatches(batches) => { + for batch in batches.iter() { + capture_recovery_batch_windows(batch, time_window_expr, &mut aligned_windows)?; + } + } + OutputData::Stream(mut stream) => { + while let Some(batch) = stream + .try_next() + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)? + { + capture_recovery_batch_windows(&batch, time_window_expr, &mut aligned_windows)?; + } + } + } + if !metrics.is_ready() { + return UnexpectedSnafu { + reason: "Recovery timestamp projection ended without terminal metrics".to_string(), + } + .fail(); + } + let participating = metrics.participating_regions().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no participating-region proof".to_string(), + })?; + let high = metrics.region_watermark_map().context(UnexpectedSnafu { + reason: "Recovery timestamp projection has no terminal watermark proof".to_string(), + })?; + if participating.len() != lower.len() + || high.len() != lower.len() + || participating + .iter() + .any(|region| !lower.contains_key(region)) + || lower + .iter() + .any(|(region, low)| high.get(region).is_none_or(|watermark| watermark < low)) + || high.keys().any(|region| !participating.contains(region)) + { + return UnexpectedSnafu { + reason: format!( + "Recovery timestamp projection returned incomplete or regressing terminal proof: lower={lower:?}, participating={participating:?}, high={high:?}" + ), + } + .fail(); + } + Ok(( + high.into_iter().collect(), + aligned_windows.into_iter().collect(), + )) + } + /// Validates that the sink table schema can accept this flow's ordinary output. pub async fn validate_sink_table_schema(&self, engine: &QueryEngineRef) -> Result<(), Error> { self.validate_sink_table_schema_with_values(engine, &BTreeMap::new()) diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 7add59664a5..29a15600f15 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::task::Poll; +use std::task::{Context, Poll}; +use std::time::Duration; use catalog::memory::MemoryCatalogManager; use catalog::{DeregisterTableRequest, RegisterTableRequest}; @@ -27,18 +29,21 @@ use common_error::status_code::StatusCode; use common_query::Output; use common_recordbatch::RecordBatch; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; +use common_time::{TimeToLive, Timestamp}; use datatypes::data_type::ConcreteDataType as CDT; -use datatypes::schema::ColumnSchema; +use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{ TimestampMillisecondVector, TimestampNanosecondVector, UInt32Vector, VectorRef, }; use pretty_assertions::assert_eq; +use prost::Message; use query::options::{ FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY, FLOW_SCHEDULED_TIME_MILLIS, FLOW_SINK_TABLE_ID, QueryOptions, }; use session::context::QueryContext; use snafu::ResultExt; +use substrait::substrait_proto_df::proto::Plan; use table::Table; use table::metadata::FilterPushDownType; use table::test_util::MemTable; @@ -815,6 +820,47 @@ fn register_auto_created_aggregate_sink(query_engine: &QueryEngineRef, table_nam memory_catalog.register_table_sync(request).unwrap(); } +async fn configure_source_ttl(query_engine: &QueryEngineRef, ttl: Option) { + let catalog_manager = query_engine.engine_state().catalog_manager(); + let source = catalog_manager + .table( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "numbers_with_ts", + None, + ) + .await + .unwrap() + .unwrap(); + let mut info = (*source.table_info()).clone(); + info.meta.options.ttl = ttl; + let source = Arc::new(Table::new( + Arc::new(info), + FilterPushDownType::Unsupported, + source.data_source(), + )); + let memory_catalog = catalog_manager + .as_any() + .downcast_ref::() + .unwrap(); + memory_catalog + .deregister_table_sync(DeregisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + }) + .unwrap(); + memory_catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "numbers_with_ts".to_string(), + table_id: source.table_info().table_id(), + table: source, + }) + .unwrap(); +} + async fn configure_source_capability( query_engine: &QueryEngineRef, engine: &str, @@ -871,6 +917,127 @@ async fn configure_source_capability( .unwrap(); } +#[tokio::test] +async fn test_validate_recovery_retention_explicit_source_ttls() { + struct Case { + name: &'static str, + ttl: Option, + expire_after: Option, + retention_age: Option, + window_age: Option, + should_pass: bool, + } + + let cases = [ + Case { + name: "finite TTL exceeding expire_after is admissible", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: true, + }, + Case { + name: "TTL equal to expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "TTL shorter than expire_after is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(59 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "aligned window older than TTL is rejected", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: Some(Duration::from_secs(3 * 60 * 60)), + should_pass: false, + }, + Case { + name: "old retention cutoff is rejected even without windows", + ttl: Some(TimeToLive::Duration(Duration::from_secs(2 * 60 * 60))), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(3 * 60 * 60)), + window_age: None, + should_pass: false, + }, + Case { + name: "explicit forever is admissible", + ttl: Some(TimeToLive::Forever), + expire_after: None, + retention_age: None, + window_age: None, + should_pass: true, + }, + Case { + name: "instant TTL is rejected", + ttl: Some(TimeToLive::Instant), + expire_after: Some(60 * 60), + retention_age: Some(Duration::from_secs(10 * 60)), + window_age: None, + should_pass: false, + }, + ]; + + for case in cases { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query( + "SELECT number, ts FROM numbers_with_ts", + "sink", + ) + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = case.expire_after; + configure_source_ttl(&query_engine, case.ttl).await; + + let now = Timestamp::current_millis(); + let retention_lower = case.retention_age.map(|age| now.sub_duration(age).unwrap()); + let windows = case + .window_age + .map(|age| { + let start = now.sub_duration(age).unwrap(); + vec![(start, now)] + }) + .unwrap_or_default(); + assert_eq!( + task.validate_recovery_retention(retention_lower, &windows) + .await + .is_ok(), + case.should_pass, + "{}", + case.name + ); + } +} + +#[tokio::test] +async fn test_validate_recovery_retention_rejects_unknown_inherited_ttl() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_test_task_engine_and_plan_with_query("SELECT number, ts FROM numbers_with_ts", "sink") + .await; + Arc::get_mut(&mut task.config).unwrap().expire_after = Some(60 * 60); + configure_source_ttl(&query_engine, None).await; + + assert!( + task.validate_recovery_retention(Some(Timestamp::current_millis()), &[],) + .await + .is_err() + ); +} + fn dirty_marker() -> DirtyTimeWindows { let mut dirty = DirtyTimeWindows::default(); dirty.set_dirty(); @@ -1018,6 +1185,313 @@ impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError } } +struct RecoveryCaptureHandler { + output: std::sync::Mutex>, + lower: String, + expire_after: Option, + expiry_lower: std::sync::Mutex>, +} + +struct RecoveryMetricsStream { + schema: Arc, + batches: Vec, + metrics: Option, + fail: bool, +} + +impl futures::Stream for RecoveryMetricsStream { + type Item = common_recordbatch::error::Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + if self.fail { + self.fail = false; + return Poll::Ready(Some(Err(common_recordbatch::error::Error::External { + source: BoxedError::new(MockError::new(StatusCode::Unexpected)), + location: snafu::Location::new(file!(), line!(), column!()), + }))); + } + Poll::Ready(self.batches.pop().map(Ok)) + } +} + +impl common_recordbatch::RecordBatchStream for RecoveryMetricsStream { + fn name(&self) -> &str { + "RecoveryMetricsStream" + } + + fn schema(&self) -> Arc { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + self.batches + .is_empty() + .then(|| self.metrics.clone()) + .flatten() + } +} + +fn field_selection(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::field_reference::ReferenceType; + use substrait::substrait_proto_df::proto::expression::reference_segment::ReferenceType as Segment; + + let RexType::Selection(field) = expr.rex_type.as_ref()? else { + return None; + }; + let ReferenceType::DirectReference(reference) = field.reference_type.as_ref()? else { + return None; + }; + let Segment::StructField(field) = reference.reference_type.as_ref()? else { + return None; + }; + Some(field.field) +} + +fn timestamp_literal(expr: &substrait::substrait_proto_df::proto::Expression) -> Option { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + + let RexType::Literal(literal) = expr.rex_type.as_ref()? else { + return None; + }; + let LiteralType::PrecisionTimestamp(timestamp) = literal.literal_type.as_ref()? else { + return None; + }; + match timestamp.precision { + 0 => Some(Timestamp::new_second(timestamp.value)), + 3 => Some(Timestamp::new_millisecond(timestamp.value)), + 6 => Some(Timestamp::new_microsecond(timestamp.value)), + 9 => Some(Timestamp::new_nanosecond(timestamp.value)), + _ => None, + } +} + +fn assert_recovery_timestamp_plan(plan: &Plan, expire_after: Option) -> Option { + use substrait::substrait_proto_df::proto::plan_rel::RelType as PlanRelType; + use substrait::substrait_proto_df::proto::rel::RelType; + + let function_names = plan + .extensions + .iter() + .filter_map(|extension| match extension.mapping_type.as_ref()? { + substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::MappingType::ExtensionFunction(function) => { + Some((function.function_anchor, function.name.as_str())) + } + _ => None, + }) + .collect::>(); + let [root] = plan.relations.as_slice() else { + panic!("recovery capture must encode exactly one relation"); + }; + let Some(PlanRelType::Root(root)) = root.rel_type.as_ref() else { + panic!("recovery capture must encode a root relation"); + }; + assert_eq!(root.names, ["ts"]); + let Some(RelType::Project(project)) = root.input.as_ref().and_then(|rel| rel.rel_type.as_ref()) + else { + panic!("recovery capture must project the raw timestamp"); + }; + assert_eq!(project.expressions.len(), 1); + assert_eq!(field_selection(&project.expressions[0]), Some(1)); + + fn visit( + rel: &substrait::substrait_proto_df::proto::Rel, + function_names: &HashMap, + scans: &mut Vec>, + has_number_filter: &mut bool, + expiry_lower: &mut Option, + ) { + use substrait::substrait_proto_df::proto::expression::RexType; + use substrait::substrait_proto_df::proto::expression::literal::LiteralType; + use substrait::substrait_proto_df::proto::function_argument::ArgType; + use substrait::substrait_proto_df::proto::read_rel::ReadType; + use substrait::substrait_proto_df::proto::rel::RelType; + + match rel.rel_type.as_ref().expect("recovery relation type") { + RelType::Aggregate(_) | RelType::Write(_) | RelType::Ddl(_) | RelType::Update(_) => { + panic!("recovery capture must not aggregate or write") + } + RelType::Read(read) => match read.read_type.as_ref() { + Some(ReadType::NamedTable(table)) => scans.push(table.names.clone()), + _ => panic!("recovery capture must read a named source table"), + }, + RelType::Filter(filter) => { + let condition = filter + .condition + .as_ref() + .expect("recovery filter condition"); + let RexType::ScalarFunction(function) = condition + .rex_type + .as_ref() + .expect("recovery filter expression") + else { + panic!("recovery filter must be a scalar function"); + }; + let args = function + .arguments + .iter() + .map(|arg| match arg.arg_type.as_ref() { + Some(ArgType::Value(expr)) => expr, + _ => panic!("recovery filter must use value arguments"), + }) + .collect::>(); + match function_names.get(&function.function_reference) { + Some(&"equal") + if args.len() == 2 + && field_selection(args[0]) == Some(0) + && matches!( + args[1].rex_type.as_ref(), + Some(RexType::Literal(literal)) + if matches!(literal.literal_type, Some(LiteralType::I64(42))) + ) => + { + *has_number_filter = true + } + Some(&"gte") if args.len() == 2 && field_selection(args[0]) == Some(1) => { + *expiry_lower = timestamp_literal(args[1]); + } + _ => panic!("unexpected recovery filter"), + } + visit( + filter.input.as_ref().expect("recovery filter input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ); + } + RelType::Project(project) => visit( + project.input.as_ref().expect("recovery project input"), + function_names, + scans, + has_number_filter, + expiry_lower, + ), + _ => panic!("unexpected recovery relation"), + } + } + + let mut scans = Vec::new(); + let mut has_number_filter = false; + let mut expiry_lower = None; + visit( + project.input.as_ref().expect("recovery projection input"), + &function_names, + &mut scans, + &mut has_number_filter, + &mut expiry_lower, + ); + assert_eq!( + scans, + vec![vec![ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + "numbers_with_ts".to_string(), + ]] + ); + assert!(has_number_filter, "recovery capture lost WHERE number = 42"); + assert_eq!(expiry_lower.is_some(), expire_after.is_some()); + expiry_lower +} + +#[async_trait::async_trait] +impl crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError + for RecoveryCaptureHandler +{ + async fn do_query( + &self, + query: api::v1::greptime_request::Request, + ctx: QueryContextRef, + ) -> std::result::Result { + assert_eq!(ctx.extension(FLOW_INCREMENTAL_MODE), Some("sequence_range")); + assert_eq!( + ctx.extension(FLOW_INCREMENTAL_AFTER_SEQS), + Some(self.lower.as_str()) + ); + assert_eq!( + ctx.extension(query::options::FLOW_RETURN_REGION_SEQ), + Some("true") + ); + let api::v1::greptime_request::Request::Query(request) = query else { + panic!("recovery capture must issue a query request"); + }; + let Some(api::v1::query_request::Query::LogicalPlan(plan)) = request.query else { + panic!("recovery capture must issue a logical plan read"); + }; + let plan = Plan::decode(plan.as_slice()).unwrap(); + *self.expiry_lower.lock().unwrap() = + assert_recovery_timestamp_plan(&plan, self.expire_after); + Ok(self.output.lock().unwrap().take().unwrap()) + } +} + +fn recovery_timestamp_schema() -> Arc { + Arc::new(Schema::new(vec![ + ColumnSchema::new("ts", CDT::timestamp_millisecond_datatype(), true).with_time_index(true), + ])) +} + +fn recovery_stream_output( + batches: Vec, + watermarks: Option)>>, + fail: bool, +) -> Output { + Output::new_with_stream(Box::pin(RecoveryMetricsStream { + schema: recovery_timestamp_schema(), + batches, + metrics: watermarks.map(|watermarks| RecordBatchMetrics { + region_watermarks: watermarks + .into_iter() + .map(|(region_id, watermark)| RegionWatermarkEntry { + region_id, + watermark, + }) + .collect(), + ..Default::default() + }), + fail, + })) +} + +fn recovery_timestamp_batch(values: Vec>) -> RecordBatch { + RecordBatch::new( + recovery_timestamp_schema(), + vec![Arc::new(TimestampMillisecondVector::from(values)) as VectorRef], + ) + .unwrap() +} + +fn seed_recovery_state( + task: &BatchingTask, + lower: &BTreeMap, +) -> (BTreeMap, String) { + let mut state = task.state.write().unwrap(); + state.advance_checkpoints(lower.iter().map(|(region, seq)| (*region, *seq)).collect()); + state.dirty_time_windows.add_window( + Timestamp::new_millisecond(20_000), + Some(Timestamp::new_millisecond(25_000)), + ); + ( + state.checkpoints().clone(), + format!("{:?}", state.dirty_time_windows), + ) +} + +fn assert_recovery_state_unchanged( + task: &BatchingTask, + checkpoints: &BTreeMap, + dirty: &str, +) { + let state = task.state.read().unwrap(); + assert_eq!(state.checkpoints(), checkpoints); + assert_eq!(format!("{:?}", state.dirty_time_windows), dirty); +} + struct CaptureScheduledNowHandler { expected_extension: String, captured_sql: Arc>>, @@ -1370,6 +1844,331 @@ fn output_with_region_watermarks( result } +#[tokio::test] +async fn test_capture_recovery_windows_decodes_raw_timestamp_projection_with_where_and_expiry() { + let TestTaskParts { + mut task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(src.number) AS output_value, date_bin(INTERVAL '5 second', src.ts) AS output_window FROM numbers_with_ts AS src WHERE src.number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let expire_after = 5; + Arc::get_mut(&mut task.config) + .expect("test task config should be uniquely owned") + .expire_after = Some(expire_after); + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(expire_after), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &lower); + let before = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + let (high, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + let after = Timestamp::new_second( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + ); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + let expiry_lower = handler.expiry_lower.lock().unwrap().unwrap(); + let time_window_expr = task.config.time_window_expr.as_ref().unwrap(); + // The helper and samples may straddle a window boundary. + let aligned_expiry = |now: Timestamp| { + time_window_expr + .eval( + now.sub_duration(Duration::from_secs(expire_after as u64)) + .unwrap(), + ) + .unwrap() + .0 + .unwrap() + }; + assert!(aligned_expiry(before) <= expiry_lower && expiry_lower <= aligned_expiry(after)); + assert_eq!( + time_window_expr.eval(expiry_lower).unwrap().0, + Some(expiry_lower) + ); + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); +} + +#[tokio::test] +async fn test_capture_recovery_windows_since_uses_supplied_retention_lower() { + let TestTaskParts { + task, + query_engine, + .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let retention_lower = Timestamp::new_second(-10); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: Some(1), + expiry_lower: std::sync::Mutex::new(None), + }); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade( + &(handler.clone() + as Arc), + ), + QueryOptions::default(), + ); + + let (high, windows) = task + .capture_recovery_windows_since(&query_engine, &frontend, &lower, Some(retention_lower)) + .await + .unwrap(); + + assert_eq!(high, lower); + assert!(windows.is_empty()); + assert_eq!(*handler.expiry_lower.lock().unwrap(), Some(retention_lower)); +} + +#[tokio::test] +async fn test_capture_recovery_windows_terminal_proof_cases_leave_state_unchanged() { + struct Case { + name: &'static str, + lower: BTreeMap, + watermarks: Option)>>, + succeeds: bool, + } + + let cases = vec![ + Case { + name: "empty_stream_complete_proof_h_equals_c", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: true, + }, + Case { + name: "missing_metrics", + lower: BTreeMap::from([(1, 10)]), + watermarks: None, + succeeds: false, + }, + Case { + name: "missing_c_region", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10))]), + succeeds: false, + }, + Case { + name: "unknown_none_watermark", + lower: BTreeMap::from([(1, 10), (2, 20)]), + watermarks: Some(vec![(1, Some(10)), (2, None)]), + succeeds: false, + }, + Case { + name: "unexpected_region", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(10)), (2, Some(20))]), + succeeds: false, + }, + Case { + name: "regressing_h", + lower: BTreeMap::from([(1, 10)]), + watermarks: Some(vec![(1, Some(9))]), + succeeds: false, + }, + ]; + + for case in cases { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + case.watermarks, + false, + ))), + lower: serde_json::to_string(&case.lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = FrontendClient::from_grpc_handler( + Arc::downgrade(&handler_dyn), + QueryOptions::default(), + ); + let (before_checkpoints, before_dirty) = seed_recovery_state(&task, &case.lower); + + let result = task + .capture_recovery_windows(&query_engine, &frontend, &case.lower) + .await; + assert_eq!(result.is_ok(), case.succeeds, "{}: {result:?}", case.name); + if case.succeeds { + assert_eq!( + result.unwrap(), + (case.lower.clone(), vec![]), + "{}", + case.name + ); + } + assert_recovery_state_unchanged(&task, &before_checkpoints, &before_dirty); + } +} + +#[tokio::test] +async fn test_capture_recovery_windows_streams_sorted_unique_windows() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![ + recovery_timestamp_batch(vec![Some(10_000), Some(-1), Some(1_000)]), + recovery_timestamp_batch(vec![Some(11_000), Some(1_000), Some(-5_000)]), + ], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + let (_, windows) = task + .capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .unwrap(); + assert_eq!( + windows, + vec![ + ( + Timestamp::new_millisecond(-5_000), + Timestamp::new_millisecond(0) + ), + ( + Timestamp::new_millisecond(0), + Timestamp::new_millisecond(5_000) + ), + ( + Timestamp::new_millisecond(10_000), + Timestamp::new_millisecond(15_000) + ), + ] + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_stream_error() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![], + Some(vec![(1, Some(10))]), + true, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + +#[tokio::test] +async fn test_capture_recovery_windows_rejects_null_timestamp() { + let TestTaskParts { + task, query_engine, .. + } = new_time_window_test_task_with_query( + "SELECT max(number) AS output_value, date_bin(INTERVAL '5 second', ts) AS output_window FROM numbers_with_ts WHERE number = 42 GROUP BY output_window", + ) + .await; + configure_source_capability(&query_engine, "mito", true).await; + let lower = BTreeMap::from([(1, 10)]); + let handler = Arc::new(RecoveryCaptureHandler { + output: std::sync::Mutex::new(Some(recovery_stream_output( + vec![recovery_timestamp_batch(vec![None])], + Some(vec![(1, Some(10))]), + false, + ))), + lower: serde_json::to_string(&lower).unwrap(), + expire_after: None, + expiry_lower: std::sync::Mutex::new(None), + }); + let handler_dyn: Arc< + dyn crate::batching_mode::frontend_client::GrpcQueryHandlerWithBoxedError, + > = handler.clone(); + let frontend = + FrontendClient::from_grpc_handler(Arc::downgrade(&handler_dyn), QueryOptions::default()); + + assert!( + task.capture_recovery_windows(&query_engine, &frontend, &lower) + .await + .is_err() + ); +} + #[test] fn test_apply_query_result_to_state_advances_full_snapshot_to_incremental() { let query_ctx = QueryContext::arc(); @@ -2327,7 +3126,7 @@ async fn test_full_snapshot_scoped_plan_marks_checkpoint_advance_safe_only_after } #[tokio::test] -async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { +async fn test_expired_fenced_repair_uses_frozen_scope() { let TestTaskParts { mut task, query_engine, @@ -2365,11 +3164,18 @@ async fn test_expired_empty_fenced_repair_generates_scoped_base_repair_plan() { ) .await .unwrap() - .expect("expired empty repair should fall back to live dirty"); + .unwrap(); - assert!(matches!(plan.coverage, QueryCoverage::ScopedBaseRepair)); - assert!(plan.coverage.snapshot_seqs().is_empty()); - assert!(task.state.read().unwrap().pending_fenced_repair().is_none()); + // A fenced repair is a frozen explicit scope: expiry cannot discard its + // old window or make this query consume the live post-fence signal. + assert!(matches!( + plan.coverage, + QueryCoverage::FencedRepairChunk { .. } + )); + let state = task.state.read().unwrap(); + assert!(state.pending_fenced_repair().is_some()); + assert!(state.fenced_repair_pending_is_empty()); + assert_eq!(state.dirty_time_windows.len(), 1); } #[tokio::test] From 37fed9f12be0495a7221e72cdd72bbdd1bcb28dd Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 23 Sep 2026 10:32:19 +0000 Subject: [PATCH 021/133] fix(ci): repair draft PR command dispatch (#9271) * fix(ci): repair draft PR command dispatch Signed-off-by: WenyXu * fix(ci): dispatch command workflows by branch ref Signed-off-by: WenyXu * fix(ci): admit PR authors and writers for CI commands Signed-off-by: WenyXu * fix(ci): preserve dispatch guard dependency semantics Signed-off-by: WenyXu * fix(ci): configure slash command permissions individually Signed-off-by: WenyXu * refactor(ci): run command workflows at dispatch branch head Signed-off-by: WenyXu * fix(ci): rerun fork PR checks and report command failures Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/scripts/ci-slash.py | 43 ++++-- .github/scripts/tests/test_ci_slash.py | 136 ++++++++++++++++++- .github/workflows/cargo-lock-check.yml | 1 + .github/workflows/check-git-deps.yml | 1 + .github/workflows/checks.yml | 10 +- .github/workflows/ci-slash.yml | 41 ++++-- .github/workflows/dependency-check.yml | 1 + .github/workflows/docs.yml | 6 +- .github/workflows/grafana.yml | 1 + .github/workflows/integration.yml | 10 +- .github/workflows/rust.yml | 14 +- .github/workflows/slash-command-dispatch.yml | 25 ++-- 12 files changed, 239 insertions(+), 50 deletions(-) diff --git a/.github/scripts/ci-slash.py b/.github/scripts/ci-slash.py index 61efa97fdf4..528df3feee7 100644 --- a/.github/scripts/ci-slash.py +++ b/.github/scripts/ci-slash.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Admit `/ci` comments and dispatch a pinned draft-PR CI workflow.""" +"""Admit `/ci` comments and dispatch a draft-PR CI workflow.""" import json, os, re, sys, urllib.request COMMAND = re.compile(r'^/ci(?:\s+(.+))?$') @@ -29,7 +29,7 @@ OPTIONS = { 'fuzz chaos': ('integration.yml', 'chaos', 'chaos fuzz'), 'fuzz all': ('integration.yml', 'all', 'all fuzz'), } -HELP = '''Available draft-PR CI commands:\n\n- `/ci` or `/ci rust` — Rust CI\n- `/ci integration` — integration CI without fuzz\n- `/ci checks` or `/ci docs`\n- `/ci fuzz standalone|distributed|chaos`\n- `/ci fuzz all` — all fuzz suites (admin only)\n\nCommands require repository admin permission and a same-repository open draft PR. CI is pinned to the current head SHA; comment again after a push.''' +HELP = '''Available draft-PR CI commands:\n\n- `/ci` — standard CI\n- `/ci rust` — Rust CI\n- `/ci integration` — integration CI without fuzz\n- `/ci checks` or `/ci docs`\n- `/ci fuzz standalone|distributed|chaos`\n- `/ci fuzz all` — all fuzz suites (admin only)\n\nCommands require the PR author or repository write/maintain/admin permission and an open draft PR. Fork PRs require repository write permission and rerun existing pull-request CI; fuzz commands are same-repository only. `/ci fuzz all` requires admin permission. Same-repository CI runs the branch head at dispatch time. Fork CI reruns the original pull-request revision (within GitHub rerun limits); sync your branch with the updated CI configuration first.''' def api(path): req=urllib.request.Request(os.environ['GITHUB_API_URL']+path, headers={'Authorization':'Bearer '+os.environ['GITHUB_TOKEN'],'Accept':'application/vnd.github+json'}) @@ -46,7 +46,7 @@ def reject(number, text): out(skip='true',pr_number=number,reply='CI command ign def main(): if os.environ.get('DISPATCH_SENDER') != 'github-actions[bot]': return reject('', 'invalid dispatch sender.') comment=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/issues/comments/'+os.environ['COMMENT_ID']) - body=comment.get('body','').splitlines()[0].strip(); m=COMMAND.fullmatch(body) + body=(comment.get('body') or '').partition('\n')[0].strip(); m=COMMAND.fullmatch(body) number=str(comment.get('issue_url','').rstrip('/').split('/')[-1]) if not m: return reject(number,'comment is not a `/ci` command.') arg=(m.group(1) or '').strip().lower() @@ -54,13 +54,38 @@ def main(): if arg=='help': out(skip='true',pr_number=number,reply=HELP); return 0 pr=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/pulls/'+number) if pr.get('state')!='open' or not pr.get('draft'): return reject(number,'PR must be open and draft.') - if pr.get('head',{}).get('repo',{}).get('full_name') != os.environ['GITHUB_REPOSITORY']: return reject(number,'fork PRs are not admitted.') - head=pr.get('head',{}).get('sha','') - if head != os.environ.get('DISPATCH_HEAD_SHA') or not re.fullmatch('[0-9a-f]{40}',head): return reject(number,'PR head changed; comment again.') + fork=pr.get('head',{}).get('repo',{}).get('full_name') != os.environ['GITHUB_REPOSITORY'] + if fork and arg.startswith('fuzz '): return reject(number,'fuzz commands require a same-repository PR.') + head=pr.get('head',{}) + head_sha=head.get('sha','') + head_ref=head.get('ref','') + if head_sha != os.environ.get('DISPATCH_HEAD_SHA') or not re.fullmatch('[0-9a-f]{40}',head_sha) or not head_ref: return reject(number,'PR head changed; comment again.') actor=comment.get('user',{}).get('login','') - permission=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/collaborators/'+actor+'/permission').get('user',{}).get('permission') - if permission!='admin': return reject(number,'repository admin permission is required.') + if not actor: return reject(number,'comment author is missing.') + is_author=actor==pr.get('user',{}).get('login') + if fork or arg=='fuzz all' or not is_author: + permission=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/collaborators/'+actor+'/permission').get('permission') + if arg=='fuzz all': + if permission!='admin': return reject(number,'repository admin permission is required for `/ci fuzz all`.') + elif permission not in ('write','maintain','admin'): + return reject(number,'PR author or repository write permission is required.') workflow, profile, label=OPTIONS[arg] - out(skip='false',pr_number=number,head_sha=head,workflow=workflow,fuzz_profile=profile,reply=f'Dispatched {label} for `{head}`.') + if fork: + runs=api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/actions/runs?event=pull_request&head_sha='+head_sha+'&per_page=100')['workflow_runs'] + selected={} + for run in runs: + name=run.get('path','').split('@',1)[0].removeprefix('.github/workflows/') + if name not in workflow.split(',') or name in selected: continue + if run.get('head_sha') != head_sha or run.get('head_repository',{}).get('full_name') != head['repo']['full_name'] or run.get('head_branch') != head_ref: continue + selected[name]=run + if not selected: return reject(number,'no pull-request CI runs found for this head; push a new commit first.') + if arg and workflow not in selected: return reject(number,'selected CI has no pull-request run for this head.') + if any(run['status'] != 'completed' for run in selected.values()): return reject(number,'pull-request CI is still pending; approve or wait for it, then comment again.') + for run in selected.values(): + original=run if run.get('run_attempt',1)==1 else api('/repos/'+os.environ['GITHUB_REPOSITORY']+'/actions/runs/'+str(run['id'])+'/attempts/1') + if original.get('conclusion') != 'skipped': return reject(number,'only originally skipped draft CI runs can be requested; push a new commit while draft.') + out(skip='false',pr_number=number,run_ids=','.join(str(run['id']) for run in selected.values()),reply=f'Requested rerun of {", ".join(selected)} for `{head_sha}`.') + return 0 + out(skip='false',pr_number=number,head_sha=head_sha,head_ref=head_ref,workflow=workflow,fuzz_profile=profile,reply=f'Dispatched {label} for branch `{head_ref}`.') return 0 if __name__ == '__main__': sys.exit(main()) diff --git a/.github/scripts/tests/test_ci_slash.py b/.github/scripts/tests/test_ci_slash.py index afcf8031cdb..c8c01ceda15 100644 --- a/.github/scripts/tests/test_ci_slash.py +++ b/.github/scripts/tests/test_ci_slash.py @@ -14,8 +14,11 @@ # limitations under the License. import importlib.util +import json import os import tempfile +import subprocess +import textwrap import unittest from pathlib import Path from unittest.mock import patch @@ -27,6 +30,9 @@ ci = importlib.util.module_from_spec(spec) spec.loader.exec_module(ci) +OPTIONS_TO_TEST = ["/ci", "/ci rust", "/ci fuzz chaos", "/ci fuzz all"] + + class CiSlashTest(unittest.TestCase): def setUp(self): self.output = tempfile.NamedTemporaryFile(delete=False) @@ -54,23 +60,145 @@ class CiSlashTest(unittest.TestCase): self.assertIn("skip=true", output) self.assertIn("Available draft-PR CI commands", output) - def test_dispatches_full_suite_at_verified_head(self): + def test_dispatches_full_suite_with_top_level_admin_permission(self): output = self.run_main([ {"body": "/ci", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, - {"state": "open", "draft": True, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, - {"user": {"permission": "admin"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + {"permission": "admin", "user": {"login": "admin"}}, ]) self.assertIn("skip=false", output) + self.assertIn("head_ref=fix/draft-ci", output) self.assertIn("workflow=rust.yml,integration.yml,checks.yml,docs.yml", output) def test_rejects_non_draft_before_permission_lookup(self): output = self.run_main([ {"body": "/ci fuzz chaos", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, - {"state": "open", "draft": False, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + {"state": "open", "draft": False, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, ]) self.assertIn("skip=true", output) self.assertIn("PR must be open and draft", output) + def test_rejects_missing_head_ref_before_permission_lookup(self): + output = self.run_main([ + {"body": "/ci", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "admin"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + ]) + self.assertIn("skip=true", output) + self.assertIn("PR head changed; comment again.", output) + + def test_permission_matrix(self): + for command in OPTIONS_TO_TEST: + for author in (True, False): + for permission in ("admin", "maintain", "write", "triage", "read", "none", None): + with self.subTest(command=command, author=author, permission=permission): + Path(self.output.name).write_text("") + responses = [ + {"body": command, "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "actor"}}, + {"state": "open", "draft": True, "user": {"login": "actor" if author else "other"}, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, + ] + if command == "/ci fuzz all" or not author: + responses.append({"permission": permission}) + output = self.run_main(responses) + allowed = permission == "admin" if command == "/ci fuzz all" else author or permission in ("admin", "maintain", "write") + self.assertIn("skip=false" if allowed else "skip=true", output) + + def test_reply_uses_issue_comment_endpoint(self): + workflow = SCRIPT.parents[1] / "workflows" / "ci-slash.yml" + step = workflow.read_text().split(" - name: Reply with command result\n", 1)[1] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + reply = "Denied: `example`\nSecond line with $variables and 'quotes'" + with tempfile.TemporaryDirectory() as directory: + gh = Path(directory) / "gh" + gh.write_text("#!/usr/bin/env python3\nimport json, sys\nprint(json.dumps(sys.argv[1:]))\n") + gh.chmod(0o755) + env = {**os.environ, "PATH": directory + os.pathsep + os.environ["PATH"], "GITHUB_REPOSITORY": "GreptimeTeam/greptimedb", "PR_NUMBER": "42", "REPLY": reply, "ADMIT_OUTCOME": "success", "DISPATCH_OUTCOME": "success"} + result = subprocess.run(["bash", "-eu", "-c", command], env=env, check=True, capture_output=True, text=True) + self.assertEqual(json.loads(result.stdout), ["api", "--method", "POST", "/repos/GreptimeTeam/greptimedb/issues/42/comments", "-f", "body=" + reply]) + + def test_failure_reply_and_partial_dispatch(self): + workflow = (SCRIPT.parents[1] / "workflows" / "ci-slash.yml").read_text() + dispatch = textwrap.dedent(workflow.split(" - name: Dispatch selected CI workflow\n", 1)[1].split(" run: |\n", 1)[1].split(" - name: Reply", 1)[0]) + reply = textwrap.dedent(workflow.split(" - name: Reply with command result\n", 1)[1].split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + gh = Path(directory) / "gh" + log = Path(directory) / "calls" + gh.write_text("#!/usr/bin/env python3\nimport json, os, sys\nwith open(os.environ['CALLS'], 'a') as f: f.write(json.dumps(sys.argv[1:]) + '\\n')\nsys.exit(1 if '/runs/8/rerun' in sys.argv[-1] else 0)\n") + gh.chmod(0o755) + env = {**os.environ, "PATH": directory + os.pathsep + os.environ["PATH"], "CALLS": str(log), "GITHUB_REPOSITORY": "GreptimeTeam/greptimedb", "RUN_IDS": "7,8,9", "PR_NUMBER": "42", "REPLY": "Success", "ADMIT_OUTCOME": "success", "DISPATCH_OUTCOME": "failure", "GITHUB_SERVER_URL": "https://github.com", "GITHUB_RUN_ID": "123"} + result = subprocess.run(["bash", "-eu", "-c", dispatch], env=env, capture_output=True) + self.assertNotEqual(result.returncode, 0) + subprocess.run(["bash", "-eu", "-c", reply], env=env, check=True, capture_output=True) + calls = [json.loads(line) for line in log.read_text().splitlines()] + self.assertEqual([call[-1] for call in calls[:2]], ["/repos/GreptimeTeam/greptimedb/actions/runs/7/rerun", "/repos/GreptimeTeam/greptimedb/actions/runs/8/rerun"]) + self.assertIn("CI trigger failed; some workflows may already have been requested", calls[2][-1]) + self.assertIn("/actions/runs/123", calls[2][-1]) + self.assertIn("!cancelled()", workflow.split(" - name: Reply with command result", 1)[1]) + + def test_fork_rerun_gates_cover_standard_jobs(self): + import re + for name, count in [("rust", 7), ("integration", 5), ("checks", 5), ("docs", 3)]: + workflow = (SCRIPT.parents[1] / "workflows" / (name + ".yml")).read_text() + gates = re.findall(r"^ if:.*github.run_attempt > 1.*$", workflow, re.MULTILINE) + self.assertEqual(len(gates), count, name) + self.assertTrue(all("always()" not in gate for gate in gates)) + + def test_command_permissions(self): + workflow = SCRIPT.parents[1] / "workflows" / "slash-command-dispatch.yml" + config = workflow.read_text().split("config: >-", 1)[1].split(" - name:", 1)[0] + self.assertEqual(json.loads(config), [ + {"command": "query-regression", "permission": "admin", "issue_type": "pull-request"}, + {"command": "ci", "permission": "none", "issue_type": "pull-request"}, + ]) + + def test_fork_requires_writer_even_for_author(self): + for permission in ("read", "write", "maintain", "admin"): + with self.subTest(permission=permission): + Path(self.output.name).write_text("") + responses = [ + {"body": "/ci", "issue_url": "https://api.github.test/issues/42", "user": {"login": "actor"}}, + {"state": "open", "draft": True, "user": {"login": "actor"}, "head": {"sha": "a" * 40, "ref": "fork-branch", "repo": {"full_name": "actor/greptimedb"}}}, + {"permission": permission}, + ] + if permission != "read": + responses.append({"workflow_runs": [ + {"id": 7, "path": ".github/workflows/rust.yml", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "actor/greptimedb"}, "status": "completed", "conclusion": "skipped"}, + {"id": 8, "path": ".github/workflows/checks.yml", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "someone/greptimedb"}, "status": "completed", "conclusion": "skipped"}, + ]}) + output = self.run_main(responses) + if permission == "read": + self.assertIn("skip=true", output) + else: + self.assertIn("run_ids=7\n", output) + self.assertNotIn("head_ref=", output) + + def test_fork_rejects_non_draft_original_run(self): + output = self.run_main([ + {"body": "/ci rust", "issue_url": "https://api.github.test/issues/42", "user": {"login": "writer"}}, + {"state": "open", "draft": True, "head": {"sha": "a" * 40, "ref": "fork-branch", "repo": {"full_name": "actor/greptimedb"}}}, + {"permission": "write"}, + {"workflow_runs": [{"id": 7, "path": ".github/workflows/rust.yml@main", "head_sha": "a" * 40, "head_branch": "fork-branch", "head_repository": {"full_name": "actor/greptimedb"}, "status": "completed", "conclusion": "success", "run_attempt": 2}]}, + {"conclusion": "success"}, + ]) + self.assertIn("only originally skipped draft CI runs", output) + self.assertIn("skip=true", output) + + def test_author_cannot_bypass_pr_guards(self): + for change, reason in [ + ({"state": "closed"}, "PR must be open and draft"), + ({"draft": False}, "PR must be open and draft"), + ({"head": {"sha": "b" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}}, "PR head changed"), + ]: + with self.subTest(change=change): + Path(self.output.name).write_text("") + pr = {"state": "open", "draft": True, "user": {"login": "actor"}, "head": {"sha": "a" * 40, "ref": "fix/draft-ci", "repo": {"full_name": "GreptimeTeam/greptimedb"}}} + pr.update(change) + output = self.run_main([ + {"body": "/ci rust", "issue_url": "https://api.github.test/repos/GreptimeTeam/greptimedb/issues/42", "user": {"login": "actor"}}, + pr, + ]) + self.assertIn("skip=true", output) + self.assertIn(reason, output) + if __name__ == "__main__": unittest.main() diff --git a/.github/workflows/cargo-lock-check.yml b/.github/workflows/cargo-lock-check.yml index 847be806c5f..cfea53c6750 100644 --- a/.github/workflows/cargo-lock-check.yml +++ b/.github/workflows/cargo-lock-check.yml @@ -5,6 +5,7 @@ name: Cargo.lock Diff Check # posts/updates a PR comment when the threshold is exceeded. on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] paths: - "Cargo.lock" diff --git a/.github/workflows/check-git-deps.yml b/.github/workflows/check-git-deps.yml index 6ebf79f2072..a4ec488e850 100644 --- a/.github/workflows/check-git-deps.yml +++ b/.github/workflows/check-git-deps.yml @@ -2,6 +2,7 @@ name: Check Git Dependencies on Main Branch on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: [main] paths: - 'Cargo.toml' diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 6e591ba9dfd..561ee2c7473 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -32,7 +32,7 @@ concurrency: jobs: check-typos-and-docs: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check typos and docs runs-on: ubuntu-latest steps: @@ -47,7 +47,7 @@ jobs: || (echo "'config/config.md' is not up-to-date, please run 'make config-docs'." && exit 1) license-header-check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} runs-on: ubuntu-latest name: Check License Header steps: @@ -63,7 +63,7 @@ jobs: config: licenserc-enterprise.toml github-script-tests: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: GitHub Script Tests runs-on: ubuntu-latest timeout-minutes: 5 @@ -100,7 +100,7 @@ jobs: python3 tests/perf/test_agent_observability_summary.py check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check runs-on: ${{ matrix.os }} strategy: @@ -129,7 +129,7 @@ jobs: run: cargo check --locked --workspace --all-targets --all-features toml: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Toml Check runs-on: ubuntu-latest timeout-minutes: 60 diff --git a/.github/workflows/ci-slash.yml b/.github/workflows/ci-slash.yml index 1af8e60d95c..e9fb392d381 100644 --- a/.github/workflows/ci-slash.yml +++ b/.github/workflows/ci-slash.yml @@ -30,23 +30,46 @@ jobs: DISPATCH_SENDER: ${{ github.event.sender.login }} DISPATCH_HEAD_SHA: ${{ github.event.client_payload.pull_request.head.sha }} run: python3 .github/scripts/ci-slash.py - - name: Reply with command result - if: ${{ steps.admit.outputs.reply != '' }} - env: - GH_TOKEN: ${{ github.token }} - REPLY: ${{ steps.admit.outputs.reply }} - PR_NUMBER: ${{ steps.admit.outputs.pr_number }} - run: printf '%s\n' "${REPLY}" | gh pr comment "${PR_NUMBER}" --body-file - - name: Dispatch selected CI workflow + id: dispatch if: ${{ steps.admit.outputs.skip == 'false' }} env: GH_TOKEN: ${{ github.token }} + RUN_IDS: ${{ steps.admit.outputs.run_ids }} WORKFLOWS: ${{ steps.admit.outputs.workflow }} - HEAD_SHA: ${{ steps.admit.outputs.head_sha }} + HEAD_REF: ${{ steps.admit.outputs.head_ref }} FUZZ_PROFILE: ${{ steps.admit.outputs.fuzz_profile }} run: | + if [[ -n "${RUN_IDS}" ]]; then + IFS=, read -ra runs <<<"${RUN_IDS}" + for run in "${runs[@]}"; do + gh api --method POST "/repos/${GITHUB_REPOSITORY}/actions/runs/${run}/rerun" + done + exit 0 + fi IFS=, read -ra workflows <<<"${WORKFLOWS}" for workflow in "${workflows[@]}"; do + inputs=(-F 'inputs[ci_command]=true') + if [[ "${workflow}" == 'integration.yml' ]]; then + inputs+=(-f "inputs[fuzz_profile]=${FUZZ_PROFILE}") + fi gh api --method POST "/repos/${GITHUB_REPOSITORY}/actions/workflows/${workflow}/dispatches" \ - -f ref="${HEAD_SHA}" -F 'inputs[ci_command]=true' -f "inputs[fuzz_profile]=${FUZZ_PROFILE}" + -f ref="${HEAD_REF}" "${inputs[@]}" done + - name: Reply with command result + if: ${{ !cancelled() && github.event.sender.login == 'github-actions[bot]' }} + env: + GH_TOKEN: ${{ github.token }} + REPLY: ${{ steps.admit.outputs.reply }} + PR_NUMBER: ${{ steps.admit.outputs.pr_number || github.event.client_payload.github.payload.issue.number }} + ADMIT_OUTCOME: ${{ steps.admit.outcome }} + DISPATCH_OUTCOME: ${{ steps.dispatch.outcome }} + # Use the issue-comment endpoint with the existing issues: write permission. + run: | + [[ "${PR_NUMBER}" =~ ^[0-9]+$ ]] || exit 1 + if [[ "${ADMIT_OUTCOME}" == failure || "${DISPATCH_OUTCOME}" == failure ]]; then + REPLY="CI trigger failed; some workflows may already have been requested. See ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for details." + fi + [[ -n "${REPLY}" ]] || exit 1 + gh api --method POST "/repos/${GITHUB_REPOSITORY}/issues/${PR_NUMBER}/comments" \ + -f "body=${REPLY}" diff --git a/.github/workflows/dependency-check.yml b/.github/workflows/dependency-check.yml index 93a366c4b08..014a41e45c6 100644 --- a/.github/workflows/dependency-check.yml +++ b/.github/workflows/dependency-check.yml @@ -2,6 +2,7 @@ name: Check Dependencies on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: - main diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3006dc7f954..c9a6e8caaf3 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,7 +38,7 @@ name: Docs CI jobs: typos: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} name: Spell Check with Typos runs-on: ubuntu-latest steps: @@ -48,7 +48,7 @@ jobs: - uses: crate-ci/typos@v1.50.2 license-header-check: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} runs-on: ubuntu-latest name: Check License Header steps: @@ -58,7 +58,7 @@ jobs: - uses: korandoru/hawkeye@v5 required-checks: - if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false }} + if: ${{ github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1 }} name: ${{ matrix.check }} runs-on: ubuntu-slim strategy: diff --git a/.github/workflows/grafana.yml b/.github/workflows/grafana.yml index 6f3f1aa8fdf..2ec57b3cf5d 100644 --- a/.github/workflows/grafana.yml +++ b/.github/workflows/grafana.yml @@ -2,6 +2,7 @@ name: Check Grafana Panels on: pull_request: + types: [opened, synchronize, reopened, ready_for_review] branches: - main paths: diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 4e0634fbc57..e33abac3330 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -43,7 +43,7 @@ concurrency: jobs: build: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Build GreptimeDB binaries runs-on: ${{ matrix.os }} strategy: @@ -104,7 +104,7 @@ jobs: version: current sqlness: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Sqlness Test (${{ matrix.mode.name }}) needs: build runs-on: ${{ matrix.os }} @@ -184,7 +184,7 @@ jobs: retention-days: 3 export-import-v2-e2e: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Export/Import V2 E2E Test needs: build runs-on: ubuntu-latest @@ -254,7 +254,7 @@ jobs: retention-days: 3 run-multi-lang-tests: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Run Multi-language SDK Tests needs: build uses: ./.github/workflows/run-multi-lang-tests.yml @@ -262,7 +262,7 @@ jobs: artifact-name: bins compat-updater-check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) && (github.event_name == 'merge_group' || github.event_name == 'pull_request') }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) && (github.event_name == 'merge_group' || github.event_name == 'pull_request') }} name: Check compat version updater runs-on: ubuntu-latest timeout-minutes: 10 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 5778d3856bb..7302ce310b1 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -32,7 +32,7 @@ concurrency: jobs: fmt: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Rustfmt runs-on: ubuntu-latest timeout-minutes: 60 @@ -50,7 +50,7 @@ jobs: run: make fmt-check clippy: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Clippy runs-on: ubuntu-latest timeout-minutes: 60 @@ -76,7 +76,7 @@ jobs: run: make clippy check-udeps: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check Unused Dependencies runs-on: ubuntu-latest timeout-minutes: 60 @@ -94,7 +94,7 @@ jobs: run: make check-udeps check-riscv64: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check riscv64 build runs-on: ubuntu-latest timeout-minutes: 90 @@ -120,7 +120,7 @@ jobs: run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu check-windows: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check Windows build runs-on: windows-2022 timeout-minutes: 60 @@ -191,7 +191,7 @@ jobs: run: cargo nextest run --locked --workspace -F dashboard --no-fail-fast conflict-check: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} name: Check for conflict runs-on: ubuntu-latest steps: @@ -202,7 +202,7 @@ jobs: uses: olivernybroe/action-conflict-finder@v4.0 test: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false) }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && github.event_name != 'merge_group' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} runs-on: ubuntu-22.04-arm timeout-minutes: 60 needs: [conflict-check, clippy, fmt, check-udeps] diff --git a/.github/workflows/slash-command-dispatch.yml b/.github/workflows/slash-command-dispatch.yml index 26fa27b79dd..c47a9ed9ced 100644 --- a/.github/workflows/slash-command-dispatch.yml +++ b/.github/workflows/slash-command-dispatch.yml @@ -1,13 +1,13 @@ name: Slash Command Dispatch -# ChatOps front door: parse `/command` on PR comments, check admin -# permission, and repository_dispatch to a per-command handler. Handlers +# ChatOps front door: parse `/command` on PR comments, apply command-specific +# permissions, and repository_dispatch to a per-command handler. Handlers # own allowlists, SHA admission, and the actual work. Same-repo dispatch # uses github.token with contents: write; GitHub starts the handler run # for GITHUB_TOKEN-created repository_dispatch events. Do not pass the # long-lived GH_PERSONAL_ACCESS_TOKEN into this third-party action. # -# To add a command: list it under `commands` and add a workflow with +# To add a command: add it to `config` and add a workflow with # `on.repository_dispatch.types: ["-command"]`. on: @@ -30,8 +30,17 @@ jobs: uses: peter-evans/slash-command-dispatch@9bdcd7914ec1b75590b790b844aa3b8eee7c683a # v5.0.2 with: token: ${{ github.token }} - permission: admin - issue-type: pull-request - commands: | - query-regression - ci + # CI admission re-fetches the comment and checks author/member permissions. + config: >- + [ + {"command": "query-regression", "permission": "admin", "issue_type": "pull-request"}, + {"command": "ci", "permission": "none", "issue_type": "pull-request"} + ] + - name: Report CI dispatch failure + if: ${{ failure() && (github.event.comment.body == '/ci' || startsWith(github.event.comment.body, '/ci ')) }} + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.issue.number }} + run: | + gh api --method POST "/repos/${GITHUB_REPOSITORY}/issues/${PR_NUMBER}/comments" \ + -f "body=CI command dispatch failed. See ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for details." From 88197f4019eb76f4d68b293b1f2b5c026a901a49 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 23 Sep 2026 11:08:18 +0000 Subject: [PATCH 022/133] fix(promql): derive vector matching result labels and reject ambiguous matchings (#9306) * fix(promql): derive vector matching result labels and reject ambiguous matchings A vector-vector binary operation projected one operand's whole tag set and inner-joined without any cardinality check, so `on()`/`ignoring()` did not reduce the result labels, `group_left`/`group_right` changed nothing, and a non-unique match group produced a cross product that PromQL cannot represent. Result labels now follow Prometheus `resultMetric`: `on(...)` keeps the matching labels, `ignoring(...)` drops them, and a group modifier keeps the many side's labels plus the `group_x(...)` labels taken from the one side. A label the one side does not carry is deleted from the result. The reduced label set no longer identifies the operand series, so `__tsid` is dropped from the context on this path. Cardinality is enforced with a `count(1) OVER (PARTITION BY match keys, ts)` window and a scalar UDF that fails the query on a repeated group: on the one side before the join, and on the result labels after it, matching where Prometheus raises each of its three errors. Series are unique by their whole tag set, so the window is only planted when the match keys drop a tag; plain arithmetic and `on()` plan exactly as before. Closes #9207, closes #9208, closes #9209. Signed-off-by: Dennis Zhuang * perf(promql): keep __tsid when the result labels are an operand's whole tag set Deriving the result labels dropped `__tsid` from the context unconditionally, so an enclosing operation fell back to joining on the tag columns even where the column still identified the result series. Keep it when every result label comes from one operand and covers that operand's whole tag set: no other operand value reaches the labels, and the matching gives each of its rows a single partner, so its `__tsid` is still one per result series. That is the common `on()` and bare `group_left` shape; a matching that actually drops a tag still clears it. The column is re-qualified as the result's own, which is how the enclosing expression and the context look it up. Signed-off-by: Dennis Zhuang * fix(promql): keep the match group count column unambiguous The cardinality check aliased its row count to a fixed `__promql_match_group_count`. An operand carrying a label of that name made the window output two fields with the same name, and planning failed with "Schema contains qualified field name collide_right.__promql_match_group_count and unqualified field name __promql_match_group_count which would be ambiguous". Pick a name the operand does not already have, the way the `or` operator allocates its match key columns. Signed-off-by: Dennis Zhuang * test(promql): cover a match group spread over several regions The cardinality check runs above the merge of the region scans, so it counts a match group globally. Nothing pinned that: every table in these cases holds a single region, and a check evaluated per region would pass them all. Partition the operand on a column outside the match keys, which puts the two series of one match group in different regions, and assert both the pre-join and the post-join check still reject it. The case runs in the distributed environment too, where the regions sit on different datanodes. Signed-off-by: Dennis Zhuang * fix(promql): group an outer aggregate on the labels the operand kept `by`/`without` planning topped up a missing grouping column by walking down to the table scan and re-projecting it. That is right for a column the scan pruned for efficiency, but the labels a matching modifier deletes are also absent from the operand's output, and they were restored the same way: sum without(host) (a / on(host) b) `on(host)` leaves the operand with `host` alone, so the sum covers everything and Prometheus answers `{} 10`. Instead `device` came back from the scan under `a` and split the result into `{device="d1"} 5` and `{device="d2"} 5`. Same for `sum by(device)` of that operand, which has no `device` to group on at all. Take the grouping labels from the operand's own label set rather than from the row keys of the scan beneath it. A label pruned from the plan is still in that set and still gets restored; a label the operand dropped is not. Signed-off-by: Dennis Zhuang * refactor(promql): drop the unreachable aggregation tag top-up `by`/`without` planning could restore a grouping column that the plan no longer carried by rewriting the scan underneath it. Once the grouping labels come from the operand's own label set, there is nothing left for it to restore: a scan projects every label of `ctx.tag_columns` (`scan_tag_columns` only ever adds matcher columns to that set), so a label in the set is always in the schema. Stubbing the rewriter to a no-op passed the whole sqlness suite, in both the standalone and the distributed environment. Signed-off-by: Dennis Zhuang * test(promql): drop cases that guarded the removed tag top-up Three plain selector aggregates were there to show that restoring a pruned grouping column still worked. With the restore gone they only repeat what the aggregate cases already cover. Also fix a comment that still said the metric engine scan prunes tag columns: it projects every label of the operand. Signed-off-by: Dennis Zhuang * docs(promql): note the duplicate a propagated matcher hides A matcher copied onto the one-side operand removes groups without a partner before the cardinality check sees them, so a duplicate in such a group is not reported. Prometheus checks every group of the one side and fails the query. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/promql/src/functions.rs | 2 + src/promql/src/functions/vector_matching.rs | 208 +++++ src/query/src/promql/planner.rs | 742 ++++++++++++------ .../src/promql/planner/matching_filters.rs | 7 +- .../common/promql/matching_groups.result | 108 ++- .../common/promql/matching_groups.sql | 8 +- .../common/promql/set_operation.result | 32 +- .../promql/tsid_binary_join_regression.result | 51 +- .../promql/tsid_binary_join_regression.sql | 5 + .../common/promql/vector_matching.result | 406 ++++++++++ .../common/promql/vector_matching.sql | 221 ++++++ .../common/tql/matching_filter.result | 79 +- .../standalone/common/tql/matching_filter.sql | 4 +- 13 files changed, 1484 insertions(+), 389 deletions(-) create mode 100644 src/promql/src/functions/vector_matching.rs create mode 100644 tests/cases/standalone/common/promql/vector_matching.result create mode 100644 tests/cases/standalone/common/promql/vector_matching.sql diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index 22cd64b9147..34618f0a953 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -27,6 +27,7 @@ mod resets; mod round; #[cfg(test)] mod test_util; +mod vector_matching; pub use aggr_over_time::{ AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime, @@ -61,6 +62,7 @@ pub use quantile::QuantileOverTime; pub use quantile_aggr::{QUANTILE_NAME, quantile_udaf}; pub use resets::Resets; pub use round::Round; +pub use vector_matching::{MatchGroupViolation, UniqueMatchGroup}; use crate::range_array::RangeArray; diff --git a/src/promql/src/functions/vector_matching.rs b/src/promql/src/functions/vector_matching.rs new file mode 100644 index 00000000000..fdef83f6f05 --- /dev/null +++ b/src/promql/src/functions/vector_matching.rs @@ -0,0 +1,208 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Cardinality checks for PromQL vector matching. + +use datafusion::arrow::array::{Array, Int64Array}; +use datafusion::arrow::util::display::array_value_to_string; +use datafusion::common::{DataFusionError, Result as DfResult}; +use datafusion::logical_expr::{ScalarUDF, Volatility}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_common::ScalarValue; +use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datatypes::arrow::datatypes::DataType; + +use crate::functions::extract_array; + +/// The rule a repeated match group violates. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MatchGroupViolation { + /// Several series share a match group on the side that must hold one series per group. + DuplicateOnOneSide { one_side_is_left: bool }, + /// One-to-one matching found several matches for the same group. + ImplicitManyToOne, + /// A group modifier left several matches with the same result label set. + AmbiguousGroupLabels, +} + +impl MatchGroupViolation { + fn message(&self, group: &str) -> String { + match self { + Self::DuplicateOnOneSide { one_side_is_left } => { + let side = if *one_side_is_left { "left" } else { "right" }; + format!( + "found duplicate series for the match group {group} on the {side} hand-side \ + of the operation; many-to-many matching not allowed: matching labels must \ + be unique on one side" + ) + } + Self::ImplicitManyToOne => format!( + "multiple matches for labels {group}: many-to-one matching must be explicit \ + (group_left/group_right)" + ), + Self::AmbiguousGroupLabels => format!( + "multiple matches for labels {group}: grouping labels must ensure unique matches" + ), + } + } +} + +/// Rejects a vector matching whose match groups are not unique. +/// +/// Takes the per-group row count as first argument and the label columns that form the group +/// as the remaining ones. Returns `true` when every group holds a single row, and fails the +/// query otherwise; PromQL has no way to express the duplicate series it would produce. +pub struct UniqueMatchGroup; + +impl UniqueMatchGroup { + pub const fn name() -> &'static str { + "prom_assert_unique_match_group" + } + + pub fn scalar_udf(labels: Vec, violation: MatchGroupViolation) -> ScalarUDF { + ScalarUDF::new_from_impl(AssertUniqueMatchGroup { + signature: Signature::variadic_any(Volatility::Volatile), + labels, + violation, + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct AssertUniqueMatchGroup { + signature: Signature, + /// Label names of the group, in the same order as the label arguments. + labels: Vec, + violation: MatchGroupViolation, +} + +impl AssertUniqueMatchGroup { + fn render_group(&self, args: &[ColumnarValue], row: usize) -> DfResult { + let mut rendered = Vec::with_capacity(self.labels.len()); + for (label, arg) in self.labels.iter().zip(args) { + let array = extract_array(arg)?; + // A scalar argument was expanded to a single row above. + let row = if array.len() == 1 { 0 } else { row }; + if row >= array.len() || array.is_null(row) { + continue; + } + let value = array_value_to_string(&array, row) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + if value.is_empty() { + continue; + } + rendered.push(format!("{label}=\"{value}\"")); + } + Ok(format!("{{{}}}", rendered.join(", "))) + } +} + +impl ScalarUDFImpl for AssertUniqueMatchGroup { + fn name(&self) -> &str { + UniqueMatchGroup::name() + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let Some((counts, labels)) = args.args.split_first() else { + return Err(DataFusionError::Execution(format!( + "{} expects the match group row count as first argument", + UniqueMatchGroup::name() + ))); + }; + let counts = extract_array(counts)?; + let counts = counts + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{} expects an Int64 match group row count, found {}", + UniqueMatchGroup::name(), + counts.data_type() + )) + })?; + + if let Some(row) = + (0..counts.len()).find(|row| !counts.is_null(*row) && counts.value(*row) > 1) + { + let group = self.render_group(labels, row)?; + return Err(DataFusionError::Execution(self.violation.message(&group))); + } + + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::StringArray; + use datafusion::arrow::datatypes::Field; + + use super::*; + + fn invoke(counts: Vec, hosts: Vec>) -> DfResult { + let udf = UniqueMatchGroup::scalar_udf( + vec!["host".to_string()], + MatchGroupViolation::ImplicitManyToOne, + ); + let number_rows = counts.len(); + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(counts))), + ColumnarValue::Array(Arc::new(StringArray::from(hosts))), + ], + arg_fields: vec![ + Arc::new(Field::new("count", DataType::Int64, true)), + Arc::new(Field::new("host", DataType::Utf8, true)), + ], + number_rows, + return_field: Arc::new(Field::new("assert", DataType::Boolean, false)), + config_options: Arc::new(Default::default()), + }) + } + + #[test] + fn unique_groups_pass() { + let result = invoke(vec![1, 1], vec![Some("a"), Some("b")]).unwrap(); + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) + )); + } + + #[test] + fn duplicate_group_reports_its_labels() { + let err = invoke(vec![1, 2], vec![Some("a"), Some("b")]).unwrap_err(); + assert!( + err.to_string() + .contains("multiple matches for labels {host=\"b\"}"), + "{err}" + ); + } + + #[test] + fn null_label_is_omitted_from_the_group() { + let err = invoke(vec![2], vec![None]).unwrap_err(); + assert!(err.to_string().contains("labels {}"), "{err}"); + } +} diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 44211b8d2b7..2142ee82371 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -50,15 +50,12 @@ use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality, TableReference}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; -use datafusion_expr::{ - ExprSchemable, Literal, Projection, SortExpr, TableScanBuilder, TableSource, col, lit, -}; +use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit}; use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; @@ -71,7 +68,7 @@ use promql::extension_plan::{ }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + IDelta, Increase, LastOverTime, MatchGroupViolation, MaxOverTime, MinOverTime, MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, @@ -82,7 +79,8 @@ use promql::functions::{ NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, - Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup, + quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -154,6 +152,8 @@ const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_"; const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_"; const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_"; +/// Per-match-group row count the vector matching cardinality check is built on. +const MATCH_GROUP_COUNT_COLUMN: &str = "__promql_match_group_count"; /// Threshold for scatter scan mode const MAX_SCATTER_POINTS: i64 = 400; @@ -295,6 +295,35 @@ struct PlannedIslandLeaf { display_table: String, } +/// Result labels a vector-vector binary operation derives from its matching modifier, projected +/// from the operand each label belongs to. +#[derive(Debug)] +struct BinaryResultLabels { + exprs: Vec, + names: Vec, + aggregation_field_labels: Vec, + /// `__tsid` column of the operand the labels come from, when that operand contributes its + /// whole tag set and the column still identifies the result series. + tsid: Option, +} + +impl BinaryResultLabels { + fn apply(&self, ctx: &mut PromPlannerContext) { + ctx.tag_columns = self.names.clone(); + ctx.aggregation_field_labels = self.aggregation_field_labels.clone(); + ctx.use_tsid = self.tsid.is_some(); + } + + /// `__tsid` is projected from whichever operand the labels come from, so it carries that + /// operand's qualifier. Re-qualify it as the result's own, which is what the context names + /// and what the enclosing expression looks the column up by. + fn tsid_projection(&self, table_ref: Option) -> Option { + self.tsid + .clone() + .map(|tsid| tsid.alias_qualified(table_ref, DATA_SCHEMA_TSID_COLUMN_NAME)) + } +} + #[derive(Debug)] struct IslandFieldExprs { exprs: Vec, @@ -622,40 +651,12 @@ impl PromPlanner { param, } = aggr_expr; - let mut input = self.prom_expr_to_plan(expr, query_engine_state).await?; + let input = self.prom_expr_to_plan(expr, query_engine_state).await?; let input_has_tsid = input.schema().fields().iter().any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME && field.data_type() == &ArrowDataType::UInt64 }); - // `__tsid` based scan projection may prune tag columns. Ensure tags referenced in - // aggregation modifiers (`by`/`without`) are available before planning group keys. - let required_group_tags = match modifier { - None => BTreeSet::new(), - Some(LabelModifier::Include(labels)) => labels - .labels - .iter() - .filter(|label| !is_metric_engine_internal_column(label.as_str())) - .cloned() - .collect(), - Some(LabelModifier::Exclude(labels)) => { - let mut all_tags = self.collect_row_key_tag_columns_from_plan(&input)?; - for label in &labels.labels { - let _ = all_tags.remove(label); - } - all_tags - } - }; - - if !required_group_tags.is_empty() - && required_group_tags - .iter() - .any(|tag| Self::find_case_sensitive_column(input.schema(), tag.as_str()).is_none()) - { - input = self.ensure_tag_columns_available(input, &required_group_tags)?; - self.refresh_tag_columns_from_schema(input.schema()); - } - match (*op).id() { token::T_TOPK | token::T_BOTTOMK => { self.prom_topk_bottomk_to_plan(aggr_expr, input).await @@ -1634,6 +1635,16 @@ impl PromPlanner { // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast. let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric; + let only_join_time_index = lhs.value_type() == ValueType::Scalar + || rhs.value_type() == ValueType::Scalar + || has_empty_metric_operand + || ((left_context.tag_columns.is_empty() + || right_context.tag_columns.is_empty()) + && !left_context + .tag_columns + .iter() + .chain(&right_context.tag_columns) + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)); let join_plan = self.join_on_non_field_columns( left_input, right_input, @@ -1641,21 +1652,34 @@ impl PromPlanner { right_table_ref.clone(), left_time_index_column, right_time_index_column, - lhs.value_type() == ValueType::Scalar - || rhs.value_type() == ValueType::Scalar - || has_empty_metric_operand - || ((left_context.tag_columns.is_empty() - || right_context.tag_columns.is_empty()) - && !left_context - .tag_columns - .iter() - .chain(&right_context.tag_columns) - .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL)), + only_join_time_index, modifier, &left_context, &right_context, )?; let join_plan_schema = join_plan.schema().clone(); + // The matching modifier derives the result labels instead of the operands keeping + // their own tag set. An operand that is broadcast rather than matched (a scalar, + // or a vector without tags) keeps the other side's labels as before. + let result_labels = if only_join_time_index { + None + } else { + Self::binary_result_labels(&left_context, &right_context, modifier) + .map(|labels| { + Self::binary_result_label_projection( + &join_plan_schema, + &left_table_ref, + &right_table_ref, + &left_context, + &right_context, + labels, + ) + }) + .transpose()? + }; + if let Some(labels) = &result_labels { + labels.apply(&mut self.ctx); + } let promql_annotations = self.promql_annotations.clone(); // These predicates always pass; they only evaluate otherwise-discarded pairs // while collecting annotations. @@ -1776,10 +1800,18 @@ impl PromPlanner { ), }; project_context.field_columns = project_field_columns; - self.project_binary_join_side(filtered, project_table_ref, &project_context) + self.project_binary_join_side( + filtered, + project_table_ref, + &project_context, + result_labels.as_ref(), + ) } else { - let projected = - self.projection_for_each_field_column(join_plan, bin_expr_builder)?; + let projected = self.projection_for_each_field_column_with_labels( + join_plan, + result_labels.as_ref(), + bin_expr_builder, + )?; let preserve_any_value = Self::field_columns_are_alternative_samples( projected.schema(), &self.ctx.field_columns, @@ -1867,6 +1899,7 @@ impl PromPlanner { input: LogicalPlan, table_ref: &TableReference, context: &PromPlannerContext, + result_labels: Option<&BinaryResultLabels>, ) -> Result { let schema = input.schema(); @@ -1891,20 +1924,27 @@ impl PromPlanner { project_exprs.push(DfExpr::Column(field_col)); } - // Project tag columns from the chosen side. - for tag_column in &context.tag_columns { - let tag_col = schema - .qualified_field_with_name(Some(table_ref), tag_column) - .context(DataFusionPlanningSnafu)? - .into(); - project_exprs.push(DfExpr::Column(tag_col)); + // Project tag columns: the labels the matching modifier derived, or the chosen side's. + match result_labels { + Some(labels) => project_exprs.extend(labels.exprs.iter().cloned()), + None => { + for tag_column in &context.tag_columns { + let tag_col = schema + .qualified_field_with_name(Some(table_ref), tag_column) + .context(DataFusionPlanningSnafu)? + .into(); + project_exprs.push(DfExpr::Column(tag_col)); + } + } } // Preserve `__tsid` if present, so it can still be used internally downstream. It's // stripped from the final output anyway. - if let Some(tsid_col) = - Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid) - { + let tsid_col = match result_labels { + Some(labels) => labels.tsid_projection(Some(table_ref.clone())), + None => Self::optional_tsid_projection(schema, Some(table_ref), context.use_tsid), + }; + if let Some(tsid_col) = tsid_col { project_exprs.push(tsid_col); } @@ -1919,6 +1959,9 @@ impl PromPlanner { self.ctx = context.clone(); self.ctx.table_name = None; self.ctx.schema_name = None; + if let Some(labels) = result_labels { + labels.apply(&mut self.ctx); + } Ok(plan) } @@ -3214,130 +3257,6 @@ impl PromPlanner { Ok(out) } - fn ensure_tag_columns_available( - &self, - plan: LogicalPlan, - required_tags: &BTreeSet, - ) -> Result { - if required_tags.is_empty() { - return Ok(plan); - } - - struct Rewriter { - required_tags: BTreeSet, - } - - impl TreeNodeRewriter for Rewriter { - type Node = LogicalPlan; - - fn f_up( - &mut self, - node: Self::Node, - ) -> datafusion_common::Result> { - match node { - LogicalPlan::TableScan(scan) => { - let schema = scan.source.schema(); - let mut projection = match scan.projection.clone() { - Some(p) => p, - None => { - // Scanning all columns already covers required tags. - return Ok(Transformed::no(LogicalPlan::TableScan(scan))); - } - }; - - let mut changed = false; - for tag in &self.required_tags { - if let Some((idx, _)) = schema - .fields() - .iter() - .enumerate() - .find(|(_, field)| field.name() == tag) - && !projection.contains(&idx) - { - projection.push(idx); - changed = true; - } - } - - if !changed { - return Ok(Transformed::no(LogicalPlan::TableScan(scan))); - } - - projection.sort_unstable(); - projection.dedup(); - - let new_scan = - TableScanBuilder::new(scan.table_name.clone(), scan.source.clone()) - .with_projection(Some(projection)) - .with_filters(scan.filters) - .with_fetch(scan.fetch) - .build()?; - Ok(Transformed::yes(LogicalPlan::TableScan(new_scan))) - } - LogicalPlan::Projection(proj) => { - let input_schema = proj.input.schema(); - - let existing = proj - .schema - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - - let mut expr = proj.expr.clone(); - let mut has_changed = false; - for tag in &self.required_tags { - if existing.contains(tag.as_str()) { - continue; - } - - if let Some(idx) = input_schema.index_of_column_by_name(None, tag) { - expr.push(DfExpr::Column(Column::from( - input_schema.qualified_field(idx), - ))); - has_changed = true; - } - } - - if !has_changed { - return Ok(Transformed::no(LogicalPlan::Projection(proj))); - } - - let new_proj = Projection::try_new(expr, proj.input)?; - Ok(Transformed::yes(LogicalPlan::Projection(new_proj))) - } - other => Ok(Transformed::no(other)), - } - } - } - - let mut rewriter = Rewriter { - required_tags: required_tags.clone(), - }; - let rewritten = plan - .rewrite(&mut rewriter) - .context(DataFusionPlanningSnafu)?; - Ok(rewritten.data) - } - - fn refresh_tag_columns_from_schema(&mut self, schema: &DFSchemaRef) { - let time_index = self.ctx.time_index_column.as_deref(); - let field_columns = self.ctx.field_columns.iter().collect::>(); - - let mut tags = schema - .fields() - .iter() - .map(|f| f.name()) - .filter(|name| Some(name.as_str()) != time_index) - .filter(|name| !field_columns.contains(name)) - .filter(|name| !is_metric_engine_internal_column(name)) - .cloned() - .collect::>(); - tags.sort_unstable(); - tags.dedup(); - self.ctx.tag_columns = tags; - } - /// Setup [PromPlannerContext]'s state fields. /// /// Returns a logical plan for an empty metric. @@ -5956,6 +5875,225 @@ impl PromPlanner { Ok((left_tag_columns, right_tag_columns, force_empty_join)) } + /// Result labels of a vector-vector binary operation, following Prometheus `resultMetric`: + /// `on(...)` keeps only the matching labels, `ignoring(...)` drops them, and a group modifier + /// keeps the "many" side's labels plus the `group_x(...)` labels taken from the "one" side. + /// + /// The flag of each entry tells which operand the label is projected from. `None` means the + /// operation keeps a whole operand tag set, which the default projection already does. + fn binary_result_labels( + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + modifier: &Option, + ) -> Option> { + let modifier = modifier.as_ref()?; + let (many_is_left, include) = match &modifier.card { + VectorMatchCardinality::OneToOne => (true, None), + VectorMatchCardinality::ManyToOne(labels) => (true, Some(labels)), + VectorMatchCardinality::OneToMany(labels) => (false, Some(labels)), + // Set operators keep their operands' labels and don't reach this path. + VectorMatchCardinality::ManyToMany => return None, + }; + + let Some(include) = include else { + let matching = modifier.matching.as_ref()?; + let reduced = |keep: bool, labels: &BTreeSet<&String>| { + left_context + .tag_columns + .iter() + .filter(|tag| labels.contains(tag) == keep) + .map(|tag| (true, tag.clone())) + .collect() + }; + return Some(match matching { + LabelModifier::Include(on) => reduced(true, &on.labels.iter().collect()), + LabelModifier::Exclude(ignoring) => { + reduced(false, &ignoring.labels.iter().collect()) + } + }); + }; + + let (many_context, one_context) = if many_is_left { + (left_context, right_context) + } else { + (right_context, left_context) + }; + let include = include.labels.iter().collect::>(); + // An included label the "one" side doesn't carry is deleted from the result, so it is + // dropped from the "many" side as well. + let mut labels = many_context + .tag_columns + .iter() + .filter(|tag| !include.contains(tag)) + .map(|tag| (many_is_left, tag.clone())) + .collect::>(); + labels.extend( + include + .into_iter() + .filter(|label| one_context.tag_columns.contains(label)) + .map(|label| (!many_is_left, label.clone())), + ); + Some(labels) + } + + /// Resolve [`Self::binary_result_labels`] against the join output. + fn binary_result_label_projection( + schema: &DFSchemaRef, + left_table_ref: &TableReference, + right_table_ref: &TableReference, + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + labels: Vec<(bool, String)>, + ) -> Result { + let mut exprs = Vec::with_capacity(labels.len()); + let mut names = Vec::with_capacity(labels.len()); + let mut aggregation_field_labels = Vec::new(); + let mut sources = HashSet::new(); + for (from_left, label) in labels { + let (table_ref, context) = if from_left { + (left_table_ref, left_context) + } else { + (right_table_ref, right_context) + }; + let field = schema + .qualified_field_with_name(Some(table_ref), &label) + .context(DataFusionPlanningSnafu)?; + exprs.push(DfExpr::Column(field.into())); + if context.aggregation_field_labels.contains(&label) { + aggregation_field_labels.push(label.clone()); + } + sources.insert(from_left); + names.push(label); + } + + // One operand contributing every one of its tags as the whole result label set keeps the + // one-to-one correspondence between its `__tsid` and a result series: no other operand + // value reaches the labels, and the matching gives each of its rows a single partner. + let source_context = match sources.into_iter().collect::>().as_slice() { + [true] => Some((left_table_ref, left_context)), + [false] => Some((right_table_ref, right_context)), + _ => None, + }; + let tsid = source_context + .filter(|(_, context)| { + context.use_tsid + && context.tag_columns.len() == names.len() + && context.tag_columns.iter().all(|tag| names.contains(tag)) + }) + .and_then(|(table_ref, _)| { + Self::optional_tsid_projection(schema, Some(table_ref), true) + }); + + Ok(BinaryResultLabels { + exprs, + names, + aggregation_field_labels, + tsid, + }) + } + + /// Whether the result of a binary operation can hold two series with the same labels, which + /// only [`Self::binary_result_labels`] can introduce: one-to-one matching on a subset of the + /// tags, or a group modifier that overwrites a label of the "many" side. + fn binary_result_labels_may_repeat( + left_context: &PromPlannerContext, + right_context: &PromPlannerContext, + modifier: &Option, + ) -> bool { + let Some(modifier) = modifier else { + return false; + }; + match &modifier.card { + VectorMatchCardinality::OneToOne => match &modifier.matching { + None => false, + Some(LabelModifier::Include(on)) => { + let on = on.labels.iter().collect::>(); + !left_context.tag_columns.iter().all(|tag| on.contains(tag)) + } + Some(LabelModifier::Exclude(ignoring)) => ignoring + .labels + .iter() + .any(|label| left_context.tag_columns.contains(label)), + }, + VectorMatchCardinality::ManyToOne(include) => include + .labels + .iter() + .any(|label| left_context.tag_columns.contains(label)), + VectorMatchCardinality::OneToMany(include) => include + .labels + .iter() + .any(|label| right_context.tag_columns.contains(label)), + VectorMatchCardinality::ManyToMany => false, + } + } + + /// Wrap `plan` in a check that fails the query when a match group holds more than one row at + /// a timestamp. `group_exprs` are the label columns of the group, resolved against `plan`. + fn assert_unique_match_group( + plan: LogicalPlan, + group_exprs: Vec, + group_labels: Vec, + time_index_expr: DfExpr, + violation: MatchGroupViolation, + ) -> Result { + let mut partition_by = group_exprs.clone(); + partition_by.push(time_index_expr); + // A label may carry the generated name, and the count column has to stay unambiguous + // against every column the operand already has. + let occupied_column_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(); + let mut next_suffix = 0; + let count_column = loop { + let name = match next_suffix { + 0 => MATCH_GROUP_COUNT_COLUMN.to_string(), + suffix => format!("{MATCH_GROUP_COUNT_COLUMN}_{suffix}"), + }; + next_suffix += 1; + if !occupied_column_names.contains(name.as_str()) { + break name; + } + }; + let count = DfExpr::WindowFunction(Box::new(WindowFunction { + fun: WindowFunctionDefinition::AggregateUDF(count_udaf()), + params: WindowFunctionParams { + args: vec![lit(1i64)], + partition_by, + order_by: vec![], + window_frame: WindowFrame::new(None), + null_treatment: None, + distinct: false, + filter: None, + }, + })) + .alias(count_column.as_str()); + + let output_exprs = plan + .schema() + .iter() + .map(|(qualifier, field)| DfExpr::Column(Column::new(qualifier.cloned(), field.name()))) + .collect::>(); + let assert_expr = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)), + args: std::iter::once(col(count_column.as_str())) + .chain(group_exprs) + .collect(), + }); + + LogicalPlanBuilder::from(plan) + .window(vec![count]) + .context(DataFusionPlanningSnafu)? + .filter(assert_expr) + .context(DataFusionPlanningSnafu)? + .project(output_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + fn binary_modifier_preserves_tsid_join_key( &self, left_context: &PromPlannerContext, @@ -6023,7 +6161,7 @@ impl PromPlanner { && !force_empty_join && left_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]) && right_tag_columns == BTreeSet::from([DATA_SCHEMA_TSID_COLUMN_NAME.to_string()]); - let (left, right) = if !only_join_time_index + let (left, right, left_matched_tags, right_matched_tags) = if !only_join_time_index && !use_tsid_join && Self::only_temporality_match_label_mismatches(left_context, right_context, modifier) { @@ -6044,28 +6182,79 @@ impl PromPlanner { false, modifier, )?; - (left, right) + ( + left, + right, + aligned_left_context.tag_columns, + aligned_right_context.tag_columns, + ) } else { + ( + left, + right, + left_context.tag_columns.clone(), + right_context.tag_columns.clone(), + ) + }; + + // A join key that covers the whole tag set of a side already makes that side's match + // groups unique, and so does a join on `__tsid`. Only the reduced keys need the check. + let checked_matching = !only_join_time_index + && !force_empty_join + && !use_tsid_join + && !matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::ManyToMany) + ); + // The right operand is the "one" side of the matching, unless `group_right` swaps them. + let one_side_is_left = matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::OneToMany(_)) + ); + let (left, right) = if !checked_matching { (left, right) + } else if one_side_is_left { + ( + Self::assert_unique_one_side( + left, + &left_tag_columns, + &left_matched_tags, + left_time_index_column.as_deref(), + true, + )?, + right, + ) + } else { + ( + left, + Self::assert_unique_one_side( + right, + &right_tag_columns, + &right_matched_tags, + right_time_index_column.as_deref(), + false, + )?, + ) }; // push time index column if it exists - if let (Some(left_time_index_column), Some(right_time_index_column)) = - (left_time_index_column, right_time_index_column) - { + if let (Some(left_time_index_column), Some(right_time_index_column)) = ( + left_time_index_column.clone(), + right_time_index_column.clone(), + ) { left_tag_columns.insert(left_time_index_column); right_tag_columns.insert(right_time_index_column); } let right = LogicalPlanBuilder::from(right) - .alias(right_table_ref) + .alias(right_table_ref.clone()) .context(DataFusionPlanningSnafu)? .build() .context(DataFusionPlanningSnafu)?; // Inner Join on time index column to concat two operator - LogicalPlanBuilder::from(left) - .alias(left_table_ref) + let join_plan = LogicalPlanBuilder::from(left) + .alias(left_table_ref.clone()) .context(DataFusionPlanningSnafu)? .join_detailed( right, @@ -6085,7 +6274,90 @@ impl PromPlanner { ) .context(DataFusionPlanningSnafu)? .build() - .context(DataFusionPlanningSnafu) + .context(DataFusionPlanningSnafu)?; + + // The result labels no longer identify a series on their own: the "many" side can hold + // several series per result label set, which PromQL cannot represent. + let labels = checked_matching + .then(|| Self::binary_result_labels(left_context, right_context, modifier)) + .flatten() + .filter(|_| { + Self::binary_result_labels_may_repeat(left_context, right_context, modifier) + }); + let Some(labels) = labels else { + return Ok(join_plan); + }; + let schema = join_plan.schema().clone(); + let group_exprs = labels + .iter() + .map(|(from_left, label)| { + let table_ref = if *from_left { + &left_table_ref + } else { + &right_table_ref + }; + schema + .qualified_field_with_name(Some(table_ref), label) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }) + .collect::>>()?; + let time_index_expr = left_time_index_column + .as_deref() + .map(|column| { + schema + .qualified_field_with_name(Some(&left_table_ref), column) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }) + .transpose()? + .with_context(|| UnexpectedPlanExprSnafu { + desc: "vector matching on a plan without a time index", + })?; + let violation = if matches!( + modifier.as_ref().map(|modifier| &modifier.card), + Some(VectorMatchCardinality::OneToOne) + ) { + MatchGroupViolation::ImplicitManyToOne + } else { + MatchGroupViolation::AmbiguousGroupLabels + }; + Self::assert_unique_match_group( + join_plan, + group_exprs, + labels.into_iter().map(|(_, label)| label).collect(), + time_index_expr, + violation, + ) + } + + /// Guard the side of a vector matching that must hold one series per match group. + fn assert_unique_one_side( + plan: LogicalPlan, + join_keys: &BTreeSet, + tag_columns: &[String], + time_index_column: Option<&str>, + one_side_is_left: bool, + ) -> Result { + let Some(time_index_column) = time_index_column else { + return Ok(plan); + }; + if tag_columns.iter().all(|tag| join_keys.contains(tag)) { + return Ok(plan); + } + + let group_labels = join_keys.iter().cloned().collect::>(); + let group_exprs = group_labels + .iter() + .map(|label| DfExpr::Column(Column::from_name(label))) + .collect(); + Self::assert_unique_match_group( + plan, + group_exprs, + group_labels, + DfExpr::Column(Column::from_name(time_index_column)), + MatchGroupViolation::DuplicateOnOneSide { one_side_is_left }, + ) } fn selected_binary_match_labels( @@ -7118,6 +7390,20 @@ impl PromPlanner { input: LogicalPlan, name_to_expr: F, ) -> Result + where + F: FnMut(&String) -> Result, + { + self.projection_for_each_field_column_with_labels(input, None, name_to_expr) + } + + /// Like [`Self::projection_for_each_field_column`], but projects `result_labels` instead of + /// the context tag columns when a binary operation derived its own result label set. + fn projection_for_each_field_column_with_labels( + &mut self, + input: LogicalPlan, + result_labels: Option<&BinaryResultLabels>, + name_to_expr: F, + ) -> Result where F: FnMut(&String) -> Result, { @@ -7128,22 +7414,30 @@ impl PromPlanner { let table_ref = self.ctx.table_name.clone().map(TableReference::bare); // Derived labels can be unqualified even when the context still names the source table. let input_schema = input.schema().clone(); - let non_field_columns_iter = self - .ctx - .tag_columns - .iter() - .chain(self.ctx.time_index_column.iter()) - .map(|col| { - input_schema - .qualified_field_with_name(table_ref.as_ref(), col) - .or_else(|_| input_schema.qualified_field_with_unqualified_name(col)) - .map(|field| DfExpr::Column(field.into())) - .context(DataFusionPlanningSnafu) - }); - let tsid_iter = - Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid) - .into_iter() - .map(Ok); + let lookup = |col: &String| { + input_schema + .qualified_field_with_name(table_ref.as_ref(), col) + .or_else(|_| input_schema.qualified_field_with_unqualified_name(col)) + .map(|field| DfExpr::Column(field.into())) + .context(DataFusionPlanningSnafu) + }; + let tag_columns_iter = match result_labels { + Some(labels) => labels.exprs.iter().cloned().map(Ok).collect::>(), + None => self.ctx.tag_columns.iter().map(lookup).collect::>(), + }; + let non_field_columns_iter = tag_columns_iter + .into_iter() + .chain(self.ctx.time_index_column.iter().map(lookup)); + let tsid_iter = match result_labels { + Some(labels) => labels.tsid_projection(table_ref.clone()), + None => Self::optional_tsid_projection( + input.schema(), + table_ref.as_ref(), + self.ctx.use_tsid, + ), + } + .into_iter() + .map(Ok); // build computation exprs let result_field_columns = self @@ -11848,20 +12142,26 @@ mod test { PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) .await .unwrap(); - let expected = "Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ - \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ - \n SubqueryAlias: http_server_requests_seconds_sum\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_sum\ - \n SubqueryAlias: http_server_requests_seconds_count\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_count"; + let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ + \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ + \n SubqueryAlias: http_server_requests_seconds_sum\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_sum\ + \n SubqueryAlias: http_server_requests_seconds_count\ + \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_count"; assert_eq!(plan.to_string(), expected); } diff --git a/src/query/src/promql/planner/matching_filters.rs b/src/query/src/promql/planner/matching_filters.rs index e38d48044da..d70e1bc4192 100644 --- a/src/query/src/promql/planner/matching_filters.rs +++ b/src/query/src/promql/planner/matching_filters.rs @@ -56,9 +56,10 @@ const LABEL_PRESERVING_RANGE_FUNCTIONS: [&str; 20] = [ /// 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 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). +/// matching label: a propagated matcher removes whole groups, so the cardinality check still +/// sees every group that takes part in the matching. A duplicate in a group without a partner +/// goes unreported as a result, where Prometheus fails the query on it. +/// Grouped matching is conservatively restricted to a provably unique one-side operand. /// /// `left_tags` and `right_tags` are the tag columns of the planned operands: a matcher may just /// as well constrain a value field, which the parsed expression does not distinguish from a diff --git a/tests/cases/standalone/common/promql/matching_groups.result b/tests/cases/standalone/common/promql/matching_groups.result index 183ae6d218c..04b1588b292 100644 --- a/tests/cases/standalone/common/promql/matching_groups.result +++ b/tests/cases/standalone/common/promql/matching_groups.result @@ -30,19 +30,16 @@ INSERT INTO matching_groups_right VALUES Affected Rows: 5 --- The result labels are the right operand's tag set rather than what the modifier implies --- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing --- only by it are indistinguishable. -- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); -+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.Float64(8) * greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ -| x | 1970-01-01T00:00:00 | 24.0 | -| x | 1970-01-01T00:00:00 | 48.0 | -+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.Float64(8) * greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 24.0 | +| x | b | 1970-01-01T00:00:00 | 48.0 | ++------+------+---------------------+--------------------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right() (matching_groups_right{host="x"} * 2); @@ -57,45 +54,45 @@ TQL EVAL (0, 0, '1s') sum by(host)(matching_groups_left) / on(host) group_right( -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host!="y"} / on(host) group_left max by(host)(matching_groups_right); -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| | 1970-01-01T00:00:00 | 6.0 | -| | 1970-01-01T00:00:00 | 6.0 | -| x | 1970-01-01T00:00:00 | 3.0 | -| x | 1970-01-01T00:00:00 | 6.0 | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host=~"x|y"} / on(host) group_left max by(host)(matching_groups_right); -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| x | 1970-01-01T00:00:00 | 3.0 | -| x | 1970-01-01T00:00:00 | 6.0 | -| y | 1970-01-01T00:00:00 | 0.36 | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | b | 1970-01-01T00:00:00 | 6.0 | +| y | a | 1970-01-01T00:00:00 | 0.36 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host=""} / on(host) group_left max by(host)(matching_groups_right); -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| | 1970-01-01T00:00:00 | 6.0 | -| | 1970-01-01T00:00:00 | 6.0 | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host!~"x|y"} / on(host) group_left max by(host)(matching_groups_right); -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| | 1970-01-01T00:00:00 | 6.0 | -| | 1970-01-01T00:00:00 | 6.0 | -+------+---------------------+-------------------------------------------------------------------------------------------------------+ ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +| | a | 1970-01-01T00:00:00 | 6.0 | +| | a | 1970-01-01T00:00:00 | 6.0 | ++------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -- bottomk partitions by host when zone is excluded from grouping. -- SQLNESS SORT_RESULT 3 1 @@ -130,46 +127,29 @@ TQL EVAL (0, 0, '1s') matching_groups_left / on(host,zone) bottomk(1, matching_g -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk(1, max by(host)(matching_groups_right)); -+------+----+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+----+-------------------------------------------------------------------------------------------------------+ -+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); -+------+----+-------------------------------------------------------------------------------------------------------+ -| host | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+----+-------------------------------------------------------------------------------------------------------+ -+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ +| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | ++------+------+----+-------------------------------------------------------------------------------------------------------+ ++------+------+----+-------------------------------------------------------------------------------------------------------+ --- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many --- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no --- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands --- alone, so the recorded output is the unchanged baseline. +-- Both one-sides carry two series per `host`, which the matching cannot resolve. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; -+------+------+---------------------+----------------------------------------------------------------------------+ -| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.greptime_value | -+------+------+---------------------+----------------------------------------------------------------------------+ -| x | a | 1970-01-01T00:00:00 | 12.0 | -| x | a | 1970-01-01T00:00:00 | 6.0 | -| x | b | 1970-01-01T00:00:00 | 3.0 | -| x | b | 1970-01-01T00:00:00 | 6.0 | -+------+------+---------------------+----------------------------------------------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="x"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left max by(host,zone)(matching_groups_right); -+------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| host | zone | ts | matching_groups_left.greptime_value / matching_groups_right.max(matching_groups_right.greptime_value) | -+------+------+---------------------+-------------------------------------------------------------------------------------------------------+ -| x | a | 1970-01-01T00:00:00 | 12.0 | -| x | a | 1970-01-01T00:00:00 | 6.0 | -| x | b | 1970-01-01T00:00:00 | 3.0 | -| x | b | 1970-01-01T00:00:00 | 6.0 | -+------+------+---------------------+-------------------------------------------------------------------------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="x"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side DROP TABLE matching_groups_left; diff --git a/tests/cases/standalone/common/promql/matching_groups.sql b/tests/cases/standalone/common/promql/matching_groups.sql index 97e46da0437..03bdf8eaaa7 100644 --- a/tests/cases/standalone/common/promql/matching_groups.sql +++ b/tests/cases/standalone/common/promql/matching_groups.sql @@ -19,9 +19,6 @@ INSERT INTO matching_groups_right VALUES ('x', 'a', 0, 2), ('x', 'b', 0, 4), ('y', 'a', 0, 100), ('', 'a', 0, 8), (NULL, 'a', 0, 10); --- The result labels are the right operand's tag set rather than what the modifier implies --- (#9207), so `zone` is absent wherever the right side aggregates it away and rows differing --- only by it are indistinguishable. -- Filtering whole ranking partitions preserves the aggregate and scalar arithmetic. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') (8 * matching_groups_left{host="x"}) / on(host) group_left topk by(host)(1, max by(host)(matching_groups_right)); @@ -52,10 +49,7 @@ TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk( -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left topk by(host)(1, topk(1, max by(host)(matching_groups_right))); --- Both one-sides carry two series per `host`. Prometheus rejects this with "many-to-many --- matching not allowed: matching labels must be unique on one side"; GreptimeDB has no --- cardinality check (#9209) and returns the cross product. The rewrite leaves these operands --- alone, so the recorded output is the unchanged baseline. +-- Both one-sides carry two series per `host`, which the matching cannot resolve. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 0, '1s') matching_groups_left{host="x"} / on(host) group_left matching_groups_right; -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result index 2f0041f1a36..db97fed7e6b 100644 --- a/tests/cases/standalone/common/promql/set_operation.result +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -688,26 +688,26 @@ tql eval (3, 4, '1s') cache_hit_with_null_label / (cache_miss_with_null_label + -- SQLNESS SORT_RESULT 3 1 tql eval (3, 4, '1s') cache_hit_with_null_label / ignoring(null_label) (cache_miss_with_null_label + ignoring(null_label) cache_hit_with_null_label); -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| job | null_label | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| read | | 1970-01-01T00:00:03 | 0.5 | -| read | | 1970-01-01T00:00:04 | 0.75 | -| write | | 1970-01-01T00:00:03 | 0.5 | -| write | | 1970-01-01T00:00:04 | 0.6666666666666666 | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| job | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| read | 1970-01-01T00:00:03 | 0.5 | +| read | 1970-01-01T00:00:04 | 0.75 | +| write | 1970-01-01T00:00:03 | 0.5 | +| write | 1970-01-01T00:00:04 | 0.6666666666666666 | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ -- SQLNESS SORT_RESULT 3 1 tql eval (3, 4, '1s') cache_hit_with_null_label / on(job) (cache_miss_with_null_label + on(job) cache_hit_with_null_label); -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| job | null_label | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ -| read | | 1970-01-01T00:00:03 | 0.5 | -| read | | 1970-01-01T00:00:04 | 0.75 | -| write | | 1970-01-01T00:00:03 | 0.5 | -| write | | 1970-01-01T00:00:04 | 0.6666666666666666 | -+-------+------------+---------------------+---------------------------------------------------------------------------------------------------------------+ ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| job | ts | lhs.greptime_value / rhs.cache_miss_with_null_label.greptime_value + cache_hit_with_null_label.greptime_value | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ +| read | 1970-01-01T00:00:03 | 0.5 | +| read | 1970-01-01T00:00:04 | 0.75 | +| write | 1970-01-01T00:00:03 | 0.5 | +| write | 1970-01-01T00:00:04 | 0.6666666666666666 | ++-------+---------------------+---------------------------------------------------------------------------------------------------------------+ drop table cache_hit_with_null_label; diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index d80d29ce000..54fd5ad5e26 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -243,12 +243,19 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[host@0 as host, job@1 as job, ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 / greptime_value@5 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@2), (ts@2, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED +| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@4, job@1), projection=[job@1, ts@3, greptime_value@0, greptime_value@2] REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, greptime_value@3 as greptime_value, ts@4 as ts, __promql_match_group_count@5 as __promql_match_group_count] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@2)], projection=[greptime_value@0, job@1, ts@2, greptime_value@3, ts@5], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([job@2, ts@4],REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@3, job@1), projection=[greptime_value@0, job@1, ts@2] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -282,12 +289,19 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, __tsid@2 as __tsid, greptime_value@3 / greptime_value@4 as tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@3)], projection=[job@4, ts@6, __tsid@5, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED +| 0_| 0_|_ProjectionExec: expr=[job@0 as job, ts@1 as ts, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value] REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@4, job@1), projection=[job@1, ts@3, greptime_value@0, greptime_value@2] REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, greptime_value@3 as greptime_value, ts@4 as ts, __promql_match_group_count@5 as __promql_match_group_count] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_CoalescePartitionsExec REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@2)], projection=[greptime_value@0, job@1, ts@2, greptime_value@3, ts@5], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@3],REDACTED +|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@1 as job, ts@3 as ts] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -566,13 +580,16 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[host@4, job@5, ts@7, __tsid@6, tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value@2, greptime_value@3], NullsEqual: true REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 / greptime_value@3 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@3)], projection=[ts@6, __tsid@5, greptime_value@0, greptime_value@3], NullsEqual: true REDACTED -|_|_|_RepartitionExec: partitioning=Hash([REDACTED -|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED +|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@3, ts@2)], projection=[ts@6, __tsid@2, greptime_value@0, greptime_value@4], NullsEqual: true REDACTED |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_FilterExec: prom_assert_unique_match_group(__promql_match_group_count@3, job@1), projection=[greptime_value@0, job@1, ts@2] REDACTED +|_|_|_WindowAggExec: wdw=[__promql_match_group_count: Ok(Field { name: "__promql_match_group_count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=Hash([REDACTED +|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| @@ -597,6 +614,20 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_| Total rows: 4_| +-+-+-+ +-- A group modifier that keeps the many side's whole tag set keeps its `__tsid` too, so the +-- enclosing operation matches on it. Same rows as the label-based plan. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; + ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ +| host | job | ts | tsid_binary_join_right.tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value / tsid_binary_join_left.greptime_value | ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ +| host1 | job1 | 1970-01-01T00:00:00 | 0.3333333333333333 | +| host1 | job1 | 1970-01-01T00:00:05 | 0.2 | +| host2 | job2 | 1970-01-01T00:00:00 | 0.16666666666666666 | +| host2 | job2 | 1970-01-01T00:00:05 | 0.14285714285714285 | ++-------+------+---------------------+--------------------------------------------------------------------------------------------------------------------------------------------+ + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql index 7a3bc897c75..627680330aa 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.sql @@ -226,6 +226,11 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; +-- A group modifier that keeps the many side's whole tag set keeps its `__tsid` too, so the +-- enclosing operation matches on it. Same rows as the label-based plan. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid_binary_join_right) / tsid_binary_join_left; + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; diff --git a/tests/cases/standalone/common/promql/vector_matching.result b/tests/cases/standalone/common/promql/vector_matching.result new file mode 100644 index 00000000000..5382e94967c --- /dev/null +++ b/tests/cases/standalone/common/promql/vector_matching.result @@ -0,0 +1,406 @@ +CREATE TABLE counter_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +CREATE TABLE gauge_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +INSERT INTO counter_metric VALUES + ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20), ('host2', 'eth0', 0, 30); + +Affected Rows: 3 + +INSERT INTO gauge_metric VALUES + ('host1', 'eth0', 0, 2), ('host1', 'eth1', 0, 4), ('host2', 'eth0', 0, 5); + +Affected Rows: 3 + +-- Matching on the whole tag set keeps it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / gauge_metric; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `on(...)` reduces the result to the matching labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) gauge_metric{device="eth0"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 5.0 | +| host2 | 1970-01-01T00:00:00 | 6.0 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `ignoring(...)` drops the ignored labels, here from operands that don't even share them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / ignoring(device) gauge_metric{device="eth1"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 2.5 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `on()` matches everything against everything, so both sides must hold a single series. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{host="host1",device="eth0"} / on() gauge_metric{host="host1",device="eth0"}; + ++---------------------+-------------------------------------------------------------+ +| ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++---------------------+-------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 5.0 | ++---------------------+-------------------------------------------------------------+ + +-- A filtering comparison keeps the left sample values and the derived labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > on(host) gauge_metric{device="eth0"}; + ++---------------------+----------------+-------+ +| ts | greptime_value | host | ++---------------------+----------------+-------+ +| 1970-01-01T00:00:00 | 10.0 | host1 | +| 1970-01-01T00:00:00 | 30.0 | host2 | ++---------------------+----------------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > bool on(host) gauge_metric{device="eth0"}; + ++-------+---------------------+-------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value > gauge_metric.greptime_value | ++-------+---------------------+-------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 1.0 | +| host2 | 1970-01-01T00:00:00 | 1.0 | ++-------+---------------------+-------------------------------------------------------------+ + +-- `host1` carries two devices on both sides: the match group repeats on the one side. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +-- The one side is unique here, the many side is not, and no group modifier makes it explicit. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) + +-- `group_left` takes the result labels from the many side. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left gauge_metric{device="eth0"}; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `group_right` swaps the sides. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_right gauge_metric; + ++-------+--------+---------------------+-------------------------------------------------------------+ +| host | device | ts | counter_metric.greptime_value / gauge_metric.greptime_value | ++-------+--------+---------------------+-------------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:00 | 2.5 | +| host2 | eth0 | 1970-01-01T00:00:00 | 6.0 | ++-------+--------+---------------------+-------------------------------------------------------------+ + +-- `group_right` makes the left operand the one side, which `host1` breaks. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_right gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the left hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +-- `group_left(device)` copies `device` from the one side, so both `host1` rows end up with the +-- same labels. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left(device) gauge_metric{device="eth0"}; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1", device="eth0"}: grouping labels must ensure unique matches + +-- An included label the one side doesn't carry is dropped from the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_left(device) sum by(host)(gauge_metric); + ++-------+---------------------+-------------------------------------------------------------------------------+ +| host | ts | counter_metric.greptime_value / gauge_metric.sum(gauge_metric.greptime_value) | ++-------+---------------------+-------------------------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 1.6666666666666667 | +| host2 | 1970-01-01T00:00:00 | 6.0 | ++-------+---------------------+-------------------------------------------------------------------------------+ + +DROP TABLE counter_metric; + +Affected Rows: 0 + +DROP TABLE gauge_metric; + +Affected Rows: 0 + +-- A label can carry the name the cardinality check generates for its row count. +CREATE TABLE collide_left ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +Affected Rows: 0 + +CREATE TABLE collide_right ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +Affected Rows: 0 + +INSERT INTO collide_left VALUES ('host1', 'a', 0, 10); + +Affected Rows: 1 + +INSERT INTO collide_right VALUES ('host1', 'b', 0, 2); + +Affected Rows: 1 + +TQL EVAL (0, 0, '5s') collide_left / on(host) collide_right; + ++-------+---------------------+------------------------------------------------------------+ +| host | ts | collide_left.greptime_value / collide_right.greptime_value | ++-------+---------------------+------------------------------------------------------------+ +| host1 | 1970-01-01T00:00:00 | 5.0 | ++-------+---------------------+------------------------------------------------------------+ + +DROP TABLE collide_left; + +Affected Rows: 0 + +DROP TABLE collide_right; + +Affected Rows: 0 + +-- Partitioning on a column outside the match keys puts one match group in several regions. The +-- cardinality check counts the merged data, so it must still see the duplicate. +CREATE TABLE spread_left ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +) +PARTITION ON COLUMNS (device) ( + device < 'eth1', + device >= 'eth1' +); + +Affected Rows: 0 + +CREATE TABLE spread_right ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +INSERT INTO spread_left VALUES ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20); + +Affected Rows: 2 + +INSERT INTO spread_right VALUES ('host1', 0, 2); + +Affected Rows: 1 + +TQL EVAL (0, 0, '5s') spread_left / on(host) spread_right; + +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) + +TQL EVAL (0, 0, '5s') spread_right / on(host) group_left spread_left; + +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side + +DROP TABLE spread_left; + +Affected Rows: 0 + +DROP TABLE spread_right; + +Affected Rows: 0 + +-- An outer aggregate groups on the labels the operand has left, not on the ones the scan +-- underneath could still produce. +CREATE TABLE outer_agg_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +Affected Rows: 0 + +CREATE TABLE outer_agg_b ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +INSERT INTO outer_agg_a VALUES ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +Affected Rows: 2 + +INSERT INTO outer_agg_b VALUES ('h1', 0, 2), ('h2', 0, 4); + +Affected Rows: 2 + +-- `on(host)` leaves only `host`, so `without(host)` groups everything together. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_a / on(host) outer_agg_b); + ++---------------------+--------------------------------------------------------------+ +| ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++---------------------+--------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+--------------------------------------------------------------+ + +-- `device` is not a label of the operand any more, so it groups everything together too. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_a / on(host) outer_agg_b); + ++---------------------+--------------------------------------------------------------+ +| ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++---------------------+--------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_a / on(host) outer_agg_b); + ++------+---------------------+--------------------------------------------------------------+ +| host | ts | sum(outer_agg_a.greptime_value / outer_agg_b.greptime_value) | ++------+---------------------+--------------------------------------------------------------+ +| h1 | 1970-01-01T00:00:00 | 5.0 | +| h2 | 1970-01-01T00:00:00 | 5.0 | ++------+---------------------+--------------------------------------------------------------+ + +DROP TABLE outer_agg_a; + +Affected Rows: 0 + +DROP TABLE outer_agg_b; + +Affected Rows: 0 + +-- Same over the metric engine, whose operands carry `__tsid`. +CREATE TABLE outer_agg_physical ( + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, +) ENGINE = metric WITH ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE outer_agg_metric_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host, device), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +Affected Rows: 0 + +CREATE TABLE outer_agg_metric_b ( + host STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +Affected Rows: 0 + +INSERT INTO outer_agg_metric_a (host, device, ts, greptime_value) VALUES + ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +Affected Rows: 2 + +INSERT INTO outer_agg_metric_b (host, ts, greptime_value) VALUES + ('h1', 0, 2), ('h2', 0, 4); + +Affected Rows: 2 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++---------------------+----------------------------------------------------------------------------+ +| ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++---------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++---------------------+----------------------------------------------------------------------------+ +| ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++---------------------+----------------------------------------------------------------------------+ +| 1970-01-01T00:00:00 | 10.0 | ++---------------------+----------------------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + ++------+---------------------+----------------------------------------------------------------------------+ +| host | ts | sum(outer_agg_metric_a.greptime_value / outer_agg_metric_b.greptime_value) | ++------+---------------------+----------------------------------------------------------------------------+ +| h1 | 1970-01-01T00:00:00 | 5.0 | +| h2 | 1970-01-01T00:00:00 | 5.0 | ++------+---------------------+----------------------------------------------------------------------------+ + +DROP TABLE outer_agg_metric_a; + +Affected Rows: 0 + +DROP TABLE outer_agg_metric_b; + +Affected Rows: 0 + +DROP TABLE outer_agg_physical; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/vector_matching.sql b/tests/cases/standalone/common/promql/vector_matching.sql new file mode 100644 index 00000000000..1fd0091c901 --- /dev/null +++ b/tests/cases/standalone/common/promql/vector_matching.sql @@ -0,0 +1,221 @@ +CREATE TABLE counter_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +CREATE TABLE gauge_metric ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +INSERT INTO counter_metric VALUES + ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20), ('host2', 'eth0', 0, 30); + +INSERT INTO gauge_metric VALUES + ('host1', 'eth0', 0, 2), ('host1', 'eth1', 0, 4), ('host2', 'eth0', 0, 5); + +-- Matching on the whole tag set keeps it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / gauge_metric; + +-- `on(...)` reduces the result to the matching labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) gauge_metric{device="eth0"}; + +-- `ignoring(...)` drops the ignored labels, here from operands that don't even share them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / ignoring(device) gauge_metric{device="eth1"}; + +-- `on()` matches everything against everything, so both sides must hold a single series. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{host="host1",device="eth0"} / on() gauge_metric{host="host1",device="eth0"}; + +-- A filtering comparison keeps the left sample values and the derived labels. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > on(host) gauge_metric{device="eth0"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} > bool on(host) gauge_metric{device="eth0"}; + +-- `host1` carries two devices on both sides: the match group repeats on the one side. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric; + +-- The one side is unique here, the many side is not, and no group modifier makes it explicit. +TQL EVAL (0, 0, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; + +-- `group_left` takes the result labels from the many side. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left gauge_metric{device="eth0"}; + +-- `group_right` swaps the sides. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_right gauge_metric; + +-- `group_right` makes the left operand the one side, which `host1` breaks. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_right gauge_metric{device="eth0"}; + +-- `group_left(device)` copies `device` from the one side, so both `host1` rows end up with the +-- same labels. +TQL EVAL (0, 0, '5s') counter_metric / on(host) group_left(device) gauge_metric{device="eth0"}; + +-- An included label the one side doesn't carry is dropped from the result. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') counter_metric{device="eth0"} / on(host) group_left(device) sum by(host)(gauge_metric); + +DROP TABLE counter_metric; + +DROP TABLE gauge_metric; + +-- A label can carry the name the cardinality check generates for its row count. +CREATE TABLE collide_left ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +CREATE TABLE collide_right ( + host STRING NULL, + __promql_match_group_count STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, __promql_match_group_count) +); + +INSERT INTO collide_left VALUES ('host1', 'a', 0, 10); + +INSERT INTO collide_right VALUES ('host1', 'b', 0, 2); + +TQL EVAL (0, 0, '5s') collide_left / on(host) collide_right; + +DROP TABLE collide_left; + +DROP TABLE collide_right; + +-- Partitioning on a column outside the match keys puts one match group in several regions. The +-- cardinality check counts the merged data, so it must still see the duplicate. +CREATE TABLE spread_left ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +) +PARTITION ON COLUMNS (device) ( + device < 'eth1', + device >= 'eth1' +); + +CREATE TABLE spread_right ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +INSERT INTO spread_left VALUES ('host1', 'eth0', 0, 10), ('host1', 'eth1', 0, 20); + +INSERT INTO spread_right VALUES ('host1', 0, 2); + +TQL EVAL (0, 0, '5s') spread_left / on(host) spread_right; + +TQL EVAL (0, 0, '5s') spread_right / on(host) group_left spread_left; + +DROP TABLE spread_left; + +DROP TABLE spread_right; + +-- An outer aggregate groups on the labels the operand has left, not on the ones the scan +-- underneath could still produce. +CREATE TABLE outer_agg_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host, device) +); + +CREATE TABLE outer_agg_b ( + host STRING NULL, + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, + PRIMARY KEY(host) +); + +INSERT INTO outer_agg_a VALUES ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +INSERT INTO outer_agg_b VALUES ('h1', 0, 2), ('h2', 0, 4); + +-- `on(host)` leaves only `host`, so `without(host)` groups everything together. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_a / on(host) outer_agg_b); + +-- `device` is not a label of the operand any more, so it groups everything together too. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_a / on(host) outer_agg_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_a / on(host) outer_agg_b); + +DROP TABLE outer_agg_a; + +DROP TABLE outer_agg_b; + +-- Same over the metric engine, whose operands carry `__tsid`. +CREATE TABLE outer_agg_physical ( + ts TIMESTAMP(3) TIME INDEX, + greptime_value DOUBLE, +) ENGINE = metric WITH ("physical_metric_table" = ""); + +CREATE TABLE outer_agg_metric_a ( + host STRING NULL, + device STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host, device), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +CREATE TABLE outer_agg_metric_b ( + host STRING NULL, + ts TIMESTAMP(3) NOT NULL, + greptime_value DOUBLE NULL, + TIME INDEX (ts), + PRIMARY KEY(host), +) +ENGINE = metric +WITH( + on_physical_table = 'outer_agg_physical' +); + +INSERT INTO outer_agg_metric_a (host, device, ts, greptime_value) VALUES + ('h1', 'd1', 0, 10), ('h2', 'd2', 0, 20); + +INSERT INTO outer_agg_metric_b (host, ts, greptime_value) VALUES + ('h1', 0, 2), ('h2', 0, 4); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum without(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(device) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 0, '5s') sum by(host) (outer_agg_metric_a / on(host) outer_agg_metric_b); + +DROP TABLE outer_agg_metric_a; + +DROP TABLE outer_agg_metric_b; + +DROP TABLE outer_agg_physical; diff --git a/tests/cases/standalone/common/tql/matching_filter.result b/tests/cases/standalone/common/tql/matching_filter.result index 8f5bebbe8c3..c0df22a738e 100644 --- a/tests/cases/standalone/common/tql/matching_filter.result +++ b/tests/cases/standalone/common/tql/matching_filter.result @@ -179,24 +179,11 @@ tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{ho | host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | +-------+--------+---------------------+---------------------------------------+ --- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB --- returns the cross product instead (#9209). Recorded here to show the rewrite does not --- change it, not to endorse it. +-- Matching on a subset of the labels: both operands hold several devices per host. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side -- `ignoring(...)`: every label that is not ignored is a matching label. -- SQLNESS SORT_RESULT 3 1 @@ -214,48 +201,19 @@ tql eval(0, 5, '5s') counter_metric / ignoring(missing_label) gauge_metric{host= -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{host="host1"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: found duplicate series for the match group {host="host1"} on the right hand-side of the operation; many-to-many matching not allowed: matching labels must be unique on one side -- An ignored label is not a matching label: the left operand keeps its `eth1` series. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / ignoring(device) gauge_metric{device="eth0"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | -| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- Same for a label left out of `on(...)`. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{device="eth0"}; -+-------+--------+---------------------+---------------------------------------+ -| host | device | ts | counter_metric.val / gauge_metric.val | -+-------+--------+---------------------+---------------------------------------+ -| host1 | eth0 | 1970-01-01T00:00:00 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:00 | 5.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 10.0 | -| host1 | eth0 | 1970-01-01T00:00:05 | 5.0 | -| host2 | eth0 | 1970-01-01T00:00:00 | 20.0 | -| host2 | eth0 | 1970-01-01T00:00:05 | 20.0 | -+-------+--------+---------------------+---------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- Matcher kinds other than equality are enforced by the join just the same. -- SQLNESS SORT_RESULT 3 1 @@ -321,16 +279,7 @@ tql eval(0, 5, '5s') sum by(host)(counter_metric{device="eth0"}) / on(host) sum -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) sum by(host)(gauge_metric{device="eth0"}); -+-------+---------------------+---------------------------------------------------------+ -| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | -+-------+---------------------+---------------------------------------------------------+ -| host1 | 1970-01-01T00:00:00 | 10.0 | -| host1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | 1970-01-01T00:00:05 | 10.0 | -| host1 | 1970-01-01T00:00:05 | 5.0 | -| host2 | 1970-01-01T00:00:00 | 20.0 | -| host2 | 1970-01-01T00:00:05 | 20.0 | -+-------+---------------------+---------------------------------------------------------+ +Error: 3001(EngineExecuteQuery), Execution error: multiple matches for labels {host="host1"}: many-to-one matching must be explicit (group_left/group_right) -- A grouping label that is a value field, not a tag: its value varies between samples of one -- series, so a matcher on it must not filter the scan before sample selection (#9242). @@ -411,14 +360,14 @@ tql eval(0, 5, '5s') counter_metric > on(host, device) gauge_metric{host="host1" -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) group_left sum by(host)(gauge_metric{host="host1"}); -+-------+---------------------+---------------------------------------------------------+ -| host | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | -+-------+---------------------+---------------------------------------------------------+ -| host1 | 1970-01-01T00:00:00 | 2.5 | -| host1 | 1970-01-01T00:00:00 | 5.0 | -| host1 | 1970-01-01T00:00:05 | 2.5 | -| host1 | 1970-01-01T00:00:05 | 5.0 | -+-------+---------------------+---------------------------------------------------------+ ++-------+--------+---------------------+---------------------------------------------------------+ +| host | device | ts | counter_metric.val / gauge_metric.sum(gauge_metric.val) | ++-------+--------+---------------------+---------------------------------------------------------+ +| host1 | eth0 | 1970-01-01T00:00:00 | 2.5 | +| host1 | eth0 | 1970-01-01T00:00:05 | 2.5 | +| host1 | eth1 | 1970-01-01T00:00:00 | 5.0 | +| host1 | eth1 | 1970-01-01T00:00:05 | 5.0 | ++-------+--------+---------------------+---------------------------------------------------------+ -- `absent()` reports the inner selector's equality matchers as labels, so a copied matcher -- must not reach the context an enclosing expression reads. diff --git a/tests/cases/standalone/common/tql/matching_filter.sql b/tests/cases/standalone/common/tql/matching_filter.sql index 72da3da70a0..1b8492c4351 100644 --- a/tests/cases/standalone/common/tql/matching_filter.sql +++ b/tests/cases/standalone/common/tql/matching_filter.sql @@ -73,9 +73,7 @@ tql eval(0, 5, '5s') count_over_time(counter_metric[10s]) / on(host, device) cou -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric offset 5s / on(host, device) gauge_metric{host="host1"}; --- Matching on a subset of the labels. Prometheus rejects this as many-to-many; GreptimeDB --- returns the cross product instead (#9209). Recorded here to show the rewrite does not --- change it, not to endorse it. +-- Matching on a subset of the labels: both operands hold several devices per host. -- SQLNESS SORT_RESULT 3 1 tql eval(0, 5, '5s') counter_metric / on(host) gauge_metric{host="host1"}; From 2bc18c483daac9f41b4ecd619d1897c30a18129f Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 23 Sep 2026 12:25:56 +0000 Subject: [PATCH 023/133] feat: add schema metadata stream to SchemaManager (#9325) Signed-off-by: shuiyisong --- src/common/meta/src/key/schema_name.rs | 127 +++++++++++++++++++++++++ 1 file changed, 127 insertions(+) diff --git a/src/common/meta/src/key/schema_name.rs b/src/common/meta/src/key/schema_name.rs index 689f72ddb32..dcff9c6b338 100644 --- a/src/common/meta/src/key/schema_name.rs +++ b/src/common/meta/src/key/schema_name.rs @@ -334,6 +334,21 @@ impl SchemaManager { Box::pin(stream) } + + /// Returns schema names and values belonging to the target `catalog`. + /// Legacy `null` values are returned as [`SchemaNameValue::default()`]. + pub fn schemas(&self, catalog: &str) -> BoxStream<'static, Result<(String, SchemaNameValue)>> { + let start_key = SchemaNameKey::range_start_key(catalog); + let req = RangeRequest::new().with_prefix(start_key.as_bytes()); + + let stream = PaginationStream::new(self.kv_backend.clone(), req, DEFAULT_PAGE_SIZE, |kv| { + let value = SchemaNameValue::try_from_raw_value(&kv.value)?.unwrap_or_default(); + Ok((schema_decoder(kv)?, value)) + }) + .into_stream(); + + Box::pin(stream) + } } #[derive(Debug, Clone, Hash, Eq, PartialEq, Deserialize, Serialize)] @@ -358,12 +373,124 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use futures::TryStreamExt; use store_api::mito_engine_options::{ TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_TRIGGER_FILE_NUM, }; use super::*; + use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; + use crate::rpc::store::PutRequest; + + #[tokio::test] + async fn test_schemas() { + let manager = SchemaManager::new(Arc::new(MemoryKvBackend::default())); + assert!( + manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + + let mut expected = BTreeMap::new(); + for i in 0..=DEFAULT_PAGE_SIZE { + let name = format!("schema_{i}"); + let value = if i == 0 { + SchemaNameValue::default() + } else { + SchemaNameValue { + ttl: Some(Duration::from_secs(i as u64).into()), + extra_options: BTreeMap::from([("foo".to_string(), i.to_string())]), + create_procedure_id: Some(format!("procedure_{i}")), + } + }; + manager + .create( + SchemaNameKey::new("catalog", &name), + Some(value.clone()), + false, + ) + .await + .unwrap(); + expected.insert(name, value); + } + manager + .create(SchemaNameKey::new("catalog_other", "schema"), None, false) + .await + .unwrap(); + + let schemas = manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap(); + assert_eq!(schemas, expected.into_iter().collect::>()); + assert_eq!( + manager + .schema_names("catalog") + .try_collect::>() + .await + .unwrap(), + schemas + .into_iter() + .map(|(name, _)| name) + .collect::>() + ); + } + + #[tokio::test] + async fn test_schemas_legacy_and_invalid_values() { + let kv_backend = Arc::new(MemoryKvBackend::default()); + let manager = SchemaManager::new(kv_backend.clone()); + let key = SchemaNameKey::new("catalog", "schema").to_bytes(); + + for (raw, expected) in [ + (b"null".as_slice(), SchemaNameValue::default()), + ( + br#"{"ttl":"10s"}"#.as_slice(), + SchemaNameValue { + ttl: Some(Duration::from_secs(10).into()), + ..Default::default() + }, + ), + ] { + kv_backend + .put(PutRequest::new().with_key(key.clone()).with_value(raw)) + .await + .unwrap(); + assert_eq!( + manager + .schemas("catalog") + .try_collect::>() + .await + .unwrap(), + vec![("schema".to_string(), expected)] + ); + } + + kv_backend + .put(PutRequest::new().with_key(key).with_value(b"invalid")) + .await + .unwrap(); + assert!( + manager + .schemas("catalog") + .try_collect::>() + .await + .is_err() + ); + assert_eq!( + manager + .schema_names("catalog") + .try_collect::>() + .await + .unwrap(), + vec!["schema".to_string()] + ); + } #[test] fn test_display_schema_value() { From 4c98fa52651299c9a5679379289a7148579e0da6 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 23 Sep 2026 12:46:31 +0000 Subject: [PATCH 024/133] ci: add optional AWS runners for observability benchmarks (#9322) * ci: add optional AWS runners for observability benchmarks Signed-off-by: WenyXu * ci: restore automatic benchmark disk sizing Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .github/actions/setup-benchmark/action.yml | 38 ++ .../scripts/aliyun-ecs-runner-provision.py | 68 +-- .github/scripts/aliyun-ecs-runner-teardown.py | 16 +- .github/scripts/aws-ec2-runner-provision.py | 265 ++++++++++++ .github/scripts/aws-ec2-runner-teardown.py | 160 +++++++ .github/scripts/ci-runner-config.py | 47 ++ .github/scripts/ci-runners.md | 73 ++++ .github/scripts/runner_utils.py | 101 +++++ .github/workflows/agent-observability.yml | 91 +++- .github/workflows/aws-ec2-runner-janitor.yml | 29 ++ .github/workflows/checks.yml | 1 + .github/workflows/vmbench-long-range.yml | 92 +++- tests/perf/test_aliyun_ecs_runner_scripts.py | 23 +- tests/perf/test_ci_runner_scripts.py | 406 ++++++++++++++++++ 14 files changed, 1296 insertions(+), 114 deletions(-) create mode 100644 .github/actions/setup-benchmark/action.yml create mode 100644 .github/scripts/aws-ec2-runner-provision.py create mode 100644 .github/scripts/aws-ec2-runner-teardown.py create mode 100644 .github/scripts/ci-runner-config.py create mode 100644 .github/scripts/ci-runners.md create mode 100644 .github/scripts/runner_utils.py create mode 100644 .github/workflows/aws-ec2-runner-janitor.yml create mode 100644 tests/perf/test_ci_runner_scripts.py diff --git a/.github/actions/setup-benchmark/action.yml b/.github/actions/setup-benchmark/action.yml new file mode 100644 index 00000000000..a446a78a2c3 --- /dev/null +++ b/.github/actions/setup-benchmark/action.yml @@ -0,0 +1,38 @@ +name: Set up benchmark dependencies +description: Reuse installed tools and install missing dependencies on an Ubuntu self-hosted runner +runs: + using: composite + steps: + - name: Check host and install missing tools + shell: bash + run: | + set -euo pipefail + [[ "$(uname -m)" == x86_64 ]] + command -v git + missing=() + for tool in curl jq python3 rsync; do + command -v "$tool" >/dev/null || missing+=("$tool") + done + if ((${#missing[@]})); then + sudo -n apt-get update + sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get install -y --no-install-recommends "${missing[@]}" + fi + - name: Install Docker if missing and verify access + shell: bash + run: | + set -euo pipefail + if ! command -v docker >/dev/null; then + . /etc/os-release + [[ "$ID" == ubuntu ]] + sudo -n install -m 0755 -d /etc/apt/keyrings + curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo -n tee /etc/apt/keyrings/docker.asc >/dev/null + sudo -n chmod a+r /etc/apt/keyrings/docker.asc + echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu $VERSION_CODENAME stable" | sudo -n tee /etc/apt/sources.list.d/docker.list >/dev/null + sudo -n apt-get update + sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get install -y --no-install-recommends docker-ce docker-ce-cli containerd.io docker-buildx-plugin + fi + if ! docker info >/dev/null 2>&1; then + sudo -n systemctl enable --now docker + fi + # Membership must be configured before the runner starts, not in this job. + docker info diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py index 1a26776b4cd..46efbffd3ec 100644 --- a/.github/scripts/aliyun-ecs-runner-provision.py +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -49,10 +49,10 @@ import base64 import json import os import time -import urllib.error -import urllib.request from dataclasses import dataclass +from runner_utils import append_github_output, create_registration_token, find_runner_by_name, github_api + RUNNER_NAME_PREFIX = "qreg-ecs" RUNNER_LABEL_PREFIX = "query-regression-ecs" MANAGED_BY_TAG_KEY = "managed-by" @@ -243,63 +243,6 @@ def encode_user_data(script: str) -> str: return base64.b64encode(script.encode("utf-8")).decode("ascii") -def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: - data = json.dumps(body).encode("utf-8") if body is not None else None - request = urllib.request.Request( - f"https://api.github.com{path}", - data=data, - method=method, - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "X-GitHub-Api-Version": "2022-11-28", - }, - ) - try: - with urllib.request.urlopen(request, timeout=30) as response: - if response.status == 204: - return {} - return json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as error: - # GitHub's error body says exactly why (e.g. "Must have admin rights to - # Repository" for a PAT without the required scope); surface it instead - # of a bare "HTTP Error 403". - body = error.read().decode("utf-8", "replace") - raise SystemExit( - f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" - "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " - "'repo' scope (classic PAT) or 'Administration: write' on the " - "repository (fine-grained PAT)." - ) from error - - -def create_registration_token(github_token: str, repo: str) -> str: - response = github_api( - github_token, - "POST", - f"/repos/{repo}/actions/runners/registration-token", - body={}, - ) - return response["token"] - - -def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: - page = 1 - while True: - response = github_api( - github_token, - "GET", - f"/repos/{repo}/actions/runners?per_page=100&page={page}", - ) - runners = response.get("runners", []) - for runner in runners: - if runner.get("name") == name: - return runner - if len(runners) < 100: - return None - page += 1 - - def make_ecs_client(config: ProvisionConfig): from alibabacloud_ecs20140526.client import Client as EcsClient from alibabacloud_tea_openapi.models import Config as OpenApiConfig @@ -388,13 +331,6 @@ def dump_console_output(client, region_id: str, instance_id: str) -> None: summary.write("\n```\n\n") -def append_github_output(name: str, value: str) -> None: - output_path = os.environ.get("GITHUB_OUTPUT") - if output_path: - with open(output_path, "a", encoding="utf-8") as output: - output.write(f"{name}={value}\n") - - def run_instance(client, config: ProvisionConfig, user_data: str) -> str: from alibabacloud_ecs20140526 import models as ecs_models diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py index 658633e0395..d5b30b2922d 100644 --- a/.github/scripts/aliyun-ecs-runner-teardown.py +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -47,6 +47,8 @@ import sys from datetime import datetime, timedelta, timezone from pathlib import Path +from runner_utils import deregister_runner + # Reuse the GitHub API client, runner lookup, and tag constants from the # provision script. _PROVISION_SPEC = importlib.util.spec_from_file_location( @@ -159,20 +161,6 @@ def delete_instance(client, instance_id: str, region_id: str | None = None) -> b return False -def deregister_runner(token: str, repo: str, runner_name: str) -> bool: - runner = provision.find_runner_by_name(token, repo, runner_name) - if runner is None: - print(f"Runner {runner_name} is not registered", flush=True) - return True - try: - provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") - print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) - return True - except Exception as error: # noqa: BLE001 - print(f"Failed to deregister runner {runner_name}: {error}", flush=True) - return False - - def list_managed_instances( client, region_id: str ) -> list[tuple[str, str, str, str | None]]: diff --git a/.github/scripts/aws-ec2-runner-provision.py b/.github/scripts/aws-ec2-runner-provision.py new file mode 100644 index 00000000000..7bf1b4091c9 --- /dev/null +++ b/.github/scripts/aws-ec2-runner-provision.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Provision an ephemeral EC2 CI runner from an Ubuntu x86_64 AMI.""" + +from __future__ import annotations + +import argparse +import os +import re +import shlex +import time +from dataclasses import dataclass + +from runner_utils import append_github_output, create_registration_token, find_runner_by_name, github_api + +MANAGED_BY = "query-regression-ci" +REPO_TAG = "github-repository" +RUN_TAG = "query-regression-run-id" +RUNNER_TAG = "github-runner-name" +TTL_TAG_KEY = "runner-ttl-hours" + + +def make_client(region: str): + import boto3 + from botocore.config import Config + + return boto3.client("ec2", region_name=region, config=Config( + connect_timeout=10, read_timeout=30, + retries={"mode": "standard", "max_attempts": 3}, + )) + + +@dataclass(frozen=True) +class ProvisionConfig: + region: str + image_id: str + instance_type: str + subnet_id: str + security_group_id: str + repo: str + run_id: str + github_token: str + system_disk_gib: int = 500 + ttl_hours: int = 8 + + def __post_init__(self): + if not re.fullmatch(r"[0-9]+-[0-9]+", self.run_id): + raise ValueError("run-id must be GitHub run_id-run_attempt") + if not 20 <= self.system_disk_gib <= 2048: + raise ValueError("system-disk-gib must be between 20 and 2048") + if not 1 <= self.ttl_hours <= 168: + raise ValueError("ttl-hours must be between 1 and 168") + + @property + def runner_name(self) -> str: + return f"ci-ec2-{self.run_id}" + + +def resolve_image(client, image_id: str) -> str: + if image_id != "auto": + return image_id + images = client.describe_images( + Owners=["099720109477"], + Filters=[ + {"Name": "name", "Values": ["ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"]}, + {"Name": "state", "Values": ["available"]}, + {"Name": "architecture", "Values": ["x86_64"]}, + ], + )["Images"] + if not images: + raise ValueError("No Canonical Ubuntu 24.04 amd64 AMI found in this region") + image = max(images, key=lambda image: image["CreationDate"]) + print(f"Using AMI {image['ImageId']}: {image['Name']}") + return image["ImageId"] + + +def root_disk(client, image_id: str, size: int) -> dict: + images = client.describe_images(ImageIds=[image_id])["Images"] + if len(images) != 1: + raise ValueError(f"AMI not found: {image_id}") + image = images[0] + if image["State"] != "available" or image["Architecture"] != "x86_64": + raise ValueError("Runner AMI must be available and x86_64") + root = next((d for d in image["BlockDeviceMappings"] + if d["DeviceName"] == image["RootDeviceName"] and "Ebs" in d), None) + if root is None: + raise ValueError("Runner AMI must have an EBS root disk") + minimum = root["Ebs"]["VolumeSize"] + if size < minimum: + raise ValueError(f"Requested root disk {size} GiB is smaller than AMI minimum {minimum} GiB") + return {"DeviceName": image["RootDeviceName"], "Ebs": { + "VolumeSize": size, "VolumeType": "gp3", "Iops": 3000, + "Throughput": 125, "DeleteOnTermination": True, + }} + + +def runner_download(token: str, repo: str) -> str: + runners = github_api(token, "GET", f"/repos/{repo}/actions/runners/downloads") + return next(r["download_url"] for r in runners if r["os"] == "linux" and r["architecture"] == "x64") + + +def render_user_data(name: str, token: str, repo: str, download_url: str) -> str: + # Quote every external value before embedding it in the root shell script. + name, token, repo, download_url = map(shlex.quote, (name, token, repo, download_url)) + return f"""#!/bin/bash +set -euo pipefail +for tool in git curl tar sudo; do command -v "$tool"; done +# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "/swapfile" ]]; then + fallocate --length 16G "/swapfile" + chmod 600 "/swapfile" + mkswap "/swapfile" +fi +swapon "/swapfile" +sysctl --write vm.swappiness=10 +swapon --show +free --human + +id runner >/dev/null 2>&1 || useradd --create-home --shell /bin/bash runner +# Establish membership before the runner service starts; setup installs Docker later. +getent group docker >/dev/null || groupadd --system docker +usermod -aG docker runner +printf 'runner ALL=(ALL) NOPASSWD: ALL\\n' > /etc/sudoers.d/ci-runner +chmod 0440 /etc/sudoers.d/ci-runner +visudo -cf /etc/sudoers.d/ci-runner +install -d -o runner -g "$(id -gn runner)" /home/runner/benchmark-data /opt/ci-github-runner +cd /opt/ci-github-runner +curl --fail --location --retry 3 {download_url} | tar xz +./bin/installdependencies.sh +chown -R runner:"$(id -gn runner)" /opt/ci-github-runner +runuser -u runner -- ./config.sh --unattended --ephemeral --disableupdate --url https://github.com/{repo} --token {token} --name {name} --labels {name} +./svc.sh install runner +# Match the ECS runner's behavior when a child process is OOM-killed. +service_name=$(cat .service) +mkdir -p "/etc/systemd/system/$service_name.d" +cat > "/etc/systemd/system/$service_name.d/oom.conf" <<'EOF' +[Service] +OOMPolicy=continue +EOF +systemctl daemon-reload +# Package installation in a job must not restart its runner service. +install -d /etc/needrestart/conf.d +cat > /etc/needrestart/conf.d/actions_runner_services.conf <<'EOF' +$nrconf{{override_rc}}{{qr(^actions\\.runner\\..+\\.service$)}} = 0; +EOF +./svc.sh start +""" + + +def run_instance(client, config: ProvisionConfig) -> str: + image_id = resolve_image(client, config.image_id) + disk = root_disk(client, image_id, config.system_disk_gib) + instance = client.describe_instance_types(InstanceTypes=[config.instance_type])["InstanceTypes"][0] + if "x86_64" not in instance["ProcessorInfo"]["SupportedArchitectures"]: + raise ValueError("Runner instance type must support x86_64") + token = create_registration_token(config.github_token, config.repo) + print(f"::add-mask::{token}") + script = render_user_data(config.runner_name, token, config.repo, runner_download(config.github_token, config.repo)) + tags = [ + {"Key": "Name", "Value": config.runner_name}, + {"Key": "managed-by", "Value": MANAGED_BY}, + {"Key": REPO_TAG, "Value": config.repo}, + {"Key": RUN_TAG, "Value": config.run_id}, + {"Key": RUNNER_TAG, "Value": config.runner_name}, + {"Key": TTL_TAG_KEY, "Value": str(config.ttl_hours)}, + ] + response = client.run_instances( + ImageId=image_id, InstanceType=config.instance_type, + MinCount=1, MaxCount=1, SubnetId=config.subnet_id, + SecurityGroupIds=[config.security_group_id], + BlockDeviceMappings=[disk], UserData=script, + ClientToken=config.runner_name, + MetadataOptions={"HttpTokens": "required", "HttpEndpoint": "enabled"}, + TagSpecifications=[{"ResourceType": kind, "Tags": tags} for kind in ("instance", "volume")], + ) + return response["Instances"][0]["InstanceId"] + + +def dump_console(client, instance_id: str) -> None: + try: + # Boto3 decodes the API's base64 console output. + output = client.get_console_output(InstanceId=instance_id, Latest=True).get("Output", "") + print(f"::group::EC2 console {instance_id}") + print("\n".join(output.splitlines()[-80:])) + print("::endgroup::") + except Exception as error: + print(f"Unable to read console for {instance_id}: {error}") + + +def provision(client, config: ProvisionConfig) -> int: + instance_id = run_instance(client, config) + # Publish before waiting: failed registration must still reach teardown. + for key, value in (("instance_id", instance_id), + ("label", config.runner_name), ("runner_name", config.runner_name)): + append_github_output(key, value) + try: + client.get_waiter("instance_running").wait( + InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60}) + deadline = time.monotonic() + 10 * 60 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, config.runner_name) + if runner and runner.get("status") == "online": + print(f"Runner {config.runner_name} online: {instance_id}") + return 0 + time.sleep(5) + raise TimeoutError(f"Runner {config.runner_name} did not come online") + finally: + dump_console(client, instance_id) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + for name, env in (("region", "AWS_REGION"), ("repo", "GITHUB_REPOSITORY"), + ("github-token", "GH_PERSONAL_ACCESS_TOKEN"), + ("run-id", "AWS_EC2_RUN_ID"), ("image-id", "AWS_EC2_IMAGE_ID"), + ("instance-type", "AWS_EC2_INSTANCE_TYPE"), + ("subnet-id", "AWS_EC2_SUBNET_ID"), + ("security-group-id", "AWS_EC2_SECURITY_GROUP_ID")): + parser.add_argument(f"--{name}", default=os.environ.get(env) or ("auto" if name == "image-id" else None)) + parser.add_argument("--system-disk-gib", type=int, default=os.environ.get("AWS_EC2_SYSTEM_DISK_GIB", "500")) + parser.add_argument("--ttl-hours", type=int, default=os.environ.get("AWS_EC2_TTL_HOURS", "8")) + args = parser.parse_args() + for name, value in vars(args).items(): + if value is None or value == "": + parser.error(f"Missing {name.replace('_', '-')}") + config = ProvisionConfig(**vars(args)) + return provision(make_client(args.region), config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aws-ec2-runner-teardown.py b/.github/scripts/aws-ec2-runner-teardown.py new file mode 100644 index 00000000000..fbfba4f639a --- /dev/null +++ b/.github/scripts/aws-ec2-runner-teardown.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# /// script +# requires-python = ">=3.10" +# dependencies = ["boto3>=1.35,<2"] +# /// + +"""Delete managed EC2 CI runners, or sweep expired instances by ownership tags.""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +import time +from pathlib import Path +from datetime import datetime, timedelta, timezone + +from runner_utils import deregister_runner + +_SPEC = importlib.util.spec_from_file_location("aws_ec2_runner_provision", Path(__file__).with_name("aws-ec2-runner-provision.py")) +assert _SPEC is not None and _SPEC.loader is not None +provision = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = provision +_SPEC.loader.exec_module(provision) +make_client = provision.make_client +dump_console = provision.dump_console +MANAGED_BY = provision.MANAGED_BY +REPO_TAG = provision.REPO_TAG +RUN_TAG = provision.RUN_TAG +RUNNER_TAG = provision.RUNNER_TAG +TTL_TAG_KEY = provision.TTL_TAG_KEY +SWEEP_TTL_HOURS = 4 + + +def managed_instances(client, repo: str, run_id: str | None = None, + instance_id: str | None = None) -> list[dict]: + filters = [{"Name": "tag:managed-by", "Values": [MANAGED_BY]}, + {"Name": f"tag:{REPO_TAG}", "Values": [repo]}] + if instance_id: + # A retried teardown may have a newer workflow attempt than creation. + filters.append({"Name": "instance-id", "Values": [instance_id]}) + elif run_id: + filters.append({"Name": f"tag:{RUN_TAG}", "Values": [run_id]}) + return [instance + for page in client.get_paginator("describe_instances").paginate(Filters=filters) + for reservation in page["Reservations"] for instance in reservation["Instances"]] + + +def delete_instance(client, instance_id: str) -> bool: + dump_console(client, instance_id) + for attempt in range(12): + try: + client.terminate_instances(InstanceIds=[instance_id]) + client.get_waiter("instance_terminated").wait( + InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60}) + return True + except Exception as error: + code = getattr(error, "response", {}).get("Error", {}).get("Code") + if code == "InvalidInstanceID.NotFound": + if attempt < 11: + time.sleep(5) + continue + # The ownership lookup observed this instance. Treat sustained + # NotFound after bounded retries as an already-deleted VM. + return True + print(f"Failed to terminate {instance_id}: {error}") + return False + return False + + +def teardown(client, repo: str, token: str, run_id: str | None = None, + instance_id: str | None = None, now: datetime | None = None, + runner_name: str | None = None) -> int: + # Lookup by ownership/run tags also recovers a create whose response or job + # outputs were lost. Sweep never touches legacy build/JSONBench instances. + instances = managed_instances(client, repo, run_id, instance_id) + # Tags can lag RunInstances. Retry before concluding a targeted VM is gone. + if run_id and not instances: + for _ in range(11): + time.sleep(5) + instances = managed_instances(client, repo, run_id, instance_id) + if instances: + break + ok = True + found = False + for instance in instances: + identity = instance["InstanceId"] + if instance_id and identity != instance_id: + continue + found = True + tags = {t["Key"]: t["Value"] for t in instance.get("Tags", [])} + name = tags.get(RUNNER_TAG, "") + if not name.startswith("ci-ec2-"): + print(f"Skipping {identity}: missing managed runner name") + ok = False + continue + if not run_id: + ttl = tags.get(TTL_TAG_KEY) + try: + hours = SWEEP_TTL_HOURS if ttl is None else int(ttl) + if not 1 <= hours <= 168: + raise ValueError("TTL outside 1..168 hours") + except (TypeError, ValueError): + print(f"Skipping {identity}: invalid TTL {ttl!r}") + continue + if (now or datetime.now(timezone.utc)) - instance["LaunchTime"] < timedelta(hours=hours): + continue + # A GitHub outage must not prevent EC2 deletion, or vice versa. + ok &= delete_instance(client, identity) + try: + ok &= deregister_runner(token, repo, name) + except (Exception, SystemExit) as error: + print(f"Failed to unregister {name}: {error}") + ok = False + if run_id and not found: + # EC2 may already be gone while its GitHub registration remains. + name = runner_name or f"ci-ec2-{run_id}" + if not name.startswith("ci-ec2-"): + raise ValueError("Invalid CI runner name") + ok &= deregister_runner(token, repo, name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + for name, env in (("region", "AWS_REGION"), ("repo", "GITHUB_REPOSITORY"), + ("github-token", "GH_PERSONAL_ACCESS_TOKEN"), + ("run-id", "AWS_EC2_RUN_ID"), ("instance-id", "AWS_EC2_INSTANCE_ID"), + ("runner-name", "AWS_EC2_RUNNER_NAME")): + parser.add_argument(f"--{name}", default=os.environ.get(env)) + parser.add_argument("--sweep", action="store_true") + args = parser.parse_args() + for name in ("region", "repo", "github_token"): + if not getattr(args, name): + parser.error(f"Missing {name.replace('_', '-')}") + if not args.sweep and not args.run_id: + parser.error("Missing run-id for targeted teardown") + return teardown(make_client(args.region), args.repo, args.github_token, + None if args.sweep else args.run_id, + None if args.sweep else args.instance_id, + runner_name=None if args.sweep else args.runner_name) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/ci-runner-config.py b/.github/scripts/ci-runner-config.py new file mode 100644 index 00000000000..a242b263c3e --- /dev/null +++ b/.github/scripts/ci-runner-config.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve dispatch provider and instance type before allocating a runner.""" + +import argparse +import re + +from runner_utils import append_github_output + +DEFAULT_INSTANCE_TYPES = {"Aliyun": "ecs.c9i.2xlarge", "AWS": "c7i.2xlarge"} + + +def resolve_instance_type(provider: str, requested: str) -> str: + if provider not in DEFAULT_INSTANCE_TYPES: + raise ValueError("provider must be Aliyun or AWS") + instance_type = DEFAULT_INSTANCE_TYPES[provider] if requested == "auto" else requested + pattern = r"ecs\.[a-z0-9-]+\.[a-z0-9-]+" if provider == "Aliyun" else r"[a-z0-9-]+\.[a-z0-9-]+" + if not re.fullmatch(pattern, instance_type): + raise ValueError(f"Invalid {provider} instance type: {instance_type!r}") + return instance_type + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--provider", required=True) + parser.add_argument("--instance-type", default="auto") + args = parser.parse_args() + instance_type = resolve_instance_type(args.provider, args.instance_type) + append_github_output("instance_type", instance_type) + print(f"CI runner: {args.provider} / {instance_type}") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/ci-runners.md b/.github/scripts/ci-runners.md new file mode 100644 index 00000000000..6e4b2b3a4f7 --- /dev/null +++ b/.github/scripts/ci-runners.md @@ -0,0 +1,73 @@ +# CI runners + +The agent observability and long-range benchmark workflows select `Aliyun` or +`AWS` through the `provider` dispatch input. `instance_type: auto` selects +`ecs.c9i.2xlarge` or `c7i.2xlarge`, respectively (8 vCPU, 16 GiB). +Existing build/release and JSONBench workflows keep their current runner actions. + +Aliyun uses the existing ECS actions and image. Its provision and teardown +scripts share only GitHub API and output helpers from `runner_utils.py` with +the new AWS scripts. + +AWS uses `aws-ec2-runner-provision.py` and `aws-ec2-runner-teardown.py` directly. +Repository variables supply `EC2_RUNNER_REGION`, `EC2_RUNNER_SUBNET_ID`, and +`EC2_RUNNER_SECURITY_GROUP_ID`; credentials use `AWS_ACCESS_KEY_ID`, +`AWS_SECRET_ACCESS_KEY`, and `GH_PERSONAL_ACCESS_TOKEN` secrets. +The current region is Singapore (`ap-southeast-1`). By default the script selects +the newest available Canonical Ubuntu 24.04 amd64 server AMI in that region +(owner `099720109477`) and logs its ID. Set `BENCHMARK_EC2_IMAGE_ID` to pin a +prepared Ubuntu AMI. Existing build CI's `EC2_RUNNER_LINUX_AMD64_IMAGE_ID` is +unaffected. Both workflows default `system_disk_gib` to `auto` (also accepting +`0`): agent observability uses S=80, P=100, M=500 GiB, and long-range uses 80 GiB. +These defaults are the same for both providers. Explicit values from 20 to 2048 +GiB override them and must be at least the selected AMI's root snapshot size. +The official 24.04 image does not have the old custom AMI's 500 GiB minimum. + +AWS user data registers a fresh ephemeral runner. Git, curl, tar, and sudo are +expected in the Ubuntu server image. Docker group membership is set before the +runner service starts, and this dedicated VM grants the runner passwordless sudo +so its job can install dependencies. Like ECS, the host creates 16 GiB swap with +swappiness 10, disables systemd-oomd and background automatic updates, and sets +`OOMPolicy=continue` on the runner service. Explicit setup package installation +remains enabled. + +Both benchmark jobs call `.github/actions/setup-benchmark` after checkout. It +reuses existing tools, installs missing curl/jq/python3/rsync, installs Docker CE +if absent, and checks Docker access. Missing packages require passwordless sudo; +Aliyun's prepared image normally provides the tools and Docker already. Existing +Docker installations are preserved, so exact versions need not match between +providers. Provisioning must establish Docker permissions before starting the +runner; adding a group inside the running job would not update its groups. + +The subnet must provide outbound access to Ubuntu and Docker packages, GitHub, +and the configured container registries. The IAM identity needs EC2 +image/instance-type discovery, instance and volume tagging, launch, describe, +console output, and termination permissions. GitHub credentials need repository +runner administration access. + +Creation publishes instance identity before waiting for registration. Teardown +can recover missing outputs using repository and workflow-attempt tags. The +daily AWS janitor (03:23 UTC) removes expired instances tagged `managed-by=query-regression-ci` +and unregisters their runners. Both providers use the existing +`managed-by=query-regression-ci`, `query-regression-run-id`, and +`runner-ttl-hours` tags. AWS retains a run-attempt suffix in its run identity to +isolate retries, plus repository and runner-name tags to restrict cleanup. +Legacy build instances do not match these ownership filters. +Both janitors use a 4-hour fallback when the TTL tag is absent and skip invalid +TTL tags. Benchmark jobs explicitly tag their configured TTL (default 8 hours). +Cleanup does not check whether the GitHub job is still running. Expired instances +are removed on the next daily sweep; normal workflow teardown runs immediately. +TTL is measured from instance launch, so it must cover provisioning and the +benchmark. The scheduled janitor becomes active after landing on the default +branch. + +Local validation: + +```sh +python3 tests/perf/test_ci_runner_scripts.py +python3 tests/perf/test_aliyun_ecs_runner_scripts.py +actionlint .github/workflows/agent-observability.yml .github/workflows/vmbench-long-range.yml .github/workflows/aws-ec2-runner-janitor.yml +``` + +A real AWS dispatch is still needed to verify AMI boot, package +installation, registry connectivity, and end-to-end benchmark execution. diff --git a/.github/scripts/runner_utils.py b/.github/scripts/runner_utils.py new file mode 100644 index 00000000000..afd827f62cf --- /dev/null +++ b/.github/scripts/runner_utils.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small shared GitHub runner API and workflow-output helpers.""" + +from __future__ import annotations + +import json +import os +import urllib.error +import urllib.request + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + if response.status == 204: + return {} + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False diff --git a/.github/workflows/agent-observability.yml b/.github/workflows/agent-observability.yml index 02c4ccb81bc..86713c911ed 100644 --- a/.github/workflows/agent-observability.yml +++ b/.github/workflows/agent-observability.yml @@ -50,15 +50,21 @@ on: type: string default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260917-8111034e@sha256:ce9dee2abb1e70f907a09ea01854904a91cc97748b0c4ef6f33d9ec34044bec0 required: true - ecs_instance_type: - description: ECS instance type (independent of Query Regression) + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' type: string - default: ecs.c9i.2xlarge + default: auto required: true system_disk_gib: - description: System disk GiB; 0 = auto (S 80, P 100, M 500), or 20..2048 - type: number - default: 0 + description: 'System disk GiB: auto (or 0) = S 80, P 100, M 500; enter 20..2048 to override' + type: string + default: auto required: true benchmark_timeout_minutes: description: Benchmark job timeout in minutes (1..360) @@ -89,7 +95,8 @@ env: VICTORIALOGS_TAG: ${{ inputs.victorialogs_tag }} O11YBENCH_REF: ${{ inputs.o11ybench_ref }} RUNTIME_IMAGE: ${{ inputs.runtime_image }} - ECS_INSTANCE_TYPE: ${{ inputs.ecs_instance_type }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} @@ -100,6 +107,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} targets: ${{ steps.inputs.outputs.targets }} system_disk_gib: ${{ steps.inputs.outputs.system_disk_gib }} greptimedb: ${{ steps.inputs.outputs.greptimedb }} @@ -110,9 +118,13 @@ jobs: with: path: greptimedb persist-credentials: false - - name: Test ECS lifecycle helpers + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Test CI runner helpers run: | python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + python3 greptimedb/tests/perf/test_ci_runner_scripts.py python3 greptimedb/tests/perf/test_agent_observability_summary.py - name: Validate requested benchmark id: inputs @@ -124,14 +136,15 @@ jobs: run: | set -euo pipefail # Validate cloud lifecycle budgets before provisioning; runtime owns workload validation. - [[ -n "$ECS_INSTANCE_TYPE" ]] case "$PROFILE" in S) auto_disk_gib=80 ;; P) auto_disk_gib=100 ;; M) auto_disk_gib=500 ;; *) echo "Unknown dataset: $PROFILE" >&2; exit 1 ;; esac - [[ "$SYSTEM_DISK_GIB" != 0 ]] || SYSTEM_DISK_GIB=$auto_disk_gib + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + SYSTEM_DISK_GIB=$auto_disk_gib + fi for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers' >&2; exit 1; } done @@ -160,15 +173,16 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 45 outputs: - label: ${{ steps.ecs.outputs.label }} - instance_id: ${{ steps.ecs.outputs.instance_id }} - runner_name: ${{ steps.ecs.outputs.runner_name }} + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} steps: - uses: actions/checkout@v4 with: path: greptimedb persist-credentials: false - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' id: ecs uses: ./greptimedb/.github/actions/aliyun-ecs-create with: @@ -179,13 +193,35 @@ jobs: vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} - instance-type: ${{ env.ECS_INSTANCE_TYPE }} + instance-type: ${{ needs.validate.outputs.instance_type }} enable-docker: 'true' system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} ttl-hours: ${{ inputs.janitor_ttl_hours }} resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py benchmark: needs: - validate @@ -193,6 +229,7 @@ jobs: runs-on: ${{ needs.provision.outputs.label }} timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: + INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} TARGETS: ${{ needs.validate.outputs.targets }} SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} OWNER: o11ybench-${{ github.run_id }}-${{ github.run_attempt }} @@ -209,6 +246,8 @@ jobs: with: path: greptimedb persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark - name: Check Docker on ephemeral Linux host shell: bash run: | @@ -252,9 +291,9 @@ jobs: done jq -n --arg sha "$actual" --arg targets "$TARGETS" --arg profile "$PROFILE" \ --arg cpus "$DB_CPUS" --arg memory "$DB_MEMORY" --argjson images "$images" \ - --arg instance_type "$ECS_INSTANCE_TYPE" --argjson disk_gib "$SYSTEM_DISK_GIB" \ + --arg provider "$PROVIDER" --arg instance_type "$INSTANCE_TYPE" --argjson disk_gib "$SYSTEM_DISK_GIB" \ --argjson timeout_minutes "$BENCHMARK_TIMEOUT_MINUTES" --argjson ttl_hours "$JANITOR_TTL_HOURS" \ - '{o11ybench_sha:$sha, targets:$targets, profile:$profile, db_cpus:$cpus, db_memory:$memory, images:$images, ecs:{instance_type:$instance_type, system_disk_gib:$disk_gib, benchmark_timeout_minutes:$timeout_minutes, janitor_ttl_hours:$ttl_hours}}' > "$MANIFEST_PATH" + '{o11ybench_sha:$sha, targets:$targets, profile:$profile, db_cpus:$cpus, db_memory:$memory, images:$images, ($provider | if . == "AWS" then "ec2" else "ecs" end):{instance_type:$instance_type, system_disk_gib:$disk_gib, benchmark_timeout_minutes:$timeout_minutes, janitor_ttl_hours:$ttl_hours}}' > "$MANIFEST_PATH" - name: Checkout benchmark tools uses: actions/checkout@v4 with: @@ -343,7 +382,7 @@ jobs: needs: - provision - benchmark - if: ${{ always() && needs.provision.outputs.instance_id != '' }} + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -352,6 +391,7 @@ jobs: path: greptimedb persist-credentials: false - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' uses: ./greptimedb/.github/actions/aliyun-ecs-delete with: access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} @@ -360,3 +400,20 @@ jobs: region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} instance-id: ${{ needs.provision.outputs.instance_id }} runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/.github/workflows/aws-ec2-runner-janitor.yml b/.github/workflows/aws-ec2-runner-janitor.yml new file mode 100644 index 00000000000..1c4435022b4 --- /dev/null +++ b/.github/workflows/aws-ec2-runner-janitor.yml @@ -0,0 +1,29 @@ +name: AWS EC2 Runner Janitor + +on: + schedule: + - cron: '23 3 * * *' + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@v4 + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + - name: Sweep expired CI instances and unregister runners + env: + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: uv run .github/scripts/aws-ec2-runner-teardown.py --sweep diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 561ee2c7473..d3bad06f41d 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -97,6 +97,7 @@ jobs: python3 tests/perf/test_query_regression_nightly_refs.py python3 tests/perf/test_query_regression_slash.py python3 tests/perf/test_aliyun_ecs_runner_scripts.py + python3 tests/perf/test_ci_runner_scripts.py python3 tests/perf/test_agent_observability_summary.py check: diff --git a/.github/workflows/vmbench-long-range.yml b/.github/workflows/vmbench-long-range.yml index 9d0088217ad..357d2b21c69 100644 --- a/.github/workflows/vmbench-long-range.yml +++ b/.github/workflows/vmbench-long-range.yml @@ -35,10 +35,16 @@ on: type: string default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260918-545dc570@sha256:0ff03e7e690575e875bb85b3054a03486dd9359073546dc715da177153cf7a29 required: true - ecs_instance_type: - description: ECS instance type (independent of Query Regression) + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' type: string - default: ecs.c9i.2xlarge + default: auto required: true benchmark_root: description: Absolute disk-backed work directory (outside the o11ybench checkout) @@ -46,9 +52,9 @@ on: default: /home/runner/benchmark-data required: true system_disk_gib: - description: System disk GiB (20..2048) - type: number - default: 80 + description: 'System disk GiB: auto (or 0) = 80; enter 20..2048 to override' + type: string + default: auto required: true benchmark_timeout_minutes: description: Benchmark job timeout in minutes (1..360) @@ -78,7 +84,8 @@ env: GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} VICTORIAMETRICS_TAG: ${{ inputs.victoriametrics_tag }} RUNTIME_IMAGE: ${{ inputs.runtime_image }} - ECS_INSTANCE_TYPE: ${{ inputs.ecs_instance_type }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} @@ -89,6 +96,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 15 outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} system_disk_gib: ${{ steps.plan.outputs.system_disk_gib }} runtime_image: ${{ steps.plan.outputs.runtime_image }} o11ybench_ref: ${{ steps.plan.outputs.o11ybench_ref }} @@ -99,8 +107,13 @@ jobs: with: path: greptimedb persist-credentials: false - - name: Test existing ECS lifecycle helpers - run: python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Test CI runner helpers + run: | + python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py + python3 greptimedb/tests/perf/test_ci_runner_scripts.py - name: Validate budgets and build warm plan id: plan env: @@ -108,11 +121,13 @@ jobs: RUN_VICTORIAMETRICS: ${{ inputs.run_victoriametrics }} run: | set -euo pipefail - [[ -n "$ECS_INSTANCE_TYPE" ]] [[ "$BENCHMARK_ROOT" == /* && "$BENCHMARK_ROOT" != / && "$BENCHMARK_ROOT" != */ ]] [[ "$BENCHMARK_ROOT" != /tmp && "$BENCHMARK_ROOT" != /tmp/* ]] [[ "$BENCHMARK_ROOT" != *$'\n'* && "$BENCHMARK_ROOT" != *$'\r'* ]] [[ "/${BENCHMARK_ROOT#/}/" != */../* && "/${BENCHMARK_ROOT#/}/" != */./* ]] + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + SYSTEM_DISK_GIB=80 + fi for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers'; exit 1; } done @@ -159,15 +174,16 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 45 outputs: - label: ${{ steps.ecs.outputs.label }} - instance_id: ${{ steps.ecs.outputs.instance_id }} - runner_name: ${{ steps.ecs.outputs.runner_name }} + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} steps: - uses: actions/checkout@v4 with: path: greptimedb persist-credentials: false - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' id: ecs uses: ./greptimedb/.github/actions/aliyun-ecs-create with: @@ -178,13 +194,35 @@ jobs: vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} - instance-type: ${{ env.ECS_INSTANCE_TYPE }} + instance-type: ${{ needs.validate.outputs.instance_type }} enable-docker: 'true' system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} ttl-hours: ${{ inputs.janitor_ttl_hours }} resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py benchmark: needs: - validate @@ -202,6 +240,12 @@ jobs: TSDG_BATCH_SIZE: '100000' O11YBENCH_GENERATOR_CPUS: '4' steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark - name: Checkout matching benchmark tools uses: actions/checkout@v4 with: @@ -283,7 +327,7 @@ jobs: needs: - provision - benchmark - if: ${{ always() && needs.provision.outputs.instance_id != '' }} + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -292,6 +336,7 @@ jobs: path: greptimedb persist-credentials: false - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' uses: ./greptimedb/.github/actions/aliyun-ecs-delete with: access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} @@ -300,3 +345,20 @@ jobs: region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} instance-id: ${{ needs.provision.outputs.instance_id }} runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py index 53257da8732..652eae77cfe 100644 --- a/tests/perf/test_aliyun_ecs_runner_scripts.py +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -30,6 +30,9 @@ from types import SimpleNamespace from unittest.mock import Mock, patch SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" +sys.path.insert(0, str(SCRIPTS_DIR)) + +import runner_utils def load_module(name: str, filename: str): @@ -58,7 +61,7 @@ class GitHubApiResponseTest(unittest.TestCase): response.read.return_value = body response.__enter__ = Mock(return_value=response) response.__exit__ = Mock(return_value=False) - with patch.object(provision.urllib.request, "urlopen", return_value=response): + with patch.object(runner_utils.urllib.request, "urlopen", return_value=response): self.assertEqual(provision.github_api("token", "DELETE", "/test"), expected) def test_invalid_json_still_fails(self): @@ -66,10 +69,26 @@ class GitHubApiResponseTest(unittest.TestCase): response.read.return_value = b"not json" response.__enter__ = Mock(return_value=response) response.__exit__ = Mock(return_value=False) - with patch.object(provision.urllib.request, "urlopen", return_value=response): + with patch.object(runner_utils.urllib.request, "urlopen", return_value=response): with self.assertRaises(ValueError): provision.github_api("token", "GET", "/test") + def test_http_errors_still_exit_instead_of_becoming_recoverable(self): + error = runner_utils.urllib.error.HTTPError( + "https://api.github.com/test", 403, "Forbidden", {}, io.BytesIO(b"denied") + ) + with patch.object(runner_utils.urllib.request, "urlopen", side_effect=error): + with self.assertRaisesRegex(SystemExit, "HTTP 403: denied"): + provision.github_api("token", "GET", "/test") + + def test_unregister_preserves_legacy_failure_semantics(self): + with patch.object(runner_utils, "find_runner_by_name", return_value={"id": 42}): + with patch.object(runner_utils, "github_api", side_effect=SystemExit("HTTP 403")): + with self.assertRaisesRegex(SystemExit, "HTTP 403"): + runner_utils.deregister_runner("token", "owner/repo", "runner") + with patch.object(runner_utils, "github_api", side_effect=RuntimeError("connection lost")): + self.assertFalse(runner_utils.deregister_runner("token", "owner/repo", "runner")) + class ProvisionNamingTest(unittest.TestCase): def test_runner_name_and_label_derive_from_run_id(self) -> None: diff --git a/tests/perf/test_ci_runner_scripts.py b/tests/perf/test_ci_runner_scripts.py new file mode 100644 index 00000000000..f9742079b49 --- /dev/null +++ b/tests/perf/test_ci_runner_scripts.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline regression coverage for CI runner selection and lifecycle.""" + +import importlib.util +import os +import re +import subprocess +import sys +import tempfile +import textwrap +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +SCRIPTS = Path(__file__).parents[2] / ".github/scripts" +sys.path.insert(0, str(SCRIPTS)) + + +def load(name, filename): + spec = importlib.util.spec_from_file_location(name, SCRIPTS / filename) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +config = load("ci_config_test", "ci-runner-config.py") +aws = load("aws_provision_test", "aws-ec2-runner-provision.py") +teardown = load("aws_teardown_test", "aws-ec2-runner-teardown.py") + + +def provision_config(**overrides): + return aws.ProvisionConfig(**({ + "region": "ap-southeast-1", "image_id": "ami-test", "instance_type": "c7i.2xlarge", + "subnet_id": "subnet-test", "security_group_id": "sg-test", "repo": "owner/repo", + "run_id": "123-2", "github_token": "secret", "system_disk_gib": 80, + } | overrides)) + + +def client_with_image(size=40, arch="x86_64"): + client = Mock() + client.describe_images.return_value = {"Images": [{ + "State": "available", "Architecture": arch, "RootDeviceName": "/dev/sda1", + "BlockDeviceMappings": [{"DeviceName": "/dev/sda1", "Ebs": {"VolumeSize": size}}], + }]} + client.describe_instance_types.return_value = {"InstanceTypes": [{ + "ProcessorInfo": {"SupportedArchitectures": ["x86_64"]}, + }]} + client.run_instances.return_value = {"Instances": [{"InstanceId": "i-test"}]} + client.get_console_output.return_value = {"Output": "boot log"} + return client + + +class ProviderConfigTest(unittest.TestCase): + def test_defaults_and_overrides(self): + for provider, requested, expected in ( + ("Aliyun", "auto", "ecs.c9i.2xlarge"), ("AWS", "auto", "c7i.2xlarge"), + ("Aliyun", "ecs.c9i.4xlarge", "ecs.c9i.4xlarge"), + ("AWS", "c6i.4xlarge", "c6i.4xlarge"), + ("AWS", "c7i.metal-24xl", "c7i.metal-24xl"), + ): + with self.subTest(provider=provider, requested=requested): + self.assertEqual(config.resolve_instance_type(provider, requested), expected) + + def test_wrong_provider_and_shell_fragments_rejected(self): + for provider, requested in (("ec2", "auto"), ("AWS", "ecs.c9i.2xlarge"), + ("Aliyun", "c7i.2xlarge"), ("AWS", ""), + ("AWS", "c7i.2xlarge; touch /tmp/unexpected")): + with self.subTest(provider=provider, requested=requested), self.assertRaises(ValueError): + config.resolve_instance_type(provider, requested) + + def test_cli_writes_resolved_output(self): + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "output" + result = subprocess.run([sys.executable, str(SCRIPTS / "ci-runner-config.py"), + "--provider", "AWS"], + env=os.environ | {"GITHUB_OUTPUT": str(output)}, + capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(output.read_text(), "instance_type=c7i.2xlarge\n") + + + +class BenchmarkSetupTest(unittest.TestCase): + def run_setup(self, missing_tools="", docker_installed=True, docker_ready=True): + action = (SCRIPTS.parent / "actions/setup-benchmark/action.yml").read_text() + steps = re.findall(r" run: \|\n((?: .*\n|\n)+)", action) + self.assertEqual(len(steps), 2) + script = "\n".join(textwrap.dedent(step) for step in steps) + # Execute the actual action twice with host commands stubbed: installation + # changes availability, so the second pass must perform no mutations. + prelude = r""" +function . { ID=ubuntu; VERSION_CODENAME=noble; } +uname() { echo x86_64; } +dpkg() { echo amd64; } +curl() { echo fixture-key; } +command() { + if [[ "$1" == -v ]]; then + if [[ "$2" == docker ]]; then [[ "$DOCKER_INSTALLED" == true ]]; return; fi + if [[ " $MISSING_TOOLS " == *" $2 "* ]]; then return 1; fi + return 0 + fi + builtin command "$@" +} +docker() { [[ "$DOCKER_READY" == true ]]; } +sudo() { + printf '%s\n' "$*" >> "$CALLS" + if [[ "$*" == *' apt-get install '* ]]; then + MISSING_TOOLS='' + [[ "$*" != *' docker-ce '* ]] || DOCKER_INSTALLED=true + fi + if [[ "$*" == *'systemctl enable --now docker' ]]; then DOCKER_READY=true; fi + if [[ "$*" == *' tee '* ]]; then while IFS= read -r line; do :; done; fi + return 0 +} +""" + with tempfile.TemporaryDirectory() as tmp: + calls = Path(tmp) / "calls" + calls.touch() + env = os.environ | {"CALLS": str(calls), "MISSING_TOOLS": missing_tools, + "DOCKER_INSTALLED": str(docker_installed).lower(), + "DOCKER_READY": str(docker_ready).lower()} + result = subprocess.run(["bash", "-c", prelude + script + + '\nprintf "SECOND PASS\\n" >> "$CALLS"\n' + script], + env=env, capture_output=True, text=True) + self.assertEqual(result.returncode, 0, result.stderr) + first, second = calls.read_text().split("SECOND PASS\n") + self.assertEqual(second, "", "Second setup must reuse installed dependencies") + return first + + def test_prepared_runner_skips_all_privileged_changes(self): + self.assertEqual(self.run_setup(), "") + + def test_missing_tools_are_installed_without_changing_docker(self): + calls = self.run_setup(missing_tools="jq rsync") + self.assertIn("--no-install-recommends jq rsync", calls) + self.assertIn("NEEDRESTART_MODE=l", calls) + self.assertNotIn("docker", calls) + + def test_missing_docker_is_installed_once(self): + calls = self.run_setup(docker_installed=False, docker_ready=False) + self.assertIn("NEEDRESTART_MODE=l", calls) + self.assertIn("docker-ce docker-ce-cli containerd.io docker-buildx-plugin", calls) + self.assertIn("systemctl enable --now docker", calls) + + def test_existing_stopped_docker_is_started_without_installation(self): + self.assertEqual(self.run_setup(docker_ready=False), "-n systemctl enable --now docker\n") + + +class AwsProvisionTest(unittest.TestCase): + def test_host_policy_matches_aliyun_before_runner_activation(self): + aliyun = load("aliyun_host_policy_test", "aliyun-ecs-runner-provision.py") + ecs_script = aliyun.render_user_data("name", "label", "token", "owner/repo") + aws_script = aws.render_user_data("name", "token", "owner/repo", "https://example.com/runner.tgz") + start = ecs_script.index("# Mask systemd-oomd") + end = ecs_script.index("free --human", start) + len("free --human") + policy = ecs_script[start:end] + self.assertIn(policy, aws_script) + self.assertLess(aws_script.index(policy), aws_script.index("./svc.sh start")) + self.assertLess(aws_script.index("OOMPolicy=continue"), aws_script.index("./svc.sh start")) + self.assertLess(aws_script.index("./svc.sh install runner"), aws_script.index("service_name=$(cat .service)")) + self.assertIn('"/etc/systemd/system/$service_name.d/oom.conf"', aws_script) + + def test_ubuntu_image_resolution_and_override(self): + client = Mock() + client.describe_images.return_value = {"Images": [ + {"ImageId": "ami-new", "CreationDate": "2026-09-04", "Name": "noble-new"}, + {"ImageId": "ami-old", "CreationDate": "2026-08-04", "Name": "noble-old"}, + ]} + self.assertEqual(aws.resolve_image(client, "auto"), "ami-new") + request = client.describe_images.call_args.kwargs + self.assertEqual(request["Owners"], ["099720109477"]) + self.assertIn({"Name": "name", "Values": [ + "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"]}, request["Filters"]) + client.reset_mock() + self.assertEqual(aws.resolve_image(client, "ami-pinned"), "ami-pinned") + client.describe_images.assert_not_called() + client.describe_images.return_value = {"Images": []} + with self.assertRaisesRegex(ValueError, "No Canonical Ubuntu"): + aws.resolve_image(client, "auto") + + def test_disk_size_and_architecture_fail_before_allocation(self): + for client, pattern in ((client_with_image(size=500), "smaller than AMI"), + (client_with_image(arch="arm64"), "x86_64")): + with self.subTest(pattern=pattern), patch.object(aws, "create_registration_token") as token: + with self.assertRaisesRegex(ValueError, pattern): + aws.run_instance(client, provision_config()) + client.run_instances.assert_not_called() + token.assert_not_called() + + def test_arm_instance_rejected_before_registration(self): + client = client_with_image() + client.describe_instance_types.return_value["InstanceTypes"][0]["ProcessorInfo"]["SupportedArchitectures"] = ["arm64"] + with patch.object(aws, "create_registration_token") as token: + with self.assertRaisesRegex(ValueError, "x86_64"): + aws.run_instance(client, provision_config(instance_type="r8g.2xlarge")) + token.assert_not_called() + client.run_instances.assert_not_called() + + def test_create_request_and_recovery_tags(self): + client = client_with_image() + with patch.object(aws, "create_registration_token", return_value="short-lived-token"), \ + patch.object(aws, "runner_download", return_value="https://example.com/runner.tar.gz"): + self.assertEqual(aws.run_instance(client, provision_config(system_disk_gib=100)), "i-test") + request = client.run_instances.call_args.kwargs + self.assertEqual(request["BlockDeviceMappings"], [{"DeviceName": "/dev/sda1", "Ebs": { + "VolumeSize": 100, "VolumeType": "gp3", "Iops": 3000, "Throughput": 125, + "DeleteOnTermination": True}}]) + self.assertEqual(request["ClientToken"], "ci-ec2-123-2") + self.assertNotIn("InstanceMarketOptions", request) + self.assertNotIn("IamInstanceProfile", request) + for specification in request["TagSpecifications"]: + tags = {t["Key"]: t["Value"] for t in specification["Tags"]} + self.assertEqual(tags["managed-by"], "query-regression-ci") + self.assertEqual(tags[aws.REPO_TAG], "owner/repo") + self.assertEqual(tags["query-regression-run-id"], "123-2") + self.assertEqual(tags["runner-ttl-hours"], "8") + self.assertNotIn("secret", request["UserData"]) + self.assertIn("--token short-lived-token", request["UserData"]) + self.assertNotIn("apt-get install", request["UserData"]) + self.assertIn('for tool in git curl tar sudo; do command -v "$tool"; done', request["UserData"]) + self.assertLess(request["UserData"].index("usermod -aG docker runner"), + request["UserData"].index("./svc.sh start")) + self.assertIn("visudo -cf /etc/sudoers.d/ci-runner", request["UserData"]) + self.assertLess(request["UserData"].index("$nrconf{override_rc}"), + request["UserData"].index("./svc.sh start")) + check = subprocess.run(["bash", "-n"], input=request["UserData"], text=True, capture_output=True) + self.assertEqual(check.returncode, 0, check.stderr) + + def test_partial_provision_keeps_outputs_and_console(self): + client = client_with_image() + client.get_waiter.return_value.wait.side_effect = TimeoutError("not running") + with tempfile.TemporaryDirectory() as tmp: + output = Path(tmp) / "output" + with patch.dict(os.environ, {"GITHUB_OUTPUT": str(output)}), patch.object(aws, "run_instance", return_value="i-test"): + with self.assertRaises(TimeoutError): + aws.provision(client, provision_config()) + self.assertEqual(output.read_text().splitlines(), [ + "instance_id=i-test", "label=ci-ec2-123-2", + "runner_name=ci-ec2-123-2"]) + client.get_console_output.assert_called_once() + + def test_online_poll_and_timeout(self): + for online in (True, False): + with self.subTest(online=online), patch.object(aws, "run_instance", return_value="i-test"), \ + patch.object(aws, "find_runner_by_name", side_effect=[None, {"status": "online"}]), \ + patch.object(aws.time, "sleep"), \ + patch.object(aws.time, "monotonic", side_effect=[0, 1, 2 if online else 601]): + client = client_with_image() + if online: + self.assertEqual(aws.provision(client, provision_config()), 0) + else: + with self.assertRaises(TimeoutError): + aws.provision(client, provision_config()) + client.get_console_output.assert_called_once() + + +class AwsTeardownTest(unittest.TestCase): + NOW = datetime(2026, 9, 23, 10, tzinfo=timezone.utc) + + def setUp(self): + sleeper = patch.object(teardown.time, "sleep") + sleeper.start() + self.addCleanup(sleeper.stop) + + def instance(self, identity="i-test", age=9, ttl="8"): + tags = [{"Key": teardown.RUNNER_TAG, "Value": "ci-ec2-123-2"}] + if ttl is not None: + tags.append({"Key": teardown.TTL_TAG_KEY, "Value": ttl}) + return {"InstanceId": identity, "LaunchTime": self.NOW - timedelta(hours=age), "Tags": tags} + + def test_pagination_and_ownership_filters(self): + client = Mock() + client.get_paginator.return_value.paginate.return_value = [ + {"Reservations": [{"Instances": [self.instance("i-a")]}]}, + {"Reservations": [{"Instances": [self.instance("i-b")]}]}, + ] + self.assertEqual([i["InstanceId"] for i in teardown.managed_instances(client, "owner/repo", "123-2")], ["i-a", "i-b"]) + self.assertEqual(client.get_paginator.return_value.paginate.call_args.kwargs["Filters"], [ + {"Name": "tag:managed-by", "Values": ["query-regression-ci"]}, + {"Name": "tag:github-repository", "Values": ["owner/repo"]}, + {"Name": "tag:query-regression-run-id", "Values": ["123-2"]}, + ]) + + def test_sweep_expiry_and_missing_or_invalid_ttl(self): + instances = [self.instance("expired", age=8), self.instance("fresh", age=7), + self.instance("invalid", ttl="bad"), + self.instance("missing-expired", age=4, ttl=None), + self.instance("missing-fresh", age=3, ttl=None)] + with patch.object(teardown, "managed_instances", return_value=instances), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", now=self.NOW), 0) + self.assertEqual([c.args[1] for c in delete.call_args_list], ["expired", "missing-expired"]) + self.assertEqual(unregister.call_count, 2) + + def test_ttl_decisions_match_aliyun(self): + aliyun = load("aliyun_ttl_policy_test", "aliyun-ecs-runner-teardown.py") + for age, ttl in [(3, None), (4, None), (7, "8"), (8, "8"), (9, "8"), + (10, "bad"), (10, ""), (10, "0"), (200, "169"), (168, "168")]: + instance = self.instance(age=age, ttl=ttl) + created = instance["LaunchTime"].strftime("%Y-%m-%dT%H:%M:%SZ") + expected = aliyun.expired_instance_names( + [(instance["InstanceId"], "runner", created, ttl)], self.NOW, timedelta(hours=4)) + with self.subTest(age=age, ttl=ttl), \ + patch.object(teardown, "managed_instances", return_value=[instance]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", now=self.NOW), 0) + self.assertEqual([call.args[1] for call in delete.call_args_list], + [identity for identity, _ in expected]) + + def test_targeted_cleanup_recovers_missing_outputs_without_waiting_for_ttl(self): + with patch.object(teardown, "managed_instances", return_value=[self.instance(age=0)]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + delete.assert_called_once() + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + + def test_independent_cloud_and_github_cleanup_failures(self): + for cloud_ok, github_result in ((False, True), (True, RuntimeError("GitHub unavailable"))): + with self.subTest(cloud_ok=cloud_ok), \ + patch.object(teardown, "managed_instances", return_value=[self.instance()]), \ + patch.object(teardown, "delete_instance", return_value=cloud_ok) as delete, \ + patch.object(teardown, "deregister_runner", side_effect=[github_result]) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 1) + delete.assert_called_once() + unregister.assert_called_once() + + def test_already_gone_instance_still_unregisters(self): + with patch.object(teardown, "managed_instances", return_value=[]), \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + unregister.assert_called_once() + + def test_retried_teardown_uses_original_instance_and_runner(self): + client = Mock() + client.get_paginator.return_value.paginate.return_value = [ + {"Reservations": [{"Instances": [self.instance()]}]}, + ] + with patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(client, "owner/repo", "token", run_id="123-3", + instance_id="i-test", runner_name="ci-ec2-123-2"), 0) + filters = client.get_paginator.return_value.paginate.call_args.kwargs["Filters"] + self.assertIn({"Name": "instance-id", "Values": ["i-test"]}, filters) + self.assertFalse(any(f["Name"] == "tag:query-regression-run-id" for f in filters)) + delete.assert_called_once_with(client, "i-test") + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + client.get_paginator.return_value.paginate.return_value = [] + with patch.object(teardown, "deregister_runner", return_value=True) as unregister: + self.assertEqual(teardown.teardown(client, "owner/repo", "token", run_id="123-3", + instance_id="i-test", runner_name="ci-ec2-123-2"), 0) + unregister.assert_called_once_with("token", "owner/repo", "ci-ec2-123-2") + + def test_terminate_waits_and_handles_not_found(self): + client = client_with_image() + self.assertTrue(teardown.delete_instance(client, "i-test")) + client.get_waiter.assert_called_with("instance_terminated") + error = RuntimeError("gone") + error.response = {"Error": {"Code": "InvalidInstanceID.NotFound"}} + client.terminate_instances.side_effect = error + with patch.object(teardown.time, "sleep"): + self.assertTrue(teardown.delete_instance(client, "i-test")) + + def test_transient_not_found_is_retried(self): + client = client_with_image() + error = RuntimeError("not visible yet") + error.response = {"Error": {"Code": "InvalidInstanceID.NotFound"}} + client.terminate_instances.side_effect = [error, {}] + self.assertTrue(teardown.delete_instance(client, "i-test")) + self.assertEqual(client.terminate_instances.call_count, 2) + + def test_missing_tags_are_polled_before_cleanup(self): + with patch.object(teardown, "managed_instances", side_effect=[[], [self.instance()]]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", return_value=True): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 0) + delete.assert_called_once() + + def test_github_system_exit_does_not_skip_other_instances(self): + with patch.object(teardown, "managed_instances", return_value=[self.instance("i-a"), self.instance("i-b")]), \ + patch.object(teardown, "delete_instance", return_value=True) as delete, \ + patch.object(teardown, "deregister_runner", side_effect=[SystemExit("HTTP 403"), True]): + self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 1) + self.assertEqual([c.args[1] for c in delete.call_args_list], ["i-a", "i-b"]) + + +if __name__ == "__main__": + unittest.main() From 20dde2601f0bc0bc7812ccc691803da407330ecd Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 23 Sep 2026 13:03:51 +0000 Subject: [PATCH 025/133] feat: enable native histogram ingestion by default (#9301) * feat: enable native histogram ingestion by default Signed-off-by: shuiyisong * chore: add comments Signed-off-by: shuiyisong * fix: test Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- config/config.md | 4 - config/frontend.example.toml | 4 - config/standalone.example.toml | 4 - docs/rfcs/2026-08-04-native-histograms.md | 20 ++- src/cmd/src/standalone.rs | 1 - src/cmd/tests/load_config_test.rs | 72 +++----- src/frontend/src/frontend.rs | 1 - src/frontend/src/server.rs | 11 +- src/frontend/src/service_config/otlp.rs | 10 -- src/frontend/src/service_config/prom_store.rs | 5 - src/metric-engine/src/engine/bulk_insert.rs | 21 +-- src/mito2/src/memtable/bulk/part.rs | 8 +- src/mito2/src/sst/parquet/flat_format.rs | 35 +++- src/servers/benches/prom_decode.rs | 7 +- src/servers/src/http.rs | 11 +- src/servers/src/http/otlp.rs | 3 - src/servers/src/http/prom_store.rs | 10 +- src/servers/src/otel_arrow.rs | 120 +++++++++---- src/servers/src/otlp/metrics.rs | 52 ++---- src/servers/src/otlp/metrics/resource_info.rs | 51 ++---- .../otlp/metrics/resource_info/tests/delta.rs | 6 +- src/servers/src/prom_remote_write/v2.rs | 93 ++++------ src/servers/tests/http/prom_store_test.rs | 68 +------- .../tests/prom_remote_write_v2_test.rs | 2 +- src/session/src/protocol_ctx.rs | 2 - tests-integration/src/otlp.rs | 3 +- tests-integration/src/test_util.rs | 47 +---- tests-integration/src/tests/instance_test.rs | 163 ++++++++++++++++++ tests-integration/tests/http.rs | 61 +------ 29 files changed, 429 insertions(+), 466 deletions(-) diff --git a/config/config.md b/config/config.md index c33918a8573..6544a28f93d 100644 --- a/config/config.md +++ b/config/config.md @@ -93,14 +93,12 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | -| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | -| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | | `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`**. | @@ -361,14 +359,12 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | -| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | | `prom_store.prom_validation_mode` | String | `strict` | Whether to enable validation for Prometheus remote write requests.
Available options:
- strict: deny invalid UTF-8 strings (default).
- lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD).
- unchecked: do not valid strings. | -| `prom_store.experimental_enable_prometheus_native_histogram` | Bool | `false` | Experimental: enable Prometheus remote write v2 native histogram ingestion. | | `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 055acbb93bd..b528da0304f 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -295,8 +295,6 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true -## Experimental: enable cumulative OTLP exponential histogram ingestion. -experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric @@ -315,8 +313,6 @@ with_metric_engine = true ## - lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD). ## - unchecked: do not valid strings. prom_validation_mode = "strict" -## Experimental: enable Prometheus remote write v2 native histogram ingestion. -experimental_enable_prometheus_native_histogram = false ## The metasrv client options. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 357ec5904f7..ef9dd561f6a 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -274,8 +274,6 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true -## Experimental: enable cumulative OTLP exponential histogram ingestion. -experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric @@ -294,8 +292,6 @@ with_metric_engine = true ## - lossy: allow invalid UTF-8 strings, replace invalid characters with REPLACEMENT_CHARACTER(U+FFFD). ## - unchecked: do not valid strings. prom_validation_mode = "strict" -## Experimental: enable Prometheus remote write v2 native histogram ingestion. -experimental_enable_prometheus_native_histogram = false ## The WAL options. diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index c714b35351b..6260e79f3d5 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -13,8 +13,8 @@ field and evaluates them as first-class PromQL samples. Prometheus Remote Write ingestion transport: accepted cumulative points are normalized into the same Struct before persistence and are queried only as native histograms. -Native histograms are experimental. Prometheus Remote Write 2.0 and cumulative -OTLP exponential histograms are supported behind separate configuration gates. +Prometheus Remote Write 2.0 native histograms and cumulative OTLP/HTTP +exponential histograms are enabled by default. Remote Write 1.0 histogram payloads are rejected instead of being acknowledged and dropped. Native-histogram Remote Read is deferred; the existing Remote Read path continues to return scalar samples only. @@ -95,8 +95,7 @@ that send native histograms must use Remote Write 2.0. ## Remote Write 2.0 -Remote Write 2.0 accepts integer and float native histograms while -`prom_store.experimental_enable_prometheus_native_histogram` is enabled. Supported +Remote Write 2.0 accepts integer and float native histograms by default. Supported exponential schemas are `-4` through `8`; schema `-53` represents native histograms with custom buckets. @@ -121,10 +120,8 @@ explicit if sampled responses are implemented. ## OTLP -OTLP exponential histograms are accepted when -`otlp.experimental_enable_exponential_histogram` is enabled. The option defaults -to false and applies to OTLP/HTTP. Disabled points are rejected rather than -silently acknowledged. OTel Arrow exponential histograms are rejected because +OTLP/HTTP exponential histograms are accepted by default. +OTel Arrow exponential histograms are rejected because the current Arrow wire format omits `zero_threshold`; accepting them would silently change the distribution. Cumulative temporality is required; delta and unspecified exponential histograms are rejected before their points are @@ -284,3 +281,10 @@ native-histogram data remains readable. - Versioned two-phase native-histogram aggregation state. - Exact Prometheus summation compensation if measured precision differences justify the extra aggregate state. + +## Configuration migration + +The `prom_store.experimental_enable_prometheus_native_histogram` and +`otlp.experimental_enable_exponential_histogram` options have been removed. +Remove them from existing configurations; both ingestion paths are always enabled +when their protocol is enabled. Old option values are ignored, including `false`. diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 4573b0f129f..60556428117 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -1139,7 +1139,6 @@ mod tests { fn test_toml() { let opts = StandaloneOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); } diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 1c4a83e1397..34c5bc8c9ba 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -199,12 +199,6 @@ fn test_load_frontend_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); - assert!( - !options - .component - .otlp - .experimental_enable_exponential_histogram - ); let expected = GreptimeOptions:: { component: FrontendOptions { pending_rows_batcher: PendingRowsBatcherOptions { @@ -386,12 +380,6 @@ fn test_load_standalone_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); - assert!( - !options - .component - .otlp - .experimental_enable_exponential_histogram - ); let expected = GreptimeOptions:: { component: StandaloneOptions { pending_rows_batcher: PendingRowsBatcherOptions { @@ -448,40 +436,34 @@ fn test_load_standalone_example_config() { } #[test] -fn test_load_otlp_exponential_histogram_option() { - let config = tempfile::NamedTempFile::new().unwrap(); - std::fs::write( - config.path(), - "[otlp]\nexperimental_enable_exponential_histogram = true\n", - ) - .unwrap(); +fn test_load_removed_histogram_options() { + for enabled in [false, true] { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + format!( + "[otlp]\nexperimental_enable_exponential_histogram = {enabled}\n\ + [prom_store]\nenable = true\nwith_metric_engine = true\n\ + experimental_enable_prometheus_native_histogram = {enabled}\n" + ), + ) + .unwrap(); - let frontend = - GreptimeOptions::::load_layered_options(config.path().to_str(), "") - .unwrap(); - assert!( - frontend - .component - .otlp - .experimental_enable_exponential_histogram - ); - - let standalone = - GreptimeOptions::::load_layered_options(config.path().to_str(), "") - .unwrap(); - assert!( - standalone - .component - .otlp - .experimental_enable_exponential_histogram - ); - assert!( - standalone - .component - .frontend_options() - .otlp - .experimental_enable_exponential_histogram - ); + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let defaults = FrontendOptions::default(); + for options in [frontend.component, standalone.component.frontend_options()] { + assert_eq!(options.otlp, defaults.otlp); + assert_eq!(options.prom_store, defaults.prom_store); + let serialized = toml::to_string(&options).unwrap(); + assert!(!serialized.contains("experimental_enable_exponential_histogram")); + assert!(!serialized.contains("experimental_enable_prometheus_native_histogram")); + } + } } #[test] diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 51e291c7906..b2178e1ff35 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -314,7 +314,6 @@ max_batch_rows = 25 let enabled: FrontendOptions = toml::from_str("experimental_metric_export = true").unwrap(); assert!(enabled.experimental_metric_export); let toml_string = toml::to_string(&opts).unwrap(); - assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); assert_eq!(parsed.otlp, opts.otlp); assert_eq!(parsed.influxdb, opts.influxdb); diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index fb1ab2f2603..da89d871ac0 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -146,19 +146,14 @@ where Some(self.instance.clone()), opts.prom_store.with_metric_engine, opts.prom_store.prom_validation_mode, - opts.prom_store - .experimental_enable_prometheus_native_histogram, pending_rows_batcher, ) .with_prometheus_handler(self.instance.clone()); } if opts.otlp.enable { - builder = builder.with_otlp_handler( - self.instance.clone(), - opts.prom_store.with_metric_engine, - opts.otlp.experimental_enable_exponential_histogram, - ); + builder = builder + .with_otlp_handler(self.instance.clone(), opts.prom_store.with_metric_engine); } if opts.jaeger.enable { @@ -623,8 +618,6 @@ mod tests { opts.prom_store.pending_rows_flush_interval = Duration::from_secs(2); opts.prom_store.with_metric_engine = metric_engine; opts.prom_store.enable = prom_enabled; - opts.prom_store - .experimental_enable_prometheus_native_histogram = true; let shared = &mut opts.pending_rows_batcher.table; shared.protocols = vec![if selected { BatchingProtocol::Prom diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index d6a36f702e4..e67f4b5bcba 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -20,7 +20,6 @@ const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; #[serde(default)] pub struct OtlpOptions { pub enable: bool, - pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, /// Whether to synthesize the `greptime_otel_resource_info` descriptor @@ -34,7 +33,6 @@ impl Default for OtlpOptions { fn default() -> Self { Self { enable: true, - experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, experimental_enable_resource_info: false, } @@ -49,13 +47,11 @@ mod tests { fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); - assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); - assert!(!options.experimental_enable_exponential_histogram); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -63,15 +59,9 @@ mod tests { let options: OtlpOptions = toml::from_str("trace_ingest_chunk_size = 0").unwrap(); assert!(options.enable); - assert!(!options.experimental_enable_exponential_histogram); assert_eq!(options.trace_ingest_chunk_size, 0); - let options: OtlpOptions = - toml::from_str("experimental_enable_exponential_histogram = true").unwrap(); - assert!(options.experimental_enable_exponential_histogram); - let serialized = toml::to_string(&options).unwrap(); - assert!(serialized.contains("experimental_enable_exponential_histogram = true")); assert_eq!(toml::from_str::(&serialized).unwrap(), options); } } diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs index 3d00f48cdda..575ec600c15 100644 --- a/src/frontend/src/service_config/prom_store.rs +++ b/src/frontend/src/service_config/prom_store.rs @@ -27,9 +27,6 @@ pub struct PromStoreOptions { /// Validation mode while decoding Prometheus remote write requests. #[serde(default)] pub prom_validation_mode: PromValidationMode, - /// Enables experimental Prometheus remote write v2 native histogram ingestion. - #[serde(default)] - pub experimental_enable_prometheus_native_histogram: bool, #[serde(default, with = "humantime_serde")] pub pending_rows_flush_interval: Duration, #[serde(default = "default_max_batch_rows")] @@ -89,7 +86,6 @@ impl Default for PromStoreOptions { enable: true, with_metric_engine: true, prom_validation_mode: PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram: false, pending_rows_flush_interval: Duration::ZERO, max_batch_rows: default_max_batch_rows(), max_concurrent_flushes: default_max_concurrent_flushes(), @@ -152,7 +148,6 @@ mod tests { assert!(default.enable); assert!(default.with_metric_engine); assert_eq!(default.prom_validation_mode, PromValidationMode::Strict); - assert!(!default.experimental_enable_prometheus_native_histogram); assert_eq!(default.pending_rows_flush_interval, Duration::ZERO); assert_eq!(default.max_batch_rows, default_max_batch_rows()); assert_eq!( diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index cb43a71406d..fa47ce48008 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -18,10 +18,9 @@ use api::v1::{ArrowIpc, SemanticType}; use bytes::Bytes; use common_grpc::flight::{FlightEncoder, FlightMessage}; use datatypes::arrow::record_batch::RecordBatch; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{OptionExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; -use store_api::region_engine::RegionEngine; use store_api::region_request::{AffectedRows, RegionBulkInsertsRequest, RegionRequest}; use store_api::storage::RegionId; @@ -74,9 +73,9 @@ impl MetricEngineInner { region_id: RegionId, mut request: RegionBulkInsertsRequest, ) -> Result { - // Simply set the aligned schema to the data region schema version to avoid filling missing columns - // because that schema should be constant and callers have ensured request has the same schema. - request.aligned_schema_version = Some(self.physical_schema_version(region_id).await?); + // Other logical tables can add fields to the shared physical schema. + // Let Mito fill fields absent from this batch before writing it. + request.aligned_schema_version = None; self.data_region .write_data(region_id, RegionRequest::BulkInserts(request)) .await @@ -121,7 +120,6 @@ impl MetricEngineInner { let (schema, data_header, payload) = record_batch_to_ipc(&modified_batch)?; let partition_expr_version = request.partition_expr_version; - let aligned_schema_version = Some(self.physical_schema_version(data_region_id).await?); let request = RegionBulkInsertsRequest { skip_wal: request.skip_wal, @@ -133,22 +131,13 @@ impl MetricEngineInner { payload, }, partition_expr_version, - aligned_schema_version, + aligned_schema_version: None, }; self.data_region .write_data(data_region_id, RegionRequest::BulkInserts(request)) .await } - async fn physical_schema_version(&self, region_id: RegionId) -> Result { - Ok(self - .mito - .get_metadata(region_id) - .await - .context(error::MitoReadOperationSnafu)? - .schema_version) - } - fn resolve_tag_columns_from_metadata( &self, logical_region_id: RegionId, diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index f994f242273..7e50fbbd4cb 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -20,7 +20,7 @@ use std::time::{Duration, Instant}; use api::helper::{ColumnDataTypeWrapper, to_grpc_value}; use api::v1::bulk_wal_entry::Body; -use api::v1::{ArrowIpc, BulkWalEntry, Mutation, OpType}; +use api::v1::{ArrowIpc, BulkWalEntry, Mutation, OpType, SemanticType}; use bytes::Bytes; use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage}; use common_recordbatch::DfRecordBatch as RecordBatch; @@ -218,6 +218,12 @@ impl BulkPart { // Finds columns that need to be filled let mut columns_to_fill = Vec::new(); for column_meta in ®ion_metadata.column_metadatas { + // Sparse bulk batches already carry tags in the encoded primary key. + if region_metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse + && column_meta.semantic_type == SemanticType::Tag + { + continue; + } // TODO(yingwen): Returns error if it is impure default after we support filling // bulk insert request in the frontend if !batch_columns.contains(column_meta.column_schema.name.as_str()) { diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 34fb8cb873c..3abba4d2db6 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -37,7 +37,7 @@ use datatypes::arrow::array::{ Array, ArrayRef, BinaryArray, DictionaryArray, UInt32Array, UInt64Array, }; use datatypes::arrow::compute::kernels::take::take; -use datatypes::arrow::datatypes::{Schema, SchemaRef}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::value::ValueRef; @@ -407,11 +407,42 @@ impl FlatReadFormat { }; // First, apply flat format conversion. - let batch = match &self.parquet_adapter { + let mut batch = match &self.parquet_adapter { ParquetAdapter::Flat(_) => record_batch, ParquetAdapter::PrimaryKeyToFlat(p) => p.convert_batch(record_batch)?, }; + // Normalize nested field names and metadata to the SST's region metadata + // before schema compatibility and merging with memtables. This removes + // Parquet-added field IDs; equals_datatype also permits nested name differences. + for index in 0..batch.num_columns() { + let array = batch.column(index); + if !matches!(array.data_type(), ArrowDataType::Struct(_)) { + continue; + } + let field = batch.schema_ref().field(index); + let Some(column) = self.metadata().column_by_name(field.name()) else { + continue; + }; + let target = column.column_schema.data_type.as_arrow_type(); + if array.data_type() != &target && array.data_type().equals_datatype(&target) { + let array = + datatypes::arrow::compute::cast(array, &target).context(ComputeArrowSnafu)?; + let mut fields = batch.schema().fields().to_vec(); + fields[index] = Arc::new(field.clone().with_data_type(target)); + let mut columns = batch.columns().to_vec(); + columns[index] = array; + batch = RecordBatch::try_new( + Arc::new(Schema::new_with_metadata( + fields, + batch.schema().metadata().clone(), + )), + columns, + ) + .context(NewRecordBatchSnafu)?; + } + } + // Then apply sequence override if provided let Some(override_array) = override_sequence_array else { return Ok(batch); diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index 06b0d02e940..1cdc7d3782e 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -288,10 +288,9 @@ fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) { group.bench_with_input(BenchmarkId::new("v2_manual", name), v2_bytes, |b, bytes| { b.iter(|| { black_box( - remote_write_v2::decode_uncompressed_write_requests( - black_box(bytes.as_slice()), - true, - ) + remote_write_v2::decode_uncompressed_write_requests(black_box( + bytes.as_slice(), + )) .unwrap(), ); }); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index f947ec1621b..a5266576cbe 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -728,7 +728,6 @@ impl HttpServerBuilder { pipeline_handler: Option, prom_store_with_metric_engine: bool, prom_validation_mode: PromValidationMode, - experimental_enable_prometheus_native_histogram: bool, pending_rows_batcher: Option>, ) -> Self { let state = PromStoreState { @@ -736,7 +735,6 @@ impl HttpServerBuilder { pipeline_handler, prom_store_with_metric_engine, prom_validation_mode, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, }; @@ -763,16 +761,11 @@ impl HttpServerBuilder { self, handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, - experimental_enable_exponential_histogram: bool, ) -> Self { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp( - handler, - with_metric_engine, - experimental_enable_exponential_histogram, - ), + HttpServer::route_otlp(handler, with_metric_engine), ), ..self } @@ -1500,7 +1493,6 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, - experimental_enable_exponential_histogram: bool, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) @@ -1516,7 +1508,6 @@ impl HttpServer { )) .with_state(OtlpState { with_metric_engine, - experimental_enable_exponential_histogram, handler: otlp_handler, }) } diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index 2d5385c5e33..717e6ca65da 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -77,7 +77,6 @@ fn content_type_to_string(content_type: Option<&TypedHeader>) -> St #[derive(Clone)] pub struct OtlpState { pub with_metric_engine: bool, - pub experimental_enable_exponential_histogram: bool, pub handler: OpenTelemetryProtocolHandlerRef, } @@ -108,7 +107,6 @@ pub async fn metrics( let OtlpState { with_metric_engine, - experimental_enable_exponential_histogram, handler, } = state; @@ -117,7 +115,6 @@ pub async fn metrics( resource_attrs: http_opts.resource_attrs, promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, - experimental_enable_exponential_histogram, // set by the frontend from its config is_legacy: false, resource_info: false, diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 17db6abae04..2c4b6927451 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -74,7 +74,6 @@ pub struct PromStoreState { pub pipeline_handler: Option, pub prom_store_with_metric_engine: bool, pub prom_validation_mode: PromValidationMode, - pub experimental_enable_prometheus_native_histogram: bool, pub pending_rows_batcher: Option>, } @@ -146,7 +145,6 @@ async fn remote_write_v1( pipeline_handler, prom_store_with_metric_engine, prom_validation_mode, - experimental_enable_prometheus_native_histogram: _, pending_rows_batcher, } = state; @@ -228,7 +226,6 @@ async fn remote_write_v2( pipeline_handler: _, prom_store_with_metric_engine, prom_validation_mode: _, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, } = state; @@ -243,11 +240,7 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let req = match decode_remote_write_v2( - is_zstd, - body, - experimental_enable_prometheus_native_histogram, - ) { + let req = match decode_remote_write_v2(is_zstd, body) { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -1017,7 +1010,6 @@ mod tests { pipeline_handler: None, prom_store_with_metric_engine: false, prom_validation_mode: PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram: false, pending_rows_batcher: None, } } diff --git a/src/servers/src/otel_arrow.rs b/src/servers/src/otel_arrow.rs index 0cf07c76468..490ae4b1c58 100644 --- a/src/servers/src/otel_arrow.rs +++ b/src/servers/src/otel_arrow.rs @@ -24,6 +24,7 @@ use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_serv use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ BatchArrowRecords, BatchStatus, StatusCode as ArrowStatusCode, }; +use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; use tonic::metadata::{Entry, MetadataValue}; @@ -50,23 +51,42 @@ impl OtelArrowServiceHandler { } } +/// Removes unsupported Arrow histograms before ingestion, preserving other metrics. +fn remove_exponential_histograms(request: &mut ExportMetricsServiceRequest) -> bool { + let mut has_data_points = false; + for scope in request + .resource_metrics + .iter_mut() + .flat_map(|resource| &mut resource.scope_metrics) + { + scope.metrics.retain(|item| { + if let Some(metric::Data::ExponentialHistogram(histogram)) = &item.data { + has_data_points |= !histogram.data_points.is_empty(); + false + } else { + true + } + }); + } + has_data_points +} + fn batch_status( batch_id: i64, outcome: MetricsIngestOutcome, has_exponential_histogram_data_points: bool, ) -> BatchStatus { - let status_code = if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + let status_code = if outcome.accepted_data_points == 0 + && (outcome.rejected_data_points > 0 || has_exponential_histogram_data_points) + { ArrowStatusCode::InvalidArgument } else { ArrowStatusCode::Ok }; - let status_message = match outcome.error_message { - // Arrow keeps the feature gate off, so these fail before per-point validation. - Some(_) if has_exponential_histogram_data_points => { - EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() - } - Some(message) => message, - None => String::new(), + let status_message = if has_exponential_histogram_data_points { + EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() + } else { + outcome.error_message.unwrap_or_default() }; BatchStatus { batch_id, @@ -117,7 +137,7 @@ impl ArrowMetricsService for OtelArrowServiceHandler outcome, Err(error::Error::InvalidOtlpMetricInput { reason }) => { @@ -202,19 +212,61 @@ mod tests { use super::*; #[test] - fn batch_status_explains_arrow_exponential_histogram_limit() { - let status = batch_status( - 7, - MetricsIngestOutcome { - rejected_data_points: 1, - error_message: Some("internal OTLP rejection detail".to_string()), - ..Default::default() - }, - true, - ); + fn removes_arrow_exponential_histograms_and_preserves_other_metrics() { + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, ResourceMetrics, + ScopeMetrics, + }; - assert_eq!(7, status.batch_id); - assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); - assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + let mut request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![ + Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint::default()], + ..Default::default() + })), + ..Default::default() + }, + Metric { + data: Some(metric::Data::Gauge(Gauge::default())), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }], + }; + assert!(remove_exponential_histograms(&mut request)); + let metrics = &request.resource_metrics[0].scope_metrics[0].metrics; + assert_eq!(metrics.len(), 1); + assert!(matches!(metrics[0].data, Some(metric::Data::Gauge(_)))); + assert!(!remove_exponential_histograms(&mut request)); + } + + #[test] + fn batch_status_explains_arrow_exponential_histogram_limit() { + for accepted_data_points in [0, 1] { + let status = batch_status( + 7, + MetricsIngestOutcome { + accepted_data_points, + ..Default::default() + }, + true, + ); + assert_eq!(7, status.batch_id); + assert_eq!( + if accepted_data_points == 0 { + ArrowStatusCode::InvalidArgument + } else { + ArrowStatusCode::Ok + } as i32, + status.status_code + ); + assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + } } } diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 6b7bbd3fab1..7862af28db3 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -138,7 +138,7 @@ pub fn to_grpc_insert_requests( && !metric_ctx.is_legacy && let Some(r) = resource.resource.as_ref() { - resource_info.observe(&r.attributes, resource, metric_ctx); + resource_info.observe(&r.attributes, resource); } let resource_attrs = resource.resource.as_ref().map(|r| { @@ -630,7 +630,7 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { + if let Err(rejection) = exponential_histogram_gate(histogram) { reject_data_points(outcome, histogram.data_points.len(), || { rejection.message(name) })?; @@ -684,7 +684,6 @@ fn encode_exponential_histogram( } pub(crate) enum ExponentialHistogramRejection { - Disabled, DeltaTemporality, UnspecifiedTemporality, } @@ -692,9 +691,6 @@ pub(crate) enum ExponentialHistogramRejection { impl ExponentialHistogramRejection { fn message(&self, name: &str) -> String { match self { - Self::Disabled => format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ), Self::DeltaTemporality => format!( "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" ), @@ -710,11 +706,7 @@ impl ExponentialHistogramRejection { /// Individual points can still fail [`exponential_histogram_value`]. pub(crate) fn exponential_histogram_gate( histogram: &ExponentialHistogram, - metric_ctx: &OtlpMetricCtx, ) -> std::result::Result<(), ExponentialHistogramRejection> { - if !metric_ctx.experimental_enable_exponential_histogram { - return Err(ExponentialHistogramRejection::Disabled); - } match AggregationTemporality::try_from(histogram.aggregation_temporality) { Ok(AggregationTemporality::Cumulative) => Ok(()), Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), @@ -2498,7 +2490,7 @@ mod tests { } #[test] - fn test_exponential_histogram_gate_and_partial_outcome() { + fn test_exponential_histogram_default_context_and_partial_outcome() { let request = metrics_request(vec![ Metric { name: "temperature".to_string(), @@ -2512,6 +2504,11 @@ mod tests { vec![exponential_point()], AggregationTemporality::Cumulative, ), + exponential_metric( + "delta_latency", + vec![exponential_point()], + AggregationTemporality::Delta, + ), ]); let MetricsConversion { requests, @@ -2520,20 +2517,20 @@ mod tests { .. } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); - assert_eq!(outcome.accepted_data_points, 1); + assert_eq!(outcome.accepted_data_points, 2); assert_eq!(outcome.rejected_data_points, 1); assert!( outcome .error_message .as_deref() .unwrap() - .contains("otlp.experimental_enable_exponential_histogram") + .contains("only cumulative temporality is supported") ); - assert_eq!(requests.inserts.len(), 1); - assert_eq!(requests.inserts[0].table_name, "temperature"); + assert_eq!(requests.inserts.len(), 2); let semantics = decode(&semantic_index); assert!(semantics.contains_key("temperature")); - assert!(!semantics.contains_key("latency")); + assert!(semantics.contains_key("latency")); + assert!(!semantics.contains_key("delta_latency")); let empty = metrics_request(vec![exponential_metric( "empty", @@ -2566,10 +2563,7 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); assert!( @@ -2589,10 +2583,7 @@ mod tests { ), histogram_metric("latency"), ]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); assert!( @@ -2657,10 +2648,7 @@ mod tests { vec![stale.clone()], temporality, )]); - let mut ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut ctx = OtlpMetricCtx::default(); let MetricsConversion { requests, rows, @@ -2688,15 +2676,11 @@ mod tests { vec![point], AggregationTemporality::Cumulative, )]); - let mut new_ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; + let mut new_ctx = OtlpMetricCtx::default(); let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) .unwrap() .requests; let mut legacy_ctx = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; @@ -2749,7 +2733,7 @@ mod tests { let request = metrics_request(vec![exponential_metric( "x".repeat(1_000), vec![ExponentialHistogramDataPoint::default()], - AggregationTemporality::Cumulative, + AggregationTemporality::Delta, )]); let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) .unwrap() diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs index c87e95176c1..4b4f2afcad6 100644 --- a/src/servers/src/otlp/metrics/resource_info.rs +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -30,7 +30,6 @@ use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ AggregationTemporality, ResourceMetrics, metric, }; -use session::protocol_ctx::OtlpMetricCtx; use crate::error::Result; use crate::otlp::metrics::{ @@ -90,12 +89,7 @@ pub struct ResourceInfoData { impl ResourceInfoData { /// Takes the raw attributes, before the promote filter runs on them. - pub fn observe( - &mut self, - raw_attrs: &[KeyValue], - resource: &ResourceMetrics, - metric_ctx: &OtlpMetricCtx, - ) { + pub fn observe(&mut self, raw_attrs: &[KeyValue], resource: &ResourceMetrics) { let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); let ServiceIdentity { job, instance } = service_identity(raw_attrs); if let Some(job) = job { @@ -119,7 +113,7 @@ impl ResourceInfoData { tags.sort_unstable(); let mut observed: BTreeMap = BTreeMap::new(); - for_each_encoded_time(resource, metric_ctx, |ts| { + for_each_encoded_time(resource, |ts| { let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); observed .entry(window) @@ -176,11 +170,7 @@ impl ResourceInfoData { /// Visits the times of the data points the encoder writes rows for, so a /// resource is described exactly where it is measured rather than wherever /// its request happens to reach. -fn for_each_encoded_time( - resource: &ResourceMetrics, - metric_ctx: &OtlpMetricCtx, - mut visit: impl FnMut(i64), -) { +fn for_each_encoded_time(resource: &ResourceMetrics, mut visit: impl FnMut(i64)) { fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { for ts in points { visit(ts as i64); @@ -214,7 +204,7 @@ fn for_each_encoded_time( visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) } Some(metric::Data::ExponentialHistogram(h)) - if exponential_histogram_gate(h, metric_ctx).is_ok() => + if exponential_histogram_gate(h).is_ok() => { for point in &h.data_points { if let Ok((_, ts)) = exponential_histogram_value(point) { @@ -285,7 +275,7 @@ mod tests { kv("k8s.node.name", "node-a"), kv("os.type", "linux"), ]; - data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + data.observe(&attrs, &gauge_at(&[100, 50])); assert_eq!(data.rows.len(), 1); let (tags, windows) = data.rows.iter().next().unwrap(); assert_eq!(windows.values().copied().collect::>(), vec![100]); @@ -298,19 +288,11 @@ mod tests { .all(|(k, _)| k != "os.type" && k != "service.instance.id") ); - data.observe( - &[kv("host.id", "h-2")], - &gauge_at(&[10]), - &OtlpMetricCtx::default(), - ); + data.observe(&[kv("host.id", "h-2")], &gauge_at(&[10])); assert_eq!(data.rows.len(), 2); let mut empty = ResourceInfoData::default(); - empty.observe( - &[kv("os.type", "linux")], - &gauge_at(&[100]), - &OtlpMetricCtx::default(), - ); + empty.observe(&[kv("os.type", "linux")], &gauge_at(&[100])); assert!(empty.into_row_insert_requests().unwrap().is_none()); } @@ -322,7 +304,6 @@ mod tests { data.observe( &[kv("service.name", "api")], &gauge_at(&[window + 1, window + 2, 3 * window + 7]), - &OtlpMetricCtx::default(), ); let windows = data.rows.values().next().unwrap(); @@ -352,20 +333,13 @@ mod tests { }], ..Default::default() }; - let enabled = OtlpMetricCtx { - experimental_enable_exponential_histogram: true, - ..Default::default() - }; - - for (resource, ctx) in [ - ( - exponential(AggregationTemporality::Cumulative), - OtlpMetricCtx::default(), - ), - (exponential(AggregationTemporality::Delta), enabled), + for temporality in [ + AggregationTemporality::Delta, + AggregationTemporality::Unspecified, ] { + let resource = exponential(temporality); let mut data = ResourceInfoData::default(); - data.observe(&[kv("service.name", "api")], &resource, &ctx); + data.observe(&[kv("service.name", "api")], &resource); assert!(data.into_row_insert_requests().unwrap().is_none()); } } @@ -377,7 +351,6 @@ mod tests { data.observe( &[kv("service.name", "api"), kv("host.id", "h-1")], &gauge_at(&[1_700_000_000_123_456_789]), - &OtlpMetricCtx::default(), ); let requests = data.into_row_insert_requests().unwrap().unwrap(); assert_eq!(requests.inserts.len(), 1); diff --git a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs index 8fa86ae51e1..eb1de565c2f 100644 --- a/src/servers/src/otlp/metrics/resource_info/tests/delta.rs +++ b/src/servers/src/otlp/metrics/resource_info/tests/delta.rs @@ -56,11 +56,7 @@ fn observe_ignores_rejected_classic_histogram_windows() { ..Default::default() }; let mut data = ResourceInfoData::default(); - data.observe( - &[kv("service.name", "api")], - &resource, - &OtlpMetricCtx::default(), - ); + data.observe(&[kv("service.name", "api")], &resource); let windows = data.rows.values().next().unwrap(); assert_eq!( diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 2f1f4d410c4..77ebe5a7299 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -130,7 +130,6 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub(crate) fn decode_remote_write_v2( is_zstd: bool, body: Bytes, - native_histograms_enabled: bool, ) -> Result { let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) @@ -151,13 +150,10 @@ pub(crate) fn decode_remote_write_v2( let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["convert", REMOTE_WRITE_V2_VERSION]) .start_timer(); - convert_remote_write_v2(request, native_histograms_enabled) + convert_remote_write_v2(request) } -fn convert_remote_write_v2( - request: BorrowedRequest<'_>, - native_histograms_enabled: bool, -) -> Result { +fn convert_remote_write_v2(request: BorrowedRequest<'_>) -> Result { ensure!( request.symbols.first().copied() == Some(""), error::InvalidPromRemoteRequestSnafu { @@ -179,14 +175,6 @@ fn convert_remote_write_v2( let counts = scan_series(series, &mut labels_refs, &mut metadata) .context(error::DecodePromRemoteRequestSnafu)?; - ensure!( - native_histograms_enabled || counts.histograms == 0, - error::InvalidPromRemoteRequestSnafu { - msg: "prometheus remote write v2 native histogram ingestion is experimental; set prom_store.experimental_enable_prometheus_native_histogram = true to enable it" - .to_string(), - } - ); - if counts.samples == 0 && counts.histograms == 0 { decode_series_leaves(series, None, Vec::new(), 0, &mut scratch)?; continue; @@ -831,15 +819,14 @@ pub mod test_util { request: Request, ) -> Result<(Vec, Vec, u64, u64)> { let body = Bytes::from(snappy_compress(&request.encode_to_vec())?); - decode_write_requests(false, body, true) + decode_write_requests(false, body) } pub fn decode_write_requests( is_zstd: bool, body: Bytes, - native_histograms_enabled: bool, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::decode_remote_write_v2(is_zstd, body, native_histograms_enabled)?; + let requests = super::decode_remote_write_v2(is_zstd, body)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -850,11 +837,10 @@ pub mod test_util { pub fn decode_uncompressed_write_requests( body: &[u8], - native_histograms_enabled: bool, ) -> Result<(Vec, Vec, u64, u64)> { let request = super::BorrowedRequest::decode(body).context(error::DecodePromRemoteRequestSnafu)?; - let requests = super::convert_remote_write_v2(request, native_histograms_enabled)?; + let requests = super::convert_remote_write_v2(request)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -953,12 +939,7 @@ mod tests { assert_eq!(decoded.timeseries[0].samples.len(), 1); assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); - assert_eq!( - decode_remote_write_v2(true, body, true) - .unwrap() - .sample_count, - 1 - ); + assert_eq!(decode_remote_write_v2(true, body).unwrap().sample_count, 1); } #[test] @@ -990,7 +971,7 @@ mod tests { wire.extend(encoded_message_field(5, &packed_u32_field(1, &[99]))); wire.extend(string_field(4, b"http_requests_total")); - let requests = decode_wire(&wire, true).unwrap(); + let requests = decode_wire(&wire).unwrap(); assert_eq!(requests.sample_count, 2); assert_eq!(requests.histogram_count, 0); let rows = requests.samples.all_req().next().unwrap().rows.unwrap(); @@ -1028,7 +1009,7 @@ mod tests { series.extend(packed_u32_field(1, &[1, 2])); let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); - let requests = decode_wire(&wire, true).unwrap(); + let requests = decode_wire(&wire).unwrap(); assert_eq!(requests.histogram_count, 1); let rows = requests.histograms.all_req().next().unwrap().rows.unwrap(); assert_eq!( @@ -1093,7 +1074,7 @@ mod tests { ("invalid sample", invalid_sample), ("invalid histogram", invalid_histogram), ] { - let error = decode_wire_error(&wire, true, name); + let error = decode_wire_error(&wire, name); assert!( matches!(error, error::Error::DecodePromRemoteRequest { .. }), "{name}: {error}" @@ -1108,7 +1089,7 @@ mod tests { series.extend(encoded_message_field(tag, &[0x08])); let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); - let error = decode_wire_error(&wire, true, "malformed ignored message"); + let error = decode_wire_error(&wire, "malformed ignored message"); assert!(matches!( error, error::Error::DecodePromRemoteRequest { .. } @@ -1174,8 +1155,8 @@ mod tests { }) .unwrap(); assert_eq!(requests.sample_count, 0); - assert!(decode_wire(&[], true).is_err()); - assert!(decode_wire(&[0x0a, 0x00], true).is_err()); + assert!(decode_wire(&[]).is_err()); + assert!(decode_wire(&[0x0a, 0x00]).is_err()); } #[test] @@ -1205,22 +1186,34 @@ mod tests { } #[test] - fn test_fused_decoder_pins_experimental_error_precedence() { - let histogram = series_wire(&[1, 2], 3, &Histogram::default().encode_to_vec()); - let malformed_sample = series_wire(&[1, 2], 2, &[0x08]); + fn test_fused_decoder_rejects_malformed_series_with_histograms() { + let histogram = series_wire( + &[1, 2], + 3, + &Histogram { + count: Some(Count::CountInt(0)), + zero_count: Some(ZeroCount::ZeroCountInt(0)), + ..Default::default() + } + .encode_to_vec(), + ); + let malformed_sample = series_wire(&[1, 3], 2, &[0x08]); let wire = request_wire( - &["", METRIC_NAME_LABEL, "metric"], + &["", METRIC_NAME_LABEL, "histogram", "sample"], &[histogram.clone(), malformed_sample.clone()], ); - let error = decode_wire_error(&wire, false, "histogram before malformed series"); - assert!(error.to_string().contains("ingestion is experimental")); + let error = decode_wire_error(&wire, "histogram before malformed series"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); let wire = request_wire( - &["", METRIC_NAME_LABEL, "metric"], + &["", METRIC_NAME_LABEL, "histogram", "sample"], &[malformed_sample, histogram.clone()], ); - let error = decode_wire_error(&wire, false, "malformed series before histogram"); + let error = decode_wire_error(&wire, "malformed series before histogram"); assert!(matches!( error, error::Error::DecodePromRemoteRequest { .. } @@ -1231,7 +1224,7 @@ mod tests { &["", METRIC_NAME_LABEL, "metric", "job", "api"], &[missing_name, histogram], ); - let error = decode_wire_error(&wire, false, "conversion error before histogram"); + let error = decode_wire_error(&wire, "conversion error before histogram"); assert!(error.to_string().contains("missing '__name__'")); } @@ -2026,16 +2019,13 @@ mod tests { )); } - fn decode_wire( - wire: &[u8], - native_histograms_enabled: bool, - ) -> Result { + fn decode_wire(wire: &[u8]) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); - decode_remote_write_v2(false, body, native_histograms_enabled) + decode_remote_write_v2(false, body) } - fn decode_wire_error(wire: &[u8], native_histograms_enabled: bool, name: &str) -> error::Error { - match decode_wire(wire, native_histograms_enabled) { + fn decode_wire_error(wire: &[u8], name: &str) -> error::Error { + match decode_wire(wire) { Ok(_) => panic!("{name}: expected decoder error"), Err(error) => error, } @@ -2104,16 +2094,9 @@ mod tests { } fn decode_test_request(request: Request) -> Result { - decode_test_request_with_histograms(request, true) - } - - fn decode_test_request_with_histograms( - request: Request, - native_histograms_enabled: bool, - ) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - decode_remote_write_v2(false, body, native_histograms_enabled) + decode_remote_write_v2(false, body) } fn assert_invalid(name: &str, request: Request, expected: &str) { diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index ffd2c9dcf7f..c7a288ccee0 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -209,34 +209,10 @@ fn make_test_app_with_write_capture( make_test_app_with_write_failure(read_tx, write_tx, None) } -fn make_test_app_with_native_histogram_write_capture( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, None, true) -} - fn make_test_app_with_write_failure( read_tx: mpsc::Sender<(String, Vec)>, write_tx: mpsc::Sender, fail_write_call: Option, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, fail_write_call, false) -} - -fn make_test_app_with_native_histogram_write_failure( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, - fail_write_call: Option, -) -> Router { - make_test_app_with_write_failure_inner(read_tx, write_tx, fail_write_call, true) -} - -fn make_test_app_with_write_failure_inner( - read_tx: mpsc::Sender<(String, Vec)>, - write_tx: mpsc::Sender, - fail_write_call: Option, - experimental_enable_prometheus_native_histogram: bool, ) -> Router { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -262,14 +238,7 @@ fn make_test_app_with_write_failure_inner( BatchingProtocol::HttpSql, ]) .with_sql_handler(instance.clone()) - .with_prom_handler( - instance, - None, - true, - PromValidationMode::Unchecked, - experimental_enable_prometheus_native_histogram, - None, - ) + .with_prom_handler(instance, None, true, PromValidationMode::Unchecked, None) .build(); server.build(server.make_app()).unwrap() } @@ -418,7 +387,7 @@ async fn test_prometheus_remote_write_v2_histogram_write_error_has_partial_writt let (read_tx, _read_rx) = mpsc::channel(100); let (write_tx, mut write_rx) = mpsc::channel(100); - let app = make_test_app_with_native_histogram_write_failure(read_tx, write_tx, Some(2)); + let app = make_test_app_with_write_failure(read_tx, write_tx, Some(2)); let client = TestClient::new(app).await; let mut write_request = remote_write_v2::request_with_labels_and_samples( @@ -671,7 +640,7 @@ async fn test_prometheus_remote_write_v2_writes_histogram_only_series() { let (read_tx, _read_rx) = mpsc::channel(100); let (write_tx, mut write_rx) = mpsc::channel(100); - let app = make_test_app_with_native_histogram_write_capture(read_tx, write_tx); + let app = make_test_app_with_write_capture(read_tx, write_tx); let client = TestClient::new(app).await; let write_request = remote_write_v2::request_with_labels_and_histograms( @@ -708,36 +677,6 @@ async fn test_prometheus_remote_write_v2_writes_histogram_only_series() { assert!(write_rx.try_recv().is_err()); } -#[tokio::test] -async fn test_prometheus_remote_write_v2_rejects_native_histogram_when_disabled() { - common_telemetry::init_default_ut_logging(); - let (read_tx, _read_rx) = mpsc::channel(100); - let (write_tx, mut write_rx) = mpsc::channel(100); - - let app = make_test_app_with_write_capture(read_tx, write_tx); - let client = TestClient::new(app).await; - - let write_request = remote_write_v2::request_with_labels_and_histograms( - vec![( - prom_store::METRIC_NAME_LABEL, - "http_request_duration_seconds", - )], - vec![remote_write_v2::histogram(1000)], - ); - - let result = post_remote_write_v2(&client, &write_request).await; - - assert_eq!(result.status(), 400); - assert_remote_write_v2_written_headers_with_histograms(&result.headers(), "0", "0"); - assert!( - result - .text() - .await - .contains("prom_store.experimental_enable_prometheus_native_histogram") - ); - assert!(write_rx.try_recv().is_err()); -} - #[tokio::test] async fn test_prometheus_remote_write_rejects_unsupported_proto() { common_telemetry::init_default_ut_logging(); @@ -873,7 +812,6 @@ async fn test_prom_batching_depends_on_protocol_and_metric_engine() { None, with_metric_engine, PromValidationMode::Unchecked, - false, None, ) .build(); diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index 75fee046f7c..7855f92b814 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -77,7 +77,7 @@ fn test_decode_remote_write_v2_native_histogram_dump() { assert_eq!(histogram.timestamp, 1782358160412); let (sample_inserts, histogram_inserts, sample_count, histogram_count) = - remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY), true).unwrap(); + remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY)).unwrap(); assert!(sample_inserts.is_empty()); assert_eq!(sample_count, 0); assert_eq!(histogram_count, 1); diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index d9a4a386f91..5165def3e49 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -42,7 +42,6 @@ impl ProtocolCtx { /// If true, all scope attributes will be promoted to the final table schema. /// Along with the scope name, scope version and scope schema URL. /// - `with_metric_engine` -/// - `experimental_enable_exponential_histogram` /// - `is_legacy` /// If the user uses OTLP metrics ingestion before v0.16, it uses the old path. /// So we call this path 'legacy'. @@ -54,7 +53,6 @@ pub struct OtlpMetricCtx { pub resource_attrs: HashSet, pub promote_scope_attrs: bool, pub with_metric_engine: bool, - pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, /// Set from the server's `otlp.experimental_enable_resource_info`; off /// means the resource descriptor is not synthesized at all. diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 9fbf80fbeea..67a4e512f7d 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -535,8 +535,7 @@ WITH( .build() .await; let instance = standalone.fe_instance(); - let mut options = standalone.opts.clone(); - options.otlp.experimental_enable_exponential_histogram = true; + let options = standalone.opts.clone(); let services = Services::new(options.clone(), instance.clone(), Plugins::default()); let server = services .http_server_builder( diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 37a1740ccd1..622d3ba6c16 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -701,7 +701,7 @@ pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true, false) + .with_otlp_handler(instance.fe_instance().clone(), true) .with_jaeger_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()) .build(); @@ -721,18 +721,6 @@ pub async fn setup_test_http_app_with_frontend_and_user_provider( user_provider, None, None, - false, - ) - .await -} - -pub async fn setup_test_http_app_with_otlp_exponential_histogram( - store_type: StorageType, - name: &str, - enabled: bool, -) -> (Router, TestGuard) { - setup_test_http_app_with_frontend_and_custom_options( - store_type, name, None, None, None, enabled, ) .await } @@ -743,7 +731,6 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( user_provider: Option, http_opts: Option, memory_limiter: Option, - experimental_enable_exponential_histogram: bool, ) -> (Router, TestGuard) { let plugins = Plugins::new(); if let Some(user_provider) = user_provider.clone() { @@ -765,11 +752,7 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler( - instance.fe_instance().clone(), - true, - experimental_enable_exponential_histogram, - ) + .with_otlp_handler(instance.fe_instance().clone(), true) .with_prometheus_handler(instance.fe_instance().clone()) .with_jaeger_handler(instance.fe_instance().clone()) .with_dashboard_handler(instance.fe_instance().clone()) @@ -801,14 +784,7 @@ pub async fn setup_test_prom_app_with_frontend( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false, false).await -} - -pub async fn setup_test_prom_app_with_frontend_native_histogram( - store_type: StorageType, - name: &str, -) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false, true).await + setup_test_prom_app_with_frontend_inner(store_type, name, false).await } /// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher, @@ -818,14 +794,13 @@ pub async fn setup_test_prom_app_with_frontend_batched( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, true, false).await + setup_test_prom_app_with_frontend_inner(store_type, name, true).await } async fn setup_test_prom_app_with_frontend_inner( store_type: StorageType, name: &str, enable_batcher: bool, - experimental_enable_prometheus_native_histogram: bool, ) -> (Router, TestGuard) { unsafe { std::env::set_var("TZ", "UTC"); @@ -874,13 +849,9 @@ async fn setup_test_prom_app_with_frontend_inner( let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)"; run_sql(sql, &instance).await; - let http_server = build_test_prom_server( - instance.fe_instance().clone(), - enable_batcher, - experimental_enable_prometheus_native_histogram, - ) - .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) - .build(); + let http_server = build_test_prom_server(instance.fe_instance().clone(), enable_batcher) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); let app = http_server.build(http_server.make_app()).unwrap(); (app, instance.guard) } @@ -889,7 +860,6 @@ async fn setup_test_prom_app_with_frontend_inner( pub fn build_test_prom_server( frontend_ref: Arc, enable_batcher: bool, - experimental_enable_prometheus_native_histogram: bool, ) -> HttpServerBuilder { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -924,7 +894,6 @@ pub fn build_test_prom_server( Some(frontend_ref.clone()), true, PromValidationMode::Strict, - experimental_enable_prometheus_native_histogram, pending_rows_batcher, ) .with_prometheus_handler(frontend_ref) @@ -1279,7 +1248,7 @@ pub async fn setup_pg_server_with_prom_native_histogram( let instance = setup_standalone_instance(name, store_type).await; // Prometheus remote-write HTTP app with native histograms enabled. - let http_server = build_test_prom_server(instance.fe_instance().clone(), false, true) + let http_server = build_test_prom_server(instance.fe_instance().clone(), false) .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) .build(); let app = http_server.build(http_server.make_app()).unwrap(); diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index caac6d33238..4b78e5fe276 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -3312,3 +3312,166 @@ CREATE TABLE b ( +-------+-----------------------------------+"#; check_output_stream(output, expected).await; } + +#[rstest] +#[case::mito(false)] +#[case::batched_metric(true)] +#[tokio::test(flavor = "multi_thread")] +async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram, Sample}; + use axum::body::{Body, to_bytes}; + use http::Request; + use prost::Message; + use servers::http::{HttpOptions, HttpServerBuilder}; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_remote_write::validation::PromValidationMode; + use servers::prom_store::snappy_compress; + use tower::ServiceExt; + + use crate::cluster::GreptimeDbClusterBuilder; + use crate::test_util::build_test_prom_server; + use crate::tests::test_util::{MockInstanceBuilder, RebuildableMockInstance, TestContext}; + + common_telemetry::init_default_ut_logging(); + let builder = MockInstanceBuilder::Distributed( + GreptimeDbClusterBuilder::new(&format!("histogram_lifecycle_{metric_engine}")) + .await + .with_datanodes(3), + ); + let mut context = TestContext::new(builder).await; + let make_router = |frontend: Arc| { + let builder = if metric_engine { + build_test_prom_server(frontend.clone(), true) + } else { + HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(frontend.clone()) + .with_prometheus_handler(frontend.clone()) + .with_prom_handler( + frontend.clone(), + None, + false, + PromValidationMode::Strict, + None, + ) + }; + let server = builder.build(); + server.build(server.make_app()).unwrap() + }; + let snapshot = |frontend: Arc| async move { + let mut result = Vec::new(); + for table in ["lifecycle_sample", "lifecycle_histogram"] { + let output = execute_sql( + &frontend, + &format!("select * from {table} order by greptime_timestamp"), + ) + .await; + let OutputData::Stream(stream) = output.data else { + panic!("expected query stream") + }; + result.push( + util::collect_batches(stream) + .await + .unwrap() + .pretty_print() + .unwrap(), + ); + } + result + }; + let mut router = make_router(context.frontend()); + for round in 0..2 { + let timestamp = 1_000 + round * 10_000; + let mut labels = vec![("__name__", "lifecycle_histogram"), ("job", "api")]; + if round == 1 { + // Exercise schema evolution after the first SST already exists. + labels.push(("zone", "east")); + } + let sample = remote_write_v2::request_with_labels_and_samples( + vec![("__name__", "lifecycle_sample"), ("job", "api")], + vec![Sample { + value: 42.0, + timestamp, + ..Default::default() + }], + ); + let histograms = remote_write_v2::request_with_labels_and_histograms( + labels, + vec![Histogram { + timestamp, + start_timestamp: 500, + count: Some(Count::CountInt(9_007_199_254_740_993)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + zero_threshold: 0.001, + sum: 12.5, + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![9_007_199_254_740_992], + ..Default::default() + }], + ); + // Create the scalar table first, then add histogram storage to its physical table. + for request in [sample, histograms] { + let response = router + .clone() + .oneshot( + Request::post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(Body::from( + snappy_compress(&request.encode_to_vec()).unwrap(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + 204, + "{:?}", + to_bytes(response.into_body(), usize::MAX).await.unwrap() + ); + } + + let expected = snapshot(context.frontend()).await; + assert!(expected[1].contains("9007199254740993"), "{}", expected[1]); + // Rebuild before flushing, including after adding fields to the physical + // schema: recovery must retain the columns filled into bulk writes. + drop(router); + context.rebuild().await; + router = make_router(context.frontend()); + assert_eq!(snapshot(context.frontend()).await, expected, "WAL recovery"); + let storage_tables = if metric_engine { + vec!["greptime_physical_table"] + } else { + vec!["lifecycle_sample", "lifecycle_histogram"] + }; + for table in &storage_tables { + execute_sql( + &context.frontend(), + &format!("admin flush_table('{table}')"), + ) + .await; + } + assert_eq!(snapshot(context.frontend()).await, expected, "SST read"); + if round == 1 { + for table in &storage_tables { + execute_sql( + &context.frontend(), + &format!("admin compact_table('{table}', 'strict_window', 'window=3600')"), + ) + .await; + } + assert_eq!(snapshot(context.frontend()).await, expected, "compaction"); + drop(router); + context.rebuild().await; + router = make_router(context.frontend()); + assert_eq!(snapshot(context.frontend()).await, expected, "SST reopen"); + } + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 4cc409deec4..3e64a6d74e3 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -85,7 +85,7 @@ use tests_integration::test_util::{ MockInstanceImpl, StorageType, assert_wal_delta, build_test_prom_server, setup_test_http_app, setup_test_http_app_with_frontend, setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, - setup_test_prom_app_with_frontend_batched, setup_test_prom_app_with_frontend_native_histogram, + setup_test_prom_app_with_frontend_batched, }; use urlencoding::encode; use yaml_rust::YamlLoader; @@ -2492,7 +2492,6 @@ enable = true [otlp] enable = true -experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 experimental_enable_resource_info = true @@ -2500,7 +2499,6 @@ experimental_enable_resource_info = true enable = true with_metric_engine = true prom_validation_mode = "strict" -experimental_enable_prometheus_native_histogram = false pending_rows_flush_interval = "0s" max_batch_rows = 100000 max_concurrent_flushes = 256 @@ -3152,7 +3150,7 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: StorageType) { common_telemetry::init_default_ut_logging(); - let (app, mut guard) = setup_test_prom_app_with_frontend_native_histogram( + let (app, mut guard) = setup_test_prom_app_with_frontend( store_type, "prometheus_remote_write_v2_native_histogram", ) @@ -3532,7 +3530,7 @@ async fn check_prometheus_remote_write_batched_skip_wal(distributed: bool, v2: b common_telemetry::init_default_ut_logging(); let mut instance = MockInstanceImpl::new(&format!("prom_bulk_skip_wal_v2_{v2}"), distributed).await; - let server = build_test_prom_server(instance.frontend(), true, false).build(); + let server = build_test_prom_server(instance.frontend(), true).build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; write_prometheus_skip_wal_sample(&client, v2, 1000, None).await; @@ -7464,7 +7462,6 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, ResourceMetrics, ScopeMetrics, exponential_histogram_data_point, metric, }; - use tests_integration::test_util::setup_test_http_app_with_otlp_exponential_histogram; common_telemetry::init_default_ut_logging(); let req = ExportMetricsServiceRequest { @@ -7511,44 +7508,8 @@ pub async fn test_otlp_exponential_histogram(store_type: StorageType) { )] }; - let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( - store_type, - "test_otlp_exponential_histogram_disabled", - false, - ) - .await; - let client = TestClient::new(app).await; - let res = send_req( - &client, - headers(), - "/v1/otlp/v1/metrics", - body.clone(), - false, - ) - .await; - assert_eq!(StatusCode::BAD_REQUEST, res.status()); - let status = GoogleRpcStatus::decode(res.bytes().await.as_ref()).unwrap(); - assert_eq!(3, status.code); - assert!( - status - .message - .contains("otlp.experimental_enable_exponential_histogram") - ); - validate_data( - "otlp_exponential_histogram_disabled_no_table", - &client, - "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency';", - "[[0]]", - ) - .await; - guard.remove_all().await; - - let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( - store_type, - "test_otlp_exponential_histogram_enabled", - true, - ) - .await; + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_exponential_histogram").await; let client = TestClient::new(app).await; let res = send_req(&client, headers(), "/v1/otlp/v1/metrics", body, false).await; assert_eq!(StatusCode::OK, res.status()); @@ -12041,16 +12002,9 @@ async fn check_http_skip_wal(name: &str, cases: &[HttpWalCase], distributed: boo .with_influxdb_handler(fe.clone()) .with_opentsdb_handler(fe.clone()) .with_log_ingest_handler(fe.clone(), None, None) - .with_otlp_handler(fe.clone(), true, false) + .with_otlp_handler(fe.clone(), true) // The pending batcher uses BulkInsert, deliberately outside this PR. - .with_prom_handler( - fe.clone(), - Some(fe), - true, - PromValidationMode::Strict, - false, - None, - ) + .with_prom_handler(fe.clone(), Some(fe), true, PromValidationMode::Strict, None) .build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; for case in cases { @@ -12395,7 +12349,6 @@ pub async fn test_http_memory_limit(store_type: StorageType) { None, Some(http_opts), Some(memory_limiter), - false, ) .await; From 1955955b2dd6e511a80f87940e3f1cdc589c42b9 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 23 Sep 2026 13:16:26 +0000 Subject: [PATCH 026/133] fix(ci): grant PR write permission for CI command replies (#9331) Signed-off-by: WenyXu --- .github/workflows/ci-slash.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci-slash.yml b/.github/workflows/ci-slash.yml index e9fb392d381..059654fd5cd 100644 --- a/.github/workflows/ci-slash.yml +++ b/.github/workflows/ci-slash.yml @@ -8,7 +8,7 @@ permissions: actions: write contents: read issues: write - pull-requests: read + pull-requests: write jobs: admit: @@ -64,7 +64,7 @@ jobs: PR_NUMBER: ${{ steps.admit.outputs.pr_number || github.event.client_payload.github.payload.issue.number }} ADMIT_OUTCOME: ${{ steps.admit.outcome }} DISPATCH_OUTCOME: ${{ steps.dispatch.outcome }} - # Use the issue-comment endpoint with the existing issues: write permission. + # PR replies require pull-requests: write even through the issue-comment endpoint. run: | [[ "${PR_NUMBER}" =~ ^[0-9]+$ ]] || exit 1 if [[ "${ADMIT_OUTCOME}" == failure || "${DISPATCH_OUTCOME}" == failure ]]; then From c7fa48ef95bc5da1f245d9e292ee53b28a4f2136 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 23 Sep 2026 14:18:42 +0000 Subject: [PATCH 027/133] feat(mito): limit approximate series index disk usage (#9313) * feat(mito): limit series index disk usage Signed-off-by: evenyag * refactor(mito): enforce series index quota during reconciliation Signed-off-by: evenyag * refactor(mito): remove series index disk budget layer Signed-off-by: evenyag * refactor(mito): simplify series index limit to estimated usage Signed-off-by: evenyag * refactor(mito): trust index catalogs when loading snapshots Signed-off-by: evenyag * refactor(mito): minimize series index disk limit changes Signed-off-by: evenyag * fix(mito2): keep series index cleanup running at capacity Signed-off-by: evenyag * fix(mito2): avoid no-op index clones and stabilize capacity tests Signed-off-by: evenyag * fix: update config API expectation and stabilize index build tests Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 2 + config/datanode.example.toml | 7 + config/standalone.example.toml | 7 + src/cmd/tests/load_config_test.rs | 2 + src/mito2/AGENTS.md | 6 +- src/mito2/src/config.rs | 22 ++ src/mito2/src/engine/index_build_test.rs | 40 --- src/mito2/src/engine/listener.rs | 8 +- src/mito2/src/engine/scan_test.rs | 11 +- src/mito2/src/metrics.rs | 8 + src/mito2/src/read/series_candidate.rs | 1 + src/mito2/src/series_index.rs | 2 +- src/mito2/src/series_index/bucket.rs | 1 + src/mito2/src/series_index/builder.rs | 51 +++- src/mito2/src/series_index/catalog.rs | 34 ++- src/mito2/src/series_index/maintenance.rs | 173 ++++++++++--- src/mito2/src/series_index/searcher.rs | 1 + src/mito2/src/series_index/task.rs | 287 +++++++++++++++++++++- src/mito2/src/series_index/tests.rs | 29 ++- src/mito2/src/series_index/version.rs | 21 +- src/mito2/src/sst/parquet/reader.rs | 2 +- src/mito2/src/worker.rs | 11 +- tests-integration/tests/http.rs | 1 + 23 files changed, 616 insertions(+), 111 deletions(-) diff --git a/config/config.md b/config/config.md index 6544a28f93d..e380f075548 100644 --- a/config/config.md +++ b/config/config.md @@ -177,6 +177,7 @@ | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | @@ -639,6 +640,7 @@ | `region_engine.mito.experimental_manifest_keep_removed_file_ttl` | String | `1h` | How long to keep removed files in the `removed_files` field of manifest
after they are removed from manifest.
files will only be removed from `removed_files` field
if both `keep_removed_file_count` and `keep_removed_file_ttl` is reached. | | `region_engine.mito.compress_manifest` | Bool | `false` | Whether to compress manifest and checkpoint file by gzip (default false). | | `region_engine.mito.experimental_enable_series_index` | Bool | `false` | Under development; do not enable. Whether to enable series indexes.
Indexes are stored on the local filesystem under `{data_home}/series_index`. | +| `region_engine.mito.experimental_series_index_max_size` | String | `5GiB` | Approximate series and range index size limit in open regions, shared across workers.
Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation
can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots
retained by readers are not counted.
Minimum: 1KiB. Takes effect on restart. | | `region_engine.mito.experimental_enable_range_index` | Bool | `false` | Whether to build and query range indexes when series indexes are enabled.
Obsolete range-index metadata and files are still cleaned up when disabled. | | `region_engine.mito.experimental_series_index_maintenance_interval` | String | `5m` | Interval between series-index maintenance runs. Zero uses the default of 5 min. | | `region_engine.mito.experimental_series_index_bucket_width` | String | `5days` | Requested minimum series-index bucket width (default: 5 days), rounded up to
an exact multiple of each region's compaction time window. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 45baaefbdfb..da2d353fef8 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -502,6 +502,13 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Approximate series and range index size limit in open regions, shared across workers. +## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation +## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots +## retained by readers are not counted. +## Minimum: 1KiB. Takes effect on restart. +#+ experimental_series_index_max_size = "5GiB" + ## Whether to build and query range indexes when series indexes are enabled. ## Obsolete range-index metadata and files are still cleaned up when disabled. #+ experimental_enable_range_index = false diff --git a/config/standalone.example.toml b/config/standalone.example.toml index ef9dd561f6a..e0a8deb9c22 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -707,6 +707,13 @@ compress_manifest = false ## Indexes are stored on the local filesystem under `{data_home}/series_index`. #+ experimental_enable_series_index = false +## Approximate series and range index size limit in open regions, shared across workers. +## Workers periodically refresh usage and skip maintenance when full. In-flight reconciliation +## can exceed the limit. Closed-region files, temporary output, catalogs, and old snapshots +## retained by readers are not counted. +## Minimum: 1KiB. Takes effect on restart. +#+ experimental_series_index_max_size = "5GiB" + ## Whether to build and query range indexes when series indexes are enabled. ## Obsolete range-index metadata and files are still cleaned up when disabled. #+ experimental_enable_range_index = false diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 34c5bc8c9ba..d4b19812df2 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -159,6 +159,7 @@ fn test_load_datanode_example_config() { }, region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { + experimental_series_index_max_size: ReadableSize::gb(5), auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), @@ -397,6 +398,7 @@ fn test_load_standalone_example_config() { }), region_engine: vec![ RegionEngineConfig::Mito(MitoConfig { + experimental_series_index_max_size: ReadableSize::gb(5), auto_flush_interval: Duration::from_secs(10 * 60), default_region_write_buffer_size: ReadableSize::mb(0), write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)), diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 583ec31bd4a..b8fd2f8892c 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned maintenance and purge tasks, and index lifecycle | +| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning, catalogs, worker-owned reconciliation (`maintenance.rs`), shared approximate usage (`task.rs`), catalogs, and index lifecycle | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | @@ -62,6 +62,10 @@ filtered `RecordBatch` stream. - **Manifest format** (`manifest/action.rs`): affects crash recovery and follower replay. Keep it backward compatible. - **SST/Parquet layout** (`sst/`): readers must stay compatible with existing files. +- **Series-index disk limit**: tasks estimate usage from entry sizes in open-region snapshots + and defer new builds when the shared estimate is full. Cleanup still prunes obsolete + range entries and expires series indexes to reclaim published usage. Concurrent builds may + overshoot; closed-region files are not counted. Series handles and the SST purger own file deletion. - **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. - **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`. diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 5164e2fd72b..035feb75a15 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -92,6 +92,10 @@ pub struct MitoConfig { /// Under development; do not enable. Whether to enable series indexes (default false). /// Indexes are stored on the local filesystem under `{data_home}/series_index`. pub experimental_enable_series_index: bool, + /// Approximate series and range index size limit in open regions (default: 5 GiB). + /// Workers share a periodically refreshed estimate and skip maintenance when full. + /// In-flight reconciliation can exceed the limit; closed-region files are not counted. + pub experimental_series_index_max_size: ReadableSize, /// Whether to build and query range indexes when series indexes are enabled (default false). /// Obsolete range-index metadata and files are still cleaned up when disabled. pub experimental_enable_range_index: bool, @@ -221,6 +225,7 @@ impl Default for MitoConfig { compress_manifest: false, max_background_index_builds: divide_num_cpus(8), experimental_enable_series_index: false, + experimental_series_index_max_size: ReadableSize::gb(5), experimental_enable_range_index: false, experimental_series_index_maintenance_interval: DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL, @@ -280,6 +285,14 @@ impl MitoConfig { /// /// Returns an error if there is a configuration that unable to sanitize. pub fn sanitize(&mut self, data_home: &str) -> Result<()> { + if self.experimental_enable_series_index { + snafu::ensure!( + self.experimental_series_index_max_size.as_bytes() >= 1024, + crate::error::InvalidConfigSnafu { + reason: "experimental_series_index_max_size must be at least 1KiB" + } + ); + } // Use default value if `num_workers` is 0. if self.num_workers == 0 { self.num_workers = divide_num_cpus(2); @@ -442,12 +455,17 @@ mod tests { let mut config: MitoConfig = toml::from_str( "experimental_enable_series_index = true experimental_enable_range_index = false + experimental_series_index_max_size = '64MiB' experimental_series_index_maintenance_interval = '30s' experimental_series_index_bucket_width = '2days'", ) .unwrap(); config.sanitize("/data").unwrap(); assert!(config.experimental_enable_series_index); + assert_eq!( + ReadableSize::mb(64), + config.experimental_series_index_max_size + ); assert!(!config.experimental_enable_range_index); assert_eq!( config.experimental_series_index_maintenance_interval, @@ -459,6 +477,10 @@ mod tests { ); let restored: MitoConfig = toml::from_str(&toml::to_string(&config).unwrap()).unwrap(); assert_eq!(config, restored); + config.experimental_series_index_max_size = ReadableSize(1023); + assert!(config.sanitize("/data").is_err()); + config.experimental_series_index_max_size = ReadableSize(1024); + config.sanitize("/data").unwrap(); let mut config: MitoConfig = toml::from_str("experimental_series_index_maintenance_interval = '0s'").unwrap(); diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index ce950f8a246..d2d59b038ff 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -1512,46 +1512,6 @@ async fn test_index_build_type_manual_consistency() { assert_listener_counts(&listener, 2, 2); } -#[tokio::test] -async fn test_gate_index_build_listener_smoke() { - use store_api::storage::{FileId, RegionId}; - - use crate::engine::listener::{EventListener, GateIndexBuildListener}; - use crate::sst::file::RegionFileId; - - let gate = Arc::new(GateIndexBuildListener::default()); - - // Initial counts are zero. - assert_eq!(gate.begin_count(), 0); - assert_eq!(gate.finish_count(), 0); - assert_eq!(gate.abort_count(), 0); - - // Spawn a task that will block in on_index_build_begin. - let gate_clone = gate.clone(); - let handle = tokio::spawn(async move { - gate_clone - .on_index_build_begin(RegionFileId::new(RegionId::new(1, 1), FileId::random())) - .await; - }); - - // Wait for begin to arrive. - tokio::time::timeout(std::time::Duration::from_secs(5), gate.wait_begin(1)) - .await - .unwrap(); - assert_eq!(gate.begin_count(), 1); - assert_eq!(gate.finish_count(), 0); - assert_eq!(gate.abort_count(), 0); - - // Release the blocked begin. - gate.release_begin(); - - // The spawned task should now complete. - tokio::time::timeout(std::time::Duration::from_secs(5), handle) - .await - .unwrap() - .unwrap(); -} - #[tokio::test] async fn test_index_build_type_manual_duplicate_in_flight() { let mut env = TestEnv::with_prefix("test_index_build_type_manual_duplicate_in_flight_").await; diff --git a/src/mito2/src/engine/listener.rs b/src/mito2/src/engine/listener.rs index 6ef4a293a2a..432e3367603 100644 --- a/src/mito2/src/engine/listener.rs +++ b/src/mito2/src/engine/listener.rs @@ -583,12 +583,12 @@ impl EventListener for GateIndexBuildListener { info!("Region {} index build begin (gated)", region_file_id); self.begin_count.fetch_add(1, Ordering::Relaxed); self.begin_notify.notify_one(); - // Block until the test releases the gate. - let _permit = self - .begin_blocker + // Consume the release so subsequent builds remain blocked. + self.begin_blocker .acquire() .await - .expect("gate semaphore should not be closed"); + .expect("gate semaphore should not be closed") + .forget(); } async fn on_index_build_finish(&self, region_file_id: RegionFileId) { diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index bf089737f35..e10968d3a32 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -1340,6 +1340,7 @@ async fn check_two_phase_series_scan( let store = region.series_index_store.clone().unwrap(); let sequence = region.flushed_sequence(); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), bucket_end: Timestamp::new_millisecond(2001), @@ -1412,7 +1413,13 @@ async fn check_two_phase_series_scan( .unwrap(); writer.write(0, &batch).await.unwrap(); writer.finish().await.unwrap(); - index_version.range_indexes.insert(file_id); + index_version.range_indexes.insert( + file_id, + crate::series_index::RangeIndexEntry { + file_id, + file_size: 0, + }, + ); } region .series_index_version_control @@ -1607,7 +1614,7 @@ async fn check_two_phase_series_scan( // succeed even when that file is unavailable. let region = engine.find_region(region_id).unwrap(); let version = region.series_index_version_control.current(); - let file_id = *version.range_indexes.iter().next().unwrap(); + let file_id = *version.range_indexes.keys().next().unwrap(); region .series_index_store .as_ref() diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index 49e91d93efd..4c275c248c9 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -318,6 +318,14 @@ lazy_static! { // Index metrics. lazy_static! { + /// Approximate published index bytes in open regions, refreshed by maintenance. + pub static ref SERIES_INDEX_DISK_BYTES: IntGauge = register_int_gauge!( + "greptime_mito_series_index_disk_bytes", "estimated series and range index bytes in open regions" + ).unwrap(); + /// Maintenance passes deferred by the estimated disk usage. + pub static ref SERIES_INDEX_CAPACITY_DEFERRED: IntCounter = register_int_counter!( + "greptime_mito_series_index_capacity_deferred_total", "series-index capacity deferrals" + ).unwrap(); // Index metrics. /// Outcomes of series-index reconciliation passes. pub static ref SERIES_INDEX_RECONCILE_TOTAL: IntCounterVec = register_int_counter_vec!( diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index f05f54e6662..3ba5512a7bd 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -778,6 +778,7 @@ mod tests { )); let store = env.access_layer.object_store().clone(); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_millisecond(0), bucket_end: Timestamp::new_millisecond(20), diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 4a4133e1a9b..3aa70256709 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -39,7 +39,7 @@ use store_api::metric_engine_consts::{ use crate::error::Result; #[cfg(test)] -pub(crate) use crate::series_index::catalog::SeriesIndexEntry; +pub(crate) use crate::series_index::catalog::{RangeIndexEntry, SeriesIndexEntry}; pub(crate) use crate::series_index::catalog::{ delete_catalogs, load_version_control, series_index_path, }; diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs index 7fecb253441..446948838f6 100644 --- a/src/mito2/src/series_index/bucket.rs +++ b/src/mito2/src/series_index/bucket.rs @@ -334,6 +334,7 @@ impl SeriesBucket { .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) .min()?; Some(SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: self.start, bucket_end: self.end, diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs index 8d361be364e..63a51badaeb 100644 --- a/src/mito2/src/series_index/builder.rs +++ b/src/mito2/src/series_index/builder.rs @@ -21,7 +21,6 @@ use common_telemetry::warn; use futures::TryStreamExt; use object_store::ObjectStore; use snafu::{OptionExt, ensure}; -use store_api::storage::FileId; use crate::error::{Result, UnexpectedSnafu}; use crate::read::BoxedRecordBatchStream; @@ -34,7 +33,7 @@ use crate::region::MitoRegionRef; use crate::region::version::VersionRef; use crate::series_index::bucket::SeriesBucket; use crate::series_index::catalog::{ - SeriesIndexEntry, range_index_path, series_index_path, series_metadata, + RangeIndexEntry, SeriesIndexEntry, range_index_path, series_index_path, series_metadata, }; use crate::series_index::purger::{IndexFilePurger, IndexFileType, file_operation}; use crate::series_index::version::SeriesIndexFileHandle; @@ -68,7 +67,7 @@ pub(crate) async fn build_range_index( region: &MitoRegionRef, version: &VersionRef, file: FileHandle, -) -> Result> { +) -> Result> { let file_id = file.file_id().file_id(); let Some((context, mut selection)) = reader_input(region, file).await? else { return Ok(None); @@ -116,9 +115,12 @@ pub(crate) async fn build_range_index( } return Err(error); } - writer.finish().await?; + let metrics = writer.finish().await?; file_operation(IndexFileType::Range, "build", "success"); - Ok(Some(file_id)) + Ok(Some(RangeIndexEntry { + file_id, + file_size: metrics.output_bytes, + })) } /// Builds only the series index. Callers build needed range indexes separately. @@ -200,11 +202,14 @@ pub(crate) async fn build_series_index( } return Err(error); } - writer.finish().await?; + let metrics = writer.finish().await?; file_operation(IndexFileType::Series, "build", "success"); Ok(SeriesIndexFileHandle::new( region.region_id, - entry.clone(), + SeriesIndexEntry { + file_size: metrics.output_bytes, + ..entry.clone() + }, purger.clone(), )) } @@ -275,10 +280,19 @@ mod tests { let mut range_bytes = HashMap::new(); if build_ranges { for file in files { - let id = build_range_index(&store, ®ion, &version, file.clone()) + let completed = build_range_index(&store, ®ion, &version, file.clone()) .await .unwrap() .unwrap(); + let id = completed.file_id; + assert_eq!( + completed.file_size, + store + .stat(&range_index_path(region.region_id, id)) + .await + .unwrap() + .content_length() + ); range_bytes.insert( id, store @@ -314,7 +328,21 @@ mod tests { assert!(result.is_err()); } else { let handle = result.unwrap(); - assert_eq!(handle.entry(), &entry); + assert_eq!( + handle.entry().file_size, + store + .stat(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + .content_length() + ); + assert_eq!( + handle.entry(), + &SeriesIndexEntry { + file_size: handle.entry().file_size, + ..entry.clone() + } + ); assert!( store .exists(&series_index_path(region.region_id, entry.index_uuid)) @@ -504,12 +532,13 @@ mod tests { let mut completed = Vec::new(); let result: Result> = async { for file in files { - let Some(id) = build_range_index(&store, ®ion, &version, file.clone()).await? + let Some(entry) = + build_range_index(&store, ®ion, &version, file.clone()).await? else { // Defer series construction if a needed range is not ready. return Ok(None); }; - completed.push(id); + completed.push(entry.file_id); } build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) .await diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index d56e10698fd..95cf87f83d3 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -60,6 +60,8 @@ pub(crate) struct WindowSequence { /// relies on this complete-coverage contract, not on `source_file_ids`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { + /// Completed size in the catalog; zero in the footer written before completion. + pub(crate) file_size: u64, pub(crate) index_uuid: FileId, /// Inclusive bucket start. pub(crate) bucket_start: Timestamp, @@ -97,6 +99,13 @@ impl SeriesIndexEntry { } } +/// A completed per-SST range index. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct RangeIndexEntry { + pub(crate) file_id: FileId, + pub(crate) file_size: u64, +} + #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct SeriesIndexCatalog { pub(crate) indexes: Vec, @@ -104,7 +113,7 @@ pub(crate) struct SeriesIndexCatalog { #[derive(Debug, Default, Serialize, Deserialize)] pub(crate) struct RangeIndexCatalog { - pub(crate) indexes: Vec, + pub(crate) indexes: Vec, } pub(crate) fn range_catalog_path(region_id: RegionId) -> String { @@ -185,9 +194,12 @@ pub(crate) async fn load_version_control( let series = load_catalog::(store, &series_catalog_path(region_id)) .await .unwrap_or_default(); - // TODO: Handle catalog entries whose index files are missing from storage. let version = SeriesIndexVersion::new( - range.indexes.into_iter().collect(), + range + .indexes + .into_iter() + .map(|entry| (entry.file_id, entry)) + .collect(), series .indexes .into_iter() @@ -216,9 +228,9 @@ mod tests { use store_api::storage::{FileId, RegionId}; use crate::series_index::catalog::{ - RangeIndexCatalog, SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, load_catalog, - load_version_control, range_catalog_path, series_catalog_path, series_metadata, - store_catalog, + RangeIndexCatalog, RangeIndexEntry, SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, + load_catalog, load_version_control, range_catalog_path, series_catalog_path, + series_metadata, store_catalog, }; use crate::series_index::purger::series_index_channel; @@ -250,6 +262,7 @@ mod tests { fn coverage_uses_exclusive_time_end_and_inclusive_file_sequences() { let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(1), bucket_end: Timestamp::new_second(2), @@ -304,7 +317,10 @@ mod tests { .write( &range_catalog_path(region_id), serde_json::to_vec(&RangeIndexCatalog { - indexes: vec![file_id], + indexes: vec![RangeIndexEntry { + file_id, + file_size: 1, + }], }) .unwrap(), ) @@ -315,7 +331,8 @@ mod tests { .await .unwrap(); let control = load_version_control(&store, region_id, &purger).await; - assert!(control.current().range_indexes.contains(&file_id)); + assert_eq!(1, control.current().range_indexes.len()); + assert_eq!(1, control.current().range_indexes[&file_id].file_size); assert!(control.current().series_indexes.is_empty()); let layer = MockLayerBuilder::default() .reader_factory(Arc::new(|_, _, _| Box::new(FailingCatalogReader))) @@ -342,6 +359,7 @@ mod tests { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), bucket_end: Timestamp::new_second(100), diff --git a/src/mito2/src/series_index/maintenance.rs b/src/mito2/src/series_index/maintenance.rs index d782cf5a2c1..f5221bb1ebe 100644 --- a/src/mito2/src/series_index/maintenance.rs +++ b/src/mito2/src/series_index/maintenance.rs @@ -20,7 +20,7 @@ use std::time::{Duration, Instant}; use common_telemetry::{debug, info}; use object_store::ObjectStore; -use store_api::storage::RegionId; +use store_api::storage::{FileId, RegionId}; use crate::error::Result; use crate::metrics::{SERIES_INDEX_RECONCILE_ELAPSED, SERIES_INDEX_RECONCILE_TOTAL}; @@ -74,6 +74,7 @@ impl ReconcileStats { } /// Reconciles indexes for one region snapshot, persists catalogs, then atomically publishes it. +#[allow(clippy::too_many_arguments)] pub(crate) async fn reconcile_series_indexes( worker_id: u32, store: ObjectStore, @@ -82,6 +83,7 @@ pub(crate) async fn reconcile_series_indexes( now_ms: i64, purger: IndexFilePurger, enable_range_index: bool, + allow_builds: bool, ) -> Result { let total_start = Instant::now(); // Use this snapshot throughout reconciliation, even if the region version advances. @@ -104,6 +106,7 @@ pub(crate) async fn reconcile_series_indexes( &purger, &mut unpublished, enable_range_index, + allow_builds, ) .await?; SERIES_INDEX_RECONCILE_ELAPSED @@ -164,6 +167,7 @@ async fn build_index_version( purger: &IndexFilePurger, unpublished: &mut UnpublishedSeriesFiles, enable_range_index: bool, + allow_builds: bool, ) -> Result<(Option, ReconcileStats)> { let mut stats = ReconcileStats::default(); let files = version @@ -179,25 +183,23 @@ async fn build_index_version( .map(|file| file.file_id().file_id()) .collect::>(); let current = region.series_index_version(); - // The SST purger deletes companion range files after final handle release. Prune - // metadata here using the captured SST snapshot, independently of physical deletion; - // a later region-version change is picked up by the next reconciliation. - stats.removed_range = current.range_indexes.difference(&visible).count(); - let buckets = match version.compaction_time_window { - Some(window) => rounded_bucket_width(requested_bucket_width, window) + // Empty build inputs still let the planner expire established bucket coverage. + let buckets = match (allow_builds, version.compaction_time_window) { + (true, Some(window)) => rounded_bucket_width(requested_bucket_width, window) // Successful rounding guarantees the window fits in i64; subsecond windows // use the same one-second minimum as rounded_bucket_width. .map(|width| { group_files_into_series_buckets(&files, width, (window.as_secs() as i64).max(1)) }) .unwrap_or_default(), - None => { + (true, None) => { debug!( "Deferring series indexes without compaction window, worker: {worker_id}, region: {}", region.region_id ); Vec::new() } + (false, _) => Vec::new(), }; let plan = plan_series_indexes( buckets, @@ -207,33 +209,41 @@ async fn build_index_version( ); stats.computed_buckets = plan.computed_buckets; stats.skipped_buckets = plan.skipped_buckets; + let mut next = prune_index_version(¤t, &visible, &plan.expired_index_ids, &mut stats); + if !allow_builds { + if let Some(next) = &mut next { + next.index_buckets = plan.index_buckets; + } + return Ok((next, stats)); + } if plan.builds.is_empty() - && plan.expired_index_ids.is_empty() - && stats.removed_range == 0 + && next.is_none() && (!enable_range_index || current.range_indexes.len() == visible.len()) { return Ok((None, stats)); } - let mut range_indexes = current.range_indexes.clone(); - range_indexes.retain(|file_id| visible.contains(file_id)); - let mut series_indexes = current.series_indexes.clone(); - for id in plan - .expired_index_ids - .iter() - .chain(&plan.superseded_index_ids) - { + let SeriesIndexVersion { + mut range_indexes, + mut series_indexes, + .. + } = next.unwrap_or_else(|| SeriesIndexVersion { + range_indexes: current.range_indexes.clone(), + series_indexes: current.series_indexes.clone(), + index_buckets: Default::default(), + }); + let index_buckets = plan.index_buckets; + for id in &plan.superseded_index_ids { series_indexes.remove(id); } for (bucket, expected) in plan.builds { // Complete companion indexes independently so a failed series build preserves them. for file in &bucket.files { if enable_range_index - && !range_indexes.contains(&file.file_id().file_id()) - && let Some(file_id) = - build_range_index(store, region, version, file.clone()).await? + && !range_indexes.contains_key(&file.file_id().file_id()) + && let Some(entry) = build_range_index(store, region, version, file.clone()).await? { stats.built_range += 1; - range_indexes.insert(file_id); + range_indexes.insert(entry.file_id, entry); } } let series_handle = @@ -246,12 +256,12 @@ async fn build_index_version( if enable_range_index { for file in files { let file_id = file.file_id().file_id(); - if range_indexes.contains(&file_id) { + if range_indexes.contains_key(&file_id) { continue; } - if let Some(file_id) = build_range_index(store, region, version, file).await? { + if let Some(entry) = build_range_index(store, region, version, file).await? { stats.built_range += 1; - range_indexes.insert(file_id); + range_indexes.insert(entry.file_id, entry); } } } @@ -268,11 +278,43 @@ async fn build_index_version( let next = SeriesIndexVersion { range_indexes, series_indexes, - index_buckets: plan.index_buckets, + index_buckets, }; Ok((Some(next), stats)) } +/// Prunes obsolete metadata without building indexes or retiring published handles. +/// Returns `None` without cloning the current version when no entries need removal. +/// Physical range deletion belongs to the SST purger; series handles are retired only +/// after the cleaned catalogs and snapshot have been published successfully. +fn prune_index_version( + current: &SeriesIndexVersion, + visible: &HashSet, + expired_index_ids: &[FileId], + stats: &mut ReconcileStats, +) -> Option { + if current.range_indexes.keys().all(|id| visible.contains(id)) + && expired_index_ids + .iter() + .all(|id| !current.series_indexes.contains_key(id)) + { + return None; + } + let mut next = SeriesIndexVersion { + range_indexes: current.range_indexes.clone(), + series_indexes: current.series_indexes.clone(), + index_buckets: current.index_buckets.clone(), + }; + next.range_indexes + .retain(|file_id, _| visible.contains(file_id)); + for id in expired_index_ids { + next.series_indexes.remove(id); + } + stats.removed_range = current.range_indexes.len() - next.range_indexes.len(); + stats.removed_series = current.series_indexes.len() - next.series_indexes.len(); + Some(next) +} + /// Writes changed catalogs in a stable order; the two writes are not atomic together. async fn persist_index_catalogs( store: &ObjectStore, @@ -281,8 +323,9 @@ async fn persist_index_catalogs( stats: &ReconcileStats, ) -> Result<()> { if stats.built_range + stats.removed_range > 0 { - let mut range_entries = next.range_indexes.iter().copied().collect::>(); - range_entries.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); + let mut range_entries = next.range_indexes.values().copied().collect::>(); + range_entries + .sort_unstable_by(|left, right| left.file_id.as_bytes().cmp(right.file_id.as_bytes())); store_catalog( store, &range_catalog_path(region_id), @@ -328,3 +371,77 @@ fn publish_index_version(region: &MitoRegionRef, next: Arc) } } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use common_time::Timestamp; + use object_store::services::Memory; + + use super::*; + use crate::series_index::catalog::{RangeIndexEntry, SeriesIndexEntry}; + use crate::series_index::purger::series_index_channel; + + #[rstest::rstest] + fn test_prune_index_version( + #[values(false, true)] remove_range: bool, + #[values(false, true)] remove_series: bool, + ) { + let range_id = FileId::random(); + let series_id = FileId::random(); + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, mut receiver) = series_index_channel(store); + let current = SeriesIndexVersion::new( + HashMap::from([( + range_id, + RangeIndexEntry { + file_id: range_id, + file_size: 10, + }, + )]), + HashMap::from([( + series_id, + SeriesIndexFileHandle::new( + RegionId::new(1, 1), + SeriesIndexEntry { + index_uuid: series_id, + file_size: 20, + bucket_start: Timestamp::new_second(0), + bucket_end: Timestamp::new_second(100), + source_file_ids: vec![range_id], + min_file_sequence: 1, + max_file_sequence: 1, + compaction_window_secs: 100, + window_sequences: Default::default(), + }, + purger, + ), + )]), + ); + let visible = if remove_range { + HashSet::new() + } else { + HashSet::from([range_id]) + }; + // Unknown and repeated IDs must not inflate removal counts. + let mut expired = vec![FileId::random()]; + if remove_series { + expired.extend([series_id, series_id]); + } + let mut stats = ReconcileStats::default(); + let next = prune_index_version(¤t, &visible, &expired, &mut stats); + assert_eq!(remove_range || remove_series, next.is_some()); + assert_eq!(usize::from(remove_range), stats.removed_range); + assert_eq!(usize::from(remove_series), stats.removed_series); + if let Some(next) = next { + assert_eq!(!remove_range, next.range_indexes.contains_key(&range_id)); + assert_eq!(!remove_series, next.series_indexes.contains_key(&series_id)); + } + assert_eq!(1, current.range_indexes.len()); + assert_eq!(1, current.series_indexes.len()); + drop(current); + // Pruning alone must never retire published files. + assert!(receiver.try_recv().is_err()); + } +} diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs index 4613edae913..0b9befe4b8a 100644 --- a/src/mito2/src/series_index/searcher.rs +++ b/src/mito2/src/series_index/searcher.rs @@ -312,6 +312,7 @@ mod tests { ) -> (SeriesIndexFileHandle, UnboundedReceiver) { let (purger, receiver) = series_index_channel(store.clone()); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: common_time::Timestamp::new_second(0), bucket_end: common_time::Timestamp::new_second(60), diff --git a/src/mito2/src/series_index/task.rs b/src/mito2/src/series_index/task.rs index 2afbb245014..5579cf87ce4 100644 --- a/src/mito2/src/series_index/task.rs +++ b/src/mito2/src/series_index/task.rs @@ -15,7 +15,7 @@ //! Worker-owned background maintenance for series indexes. use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use common_telemetry::{info, warn}; @@ -25,7 +25,9 @@ use tokio::sync::mpsc::UnboundedReceiver; use tokio::task::JoinHandle; use tokio::time::{Instant, MissedTickBehavior}; -use crate::metrics::SERIES_INDEX_RECONCILE_TOTAL; +use crate::metrics::{ + SERIES_INDEX_CAPACITY_DEFERRED, SERIES_INDEX_DISK_BYTES, SERIES_INDEX_RECONCILE_TOTAL, +}; use crate::region::{RegionLeaderState, RegionMapRef, RegionRoleState}; use crate::series_index::maintenance::reconcile_series_indexes; use crate::series_index::purger::{IndexFilePurger, PurgeRequest, run_index_purge_task}; @@ -75,6 +77,8 @@ pub(crate) fn spawn_series_index_tasks( bucket_width: Duration, purger: IndexFilePurger, purge_receiver: UnboundedReceiver, + disk_usage: Arc, + max_size: u64, interval: Duration, time_provider: TimeProviderRef, enable_range_index: bool, @@ -92,6 +96,9 @@ pub(crate) fn spawn_series_index_tasks( regions, bucket_width, purger, + disk_usage, + max_size, + reported_usage: 0, state, interval, time_provider, @@ -108,6 +115,9 @@ struct SeriesIndexTask { regions: RegionMapRef, bucket_width: Duration, purger: IndexFilePurger, + disk_usage: Arc, + max_size: u64, + reported_usage: u64, worker_id: u32, state: Arc, interval: Duration, @@ -136,8 +146,33 @@ impl SeriesIndexTask { info!("Stop series-index background task, worker: {worker_id}"); } + /// Each worker contributes only its open regions, refreshed at maintenance boundaries. + fn refresh_usage(&mut self) { + let usage = self + .regions + .list_regions() + .iter() + .map(|region| region.series_index_version().disk_usage()) + .sum(); + self.report_usage(usage); + } + + fn report_usage(&mut self, usage: u64) { + if usage >= self.reported_usage { + let delta = usage - self.reported_usage; + self.disk_usage.fetch_add(delta, Ordering::Relaxed); + SERIES_INDEX_DISK_BYTES.add(delta as i64); + } else { + let delta = self.reported_usage - usage; + self.disk_usage.fetch_sub(delta, Ordering::Relaxed); + SERIES_INDEX_DISK_BYTES.sub(delta as i64); + } + self.reported_usage = usage; + } + /// Runs periodic maintenance independently of incoming deletion requests. async fn maintain(&mut self) { + self.refresh_usage(); for region in self.regions.list_regions() { if !self.state.is_running() { break; @@ -150,6 +185,11 @@ impl SeriesIndexTask { ) { continue; } + // Full capacity defers builds, but cleanup must still reclaim published usage. + let allow_builds = self.disk_usage.load(Ordering::Relaxed) < self.max_size; + if !allow_builds { + SERIES_INDEX_CAPACITY_DEFERRED.inc(); + } if let Err(error) = reconcile_series_indexes( self.worker_id, self.store.clone(), @@ -158,6 +198,7 @@ impl SeriesIndexTask { self.time_provider.current_time_millis(), self.purger.clone(), self.enable_range_index, + allow_builds, ) .await { @@ -166,10 +207,17 @@ impl SeriesIndexTask { .inc(); warn!(error; "Failed to reconcile series indexes, worker: {}, region: {}", self.worker_id, region.region_id); } + self.refresh_usage(); } } } +impl Drop for SeriesIndexTask { + fn drop(&mut self) { + self.report_usage(0); + } +} + #[cfg(test)] mod tests { use object_store::services::Memory; @@ -211,6 +259,9 @@ mod tests { regions, bucket_width: Duration::from_secs(100), purger, + disk_usage: Arc::default(), + max_size: u64::MAX, + reported_usage: 0, worker_id: 0, state: Arc::new(SeriesIndexTaskState::new()), interval: Duration::from_secs(3600), @@ -230,4 +281,236 @@ mod tests { } engine.stop().await.unwrap(); } + + #[rstest::rstest] + #[case::at_capacity(0, true)] + #[case::above_capacity(1, true)] + #[case::at_capacity_without_range(0, false)] + #[case::above_capacity_without_range(1, false)] + #[tokio::test] + async fn test_full_capacity_cleanup_and_recovery( + #[case] excess: u64, + #[case] enable_range_index: bool, + ) { + use std::sync::Mutex; + + use object_store::layers::mock::MockLayerBuilder; + + use crate::series_index::catalog::load_version_control; + use crate::series_index::tests::prepare_region_with_timestamps; + use crate::time_provider::mock::MockTimeProvider; + + let mut env = TestEnv::with_prefix("series-capacity-cleanup").await; + let (engine, region) = + prepare_region_with_timestamps(&mut env, &[1000, 2000, 3000, 4000, 5000]).await; + assert_eq!( + 5, + region + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .count() + ); + let mut options = region.version().options.clone(); + options.ttl = Some(common_time::TimeToLive::Duration(Duration::from_secs(100))); + region.version_control.alter_options(options); + let writes = Arc::new(Mutex::new(Vec::new())); + let captured = writes.clone(); + let store = ObjectStore::new(Memory::default()).unwrap().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + captured.lock().unwrap().push(path.to_string()); + inner + })) + .build() + .unwrap(), + ); + let (purger, mut receiver) = series_index_channel(store.clone()); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region.clone()); + let clock = Arc::new(MockTimeProvider::new(0)); + let usage = Arc::new(AtomicU64::new(0)); + let mut task = SeriesIndexTask { + store: store.clone(), + regions, + bucket_width: Duration::from_secs(100), + purger, + disk_usage: usage.clone(), + max_size: u64::MAX, + reported_usage: 0, + worker_id: 0, + state: Arc::new(SeriesIndexTaskState::new()), + interval: Duration::from_secs(3600), + time_provider: clock.clone(), + enable_range_index, + }; + task.maintain().await; + let previous = region.series_index_version(); + assert_eq!(1, previous.series_indexes.len()); + let old_id = *previous.series_indexes.keys().next().unwrap(); + task.max_size = previous.disk_usage() - excess; + writes.lock().unwrap().clear(); + task.maintain().await; + assert!(Arc::ptr_eq(&previous, ®ion.series_index_version())); + assert!(writes.lock().unwrap().is_empty()); + + // Removing the newest SST makes range coverage obsolete and would normally + // trigger a series replacement over the four remaining SSTs. + let sources = region.version(); + let newest = sources + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .max_by_key(|file| file.meta_ref().sequence) + .unwrap() + .meta_ref() + .clone(); + region.version_control.apply_edit( + Some(crate::manifest::action::RegionEdit { + files_to_remove: vec![newest.clone()], + files_to_add: Vec::new(), + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }), + &[], + crate::test_util::new_noop_file_purger(), + ); + task.maintain().await; + let cleaned = region.series_index_version(); + assert!(!cleaned.range_indexes.contains_key(&newest.file_id)); + assert_eq!( + usize::from(enable_range_index) * 4, + cleaned.range_indexes.len() + ); + assert_eq!(previous.index_buckets, cleaned.index_buckets); + assert_eq!(1, cleaned.series_indexes.len()); + assert!(cleaned.series_indexes.contains_key(&old_id)); + assert_eq!(cleaned.disk_usage(), usage.load(Ordering::Relaxed)); + assert_eq!( + previous.disk_usage() + - previous + .range_indexes + .get(&newest.file_id) + .map_or(0, |e| e.file_size), + cleaned.disk_usage() + ); + assert!( + writes + .lock() + .unwrap() + .iter() + .all(|path| path == &range_catalog_path(region.region_id)) + ); + let restored = load_version_control(&store, region.region_id, &task.purger).await; + assert_eq!(cleaned.range_indexes, restored.current().range_indexes); + assert_eq!(cleaned.index_buckets, restored.current().index_buckets); + drop(restored); + + // Expiration must work even when the shared estimate remains full. + task.max_size = cleaned.disk_usage() - excess; + clock.set_now(201_000); + writes.lock().unwrap().clear(); + task.maintain().await; + let expired = region.series_index_version(); + assert!(expired.series_indexes.is_empty()); + assert!(expired.index_buckets.is_empty()); + assert_eq!(cleaned.range_indexes, expired.range_indexes); + assert_eq!(expired.disk_usage(), usage.load(Ordering::Relaxed)); + assert!(usage.load(Ordering::Relaxed) < task.max_size); + assert_eq!( + *writes.lock().unwrap(), + vec![series_catalog_path(region.region_id)] + ); + let restored = load_version_control(&store, region.region_id, &task.purger).await; + assert!(restored.current().series_indexes.is_empty()); + assert_eq!(expired.range_indexes, restored.current().range_indexes); + assert!(receiver.try_recv().is_err()); + drop(previous); + assert!(receiver.try_recv().is_err()); + drop(cleaned); + assert_eq!(old_id, receiver.try_recv().unwrap().file_id.file_id()); + assert!(receiver.try_recv().is_err()); + + // Make the remaining SSTs eligible again; reclaimed capacity admits a build + // without closing the region or increasing the configured limit. + clock.set_now(0); + task.maintain().await; + let rebuilt = region.series_index_version(); + assert_eq!(1, rebuilt.series_indexes.len()); + assert!(!rebuilt.series_indexes.contains_key(&old_id)); + assert_eq!( + 4, + rebuilt + .series_indexes + .values() + .next() + .unwrap() + .entry() + .source_file_ids + .len() + ); + assert_eq!(rebuilt.disk_usage(), usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_shared_usage_defers_builds_and_allows_overshoot() { + let mut env = TestEnv::with_prefix("series-approximate-usage").await; + let (engine, region) = prepare_region(&mut env).await; + let usage = Arc::new(AtomicU64::new(0)); + let store = ObjectStore::new(Memory::default()).unwrap(); + let make_task = || SeriesIndexTask { + store: store.clone(), + regions: Arc::new(RegionMap::default()), + bucket_width: Duration::from_secs(100), + purger: series_index_channel(store.clone()).0, + disk_usage: usage.clone(), + max_size: 1024, + reported_usage: 0, + worker_id: 0, + state: Arc::new(SeriesIndexTaskState::new()), + interval: Duration::from_secs(3600), + time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, + }; + let mut task = make_task(); + task.regions.insert_region(region.clone()); + task.maintain().await; + let bytes = region.series_index_version().disk_usage(); + assert!( + bytes > task.max_size, + "a started reconciliation may exceed the limit" + ); + assert_eq!(bytes, usage.load(Ordering::Relaxed)); + // Another worker sees the same estimate and defers builds without creating coverage. + let mut other_env = TestEnv::with_prefix("series-approximate-other").await; + let (other_engine, other_region) = prepare_region(&mut other_env).await; + let mut other = make_task(); + other.regions.insert_region(other_region.clone()); + other.max_size = bytes; // Exact equality also skips. + other.maintain().await; + assert_eq!(0, other_region.series_index_version().disk_usage()); + assert_eq!(bytes, usage.load(Ordering::Relaxed)); + + // Closing a region reduces the estimate on the next pass, even while full. + task.regions.remove_region(region.region_id); + task.maintain().await; + assert_eq!(0, usage.load(Ordering::Relaxed)); + other.maintain().await; + assert!(other_region.series_index_version().disk_usage() > 0); + assert_eq!( + other_region.series_index_version().disk_usage(), + usage.load(Ordering::Relaxed) + ); + drop(other); + assert_eq!(0, usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + other_engine.stop().await.unwrap(); + } } diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index c1efe3afdcf..8ee898277bc 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -48,11 +48,16 @@ pub(super) async fn prepare_region(env: &mut TestEnv) -> (MitoEngine, MitoRegion prepare_region_with_timestamps(env, &[1000, 2000, 3000, 4000]).await } -async fn prepare_region_with_timestamps( +pub(crate) async fn prepare_region_with_timestamps( env: &mut TestEnv, timestamps: &[i64], ) -> (MitoEngine, MitoRegionRef) { - let engine = env.create_engine(MitoConfig::default()).await; + let engine = env + .create_engine(MitoConfig { + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }) + .await; let metadata = Arc::new(sst_region_metadata_with_encoding( PrimaryKeyEncoding::Sparse, )); @@ -93,6 +98,10 @@ async fn prepare_region_with_timestamps( .handle_request(region_id, RegionRequest::Create(request)) .await .unwrap(); + // Flush completion can precede compaction scheduling. Keep background TTL + // cleanup from racing tests that alter options after preparing the SSTs. + let region = engine.get_region(region_id).unwrap(); + region.update_schedule_compaction_millis(); for &ts in timestamps { engine .handle_request( @@ -122,7 +131,6 @@ async fn prepare_region_with_timestamps( .unwrap(); flush_region(&engine, region_id, None).await; } - let region = engine.get_region(region_id).unwrap(); (engine, region) } @@ -171,6 +179,7 @@ impl IndexTest { 0, self.purger.clone(), enable_range_index, + true, ) .await } @@ -311,7 +320,7 @@ async fn test_reconcile_restores_and_reuses_indexes() { let missing_paths = [ range_index_path( region.region_id, - *first.range_indexes.iter().next().unwrap(), + *first.range_indexes.keys().next().unwrap(), ), series_index_path(region.region_id, first_id), ]; @@ -325,6 +334,7 @@ async fn test_reconcile_restores_and_reuses_indexes() { .current(); assert_eq!(first.range_indexes, restored.range_indexes); assert_eq!(first.index_buckets, restored.index_buckets); + assert_eq!(first.disk_usage(), restored.disk_usage()); assert_eq!( first.series_indexes[&first_id].entry(), restored.series_indexes[&first_id].entry() @@ -697,13 +707,13 @@ async fn test_failed_reconcile_retires_only_unpublished_series() { ); assert_eq!(if failure == "range" { 2 } else { 3 }, plan.builds.len()); let (bucket, entry) = &plan.builds[0]; - let mut ranges = HashSet::new(); + let mut ranges = HashMap::new(); for file in &bucket.files { - let id = build_range_index(store, ®ion, &version, file.clone()) + let entry = build_range_index(store, ®ion, &version, file.clone()) .await .unwrap() .unwrap(); - ranges.insert(id); + ranges.insert(entry.file_id, entry); } let handle = build_series_index(store, ®ion, &version, bucket, entry, purger) .await @@ -838,6 +848,8 @@ async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { Duration::from_secs(100), purger, receiver, + Arc::default(), + u64::MAX, Duration::from_secs(3600), clock.clone(), enable_range_index, @@ -866,6 +878,7 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { let store = ObjectStore::new(Memory::default()).unwrap(); let region_id = RegionId::new(1, 1); let entry = SeriesIndexEntry { + file_size: 0, index_uuid: FileId::random(), bucket_start: Timestamp::new_second(0), bucket_end: Timestamp::new_second(100), @@ -914,6 +927,8 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { Duration::from_secs(100), purger, receiver, + Arc::default(), + u64::MAX, Duration::from_secs(3600), Arc::new(StdTimeProvider), true, diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index c31dfde906b..fa2115fc1f2 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -14,7 +14,7 @@ //! Immutable index snapshots and aggregate series-file handles. -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap}; use std::fmt::{self, Debug, Formatter}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; @@ -23,7 +23,7 @@ use common_time::Timestamp; use store_api::storage::{FileId, RegionId}; use crate::series_index::bucket::IndexBucket; -use crate::series_index::catalog::SeriesIndexEntry; +use crate::series_index::catalog::{RangeIndexEntry, SeriesIndexEntry}; use crate::series_index::purger::{IndexFilePurger, PurgeRequest}; use crate::sst::file::RegionFileId; @@ -94,7 +94,7 @@ impl Drop for SeriesIndexFileHandleInner { pub(crate) struct SeriesIndexVersion { /// Range indexes for visible SSTs; reconciliation removes IDs absent from its SST snapshot. /// Physical deletion is independently handled by the SST file purger. - pub(crate) range_indexes: HashSet, + pub(crate) range_indexes: HashMap, pub(crate) series_indexes: HashMap, pub(crate) index_buckets: BTreeMap, } @@ -102,7 +102,7 @@ pub(crate) struct SeriesIndexVersion { impl SeriesIndexVersion { /// Restores bucket lookup from immutable index coverage stored in the catalog. pub(crate) fn new( - range_indexes: HashSet, + range_indexes: HashMap, series_indexes: HashMap, ) -> Self { let mut index_buckets = BTreeMap::new(); @@ -116,6 +116,19 @@ impl SeriesIndexVersion { } } + /// Approximate installed usage; old snapshots and unpublished outputs are excluded. + pub(crate) fn disk_usage(&self) -> u64 { + self.range_indexes + .values() + .map(|entry| entry.file_size) + .sum::() + + self + .series_indexes + .values() + .map(|handle| handle.entry().file_size) + .sum::() + } + fn mark_all_deleted(&self) { self.series_indexes .values() diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 9584ff0a6af..714e62ae4af 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -713,7 +713,7 @@ impl ParquetReaderBuilder { && context .version .range_indexes - .contains(&self.file_handle.file_id().file_id())) + .contains_key(&self.file_handle.file_id().file_id())) .then(|| context.store.clone()) }); let context = FileRangeContext::new( diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 24eec5dc2b3..0f67d8d2a1b 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -35,7 +35,7 @@ mod handle_write; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use common_base::Plugins; @@ -198,6 +198,7 @@ impl WorkerGroup { let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); let series_index_store = series_index_store_from_config(&config, data_home).await?; + let series_index_disk_usage = Arc::new(AtomicU64::new(0)); let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes)); let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions)); let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes)); @@ -251,6 +252,7 @@ impl WorkerGroup { write_buffer_manager: write_buffer_manager.clone(), index_build_job_pool: index_build_job_pool.clone(), series_index_store: series_index_store.clone(), + series_index_disk_usage: series_index_disk_usage.clone(), flush_job_pool: flush_job_pool.clone(), compact_job_pool: compact_job_pool.clone(), purge_scheduler: purge_scheduler.clone(), @@ -410,6 +412,7 @@ impl WorkerGroup { let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); let series_index_store = series_index_store_from_config(&config, data_home).await?; + let series_index_disk_usage = Arc::new(AtomicU64::new(0)); let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes)); let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions)); let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes)); @@ -464,6 +467,7 @@ impl WorkerGroup { write_buffer_manager: write_buffer_manager.clone(), index_build_job_pool: index_build_job_pool.clone(), series_index_store: series_index_store.clone(), + series_index_disk_usage: series_index_disk_usage.clone(), flush_job_pool: flush_job_pool.clone(), compact_job_pool: compact_job_pool.clone(), purge_scheduler: purge_scheduler.clone(), @@ -572,6 +576,7 @@ struct WorkerStarter { compact_job_pool: SchedulerRef, index_build_job_pool: SchedulerRef, series_index_store: Option, + series_index_disk_usage: Arc, flush_job_pool: SchedulerRef, purge_scheduler: SchedulerRef, listener: WorkerListener, @@ -620,6 +625,8 @@ impl WorkerStarter { self.config.experimental_series_index_bucket_width, purger, purge_receiver, + self.series_index_disk_usage.clone(), + self.config.experimental_series_index_max_size.as_bytes(), self.config.experimental_series_index_maintenance_interval, self.time_provider.clone(), self.config.experimental_enable_range_index, @@ -951,7 +958,7 @@ struct RegionWorkerLoop { index_build_scheduler: IndexBuildScheduler, /// Controls the worker-owned series-index task. series_index_task_state: Option>, - /// Store for companion range indexes deleted by the region SST purger. + /// Local store for series and range indexes managed by reconciliation. series_index_store: Option, series_index_purger: Option, /// Schedules background flush requests. diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 3e64a6d74e3..a4841afad93 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2561,6 +2561,7 @@ experimental_manifest_keep_removed_file_count = 256 experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false experimental_enable_series_index = false +experimental_series_index_max_size = "5GiB" experimental_enable_range_index = false experimental_series_index_maintenance_interval = "5m" experimental_series_index_bucket_width = "5days" From cc1be828588f1754ebcb2d27224191d37e2d0ed6 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 02:12:11 +0000 Subject: [PATCH 028/133] fix: merge every state row in geo_path and json_encode_path (#9339) Signed-off-by: Dennis Zhuang --- src/common/function/src/aggrs/geo/encoding.rs | 34 ++++++++------ src/common/function/src/aggrs/geo/geo_path.rs | 34 ++++++++------ .../standalone/common/function/geo.result | 44 +++++++++++++++++++ .../cases/standalone/common/function/geo.sql | 24 ++++++++++ 4 files changed, 108 insertions(+), 28 deletions(-) diff --git a/src/common/function/src/aggrs/geo/encoding.rs b/src/common/function/src/aggrs/geo/encoding.rs index 35a59d1ecfb..67dfd493a77 100644 --- a/src/common/function/src/aggrs/geo/encoding.rs +++ b/src/common/function/src/aggrs/geo/encoding.rs @@ -221,18 +221,24 @@ impl DfAccumulator for JsonEncodePathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -331,9 +337,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/src/common/function/src/aggrs/geo/geo_path.rs b/src/common/function/src/aggrs/geo/geo_path.rs index b09745eaeb7..c4e366c83b0 100644 --- a/src/common/function/src/aggrs/geo/geo_path.rs +++ b/src/common/function/src/aggrs/geo/geo_path.rs @@ -250,18 +250,24 @@ impl DfAccumulator for GeoPathAccumulator { ))); } - for state in states { - let state = as_struct_array(state)?; - let lat_list = as_list_array(state.column(0))?.value(0); - let lat_array = as_primitive_array::(&lat_list)?; - let lng_list = as_list_array(state.column(1))?.value(0); - let lng_array = as_primitive_array::(&lng_list)?; - let ts_list = as_list_array(state.column(2))?.value(0); - let ts_array = as_primitive_array::(&ts_list)?; + let state = as_struct_array(&states[0])?; + let lat_lists = as_list_array(state.column(0))?; + let lng_lists = as_list_array(state.column(1))?; + let ts_lists = as_list_array(state.column(2))?; + for row in 0..state.len() { + if state.is_null(row) { + continue; + } + let lat_list = lat_lists.value(row); + let lng_list = lng_lists.value(row); + let ts_list = ts_lists.value(row); - self.lat.extend(lat_array); - self.lng.extend(lng_array); - self.timestamp.extend(ts_array); + self.lat + .extend(as_primitive_array::(&lat_list)?); + self.lng + .extend(as_primitive_array::(&lng_list)?); + self.timestamp + .extend(as_primitive_array::(&ts_list)?); } Ok(()) @@ -403,9 +409,9 @@ mod tests { _ => panic!("Expected Struct scalar value"), }; - // Merge state arrays - merged.merge_batch(&[state_array1]).unwrap(); - merged.merge_batch(&[state_array2]).unwrap(); + // States from different partitions arrive as rows of one batch + let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap(); + merged.merge_batch(&[states]).unwrap(); // Evaluate merged result let result = merged.evaluate().unwrap(); diff --git a/tests/cases/standalone/common/function/geo.result b/tests/cases/standalone/common/function/geo.result index 31efbf401da..3fc73a3b98b 100644 --- a/tests/cases/standalone/common/function/geo.result +++ b/tests/cases/standalone/common/function/geo.result @@ -437,3 +437,47 @@ FROM | true | false | true | false | false | true | +--------------------------+--------------------------+------------------------+------------------------+----------------------------------+----------------------------------+ +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +Affected Rows: 7 + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------+ +| g | geo_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+-------------------------------------------------------------------------------------+ +| false | {lat: [1.0, 2.0, 3.0], lng: [11.0, 12.0, 13.0]} | +| true | {lat: [4.0, 5.0, 6.0, 7.0], lng: [14.0, 15.0, 16.0, 17.0]} | ++-------+-------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + ++-------+---------------------------------------------------------------------------------------------+ +| g | json_encode_path(geo_path_partitioned.lat,geo_path_partitioned.lon,geo_path_partitioned.ts) | ++-------+---------------------------------------------------------------------------------------------+ +| false | [[11.0,1.0],[12.0,2.0],[13.0,3.0]] | +| true | [[14.0,4.0],[15.0,5.0],[16.0,6.0],[17.0,7.0]] | ++-------+---------------------------------------------------------------------------------------------+ + +DROP TABLE geo_path_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/geo.sql b/tests/cases/standalone/common/function/geo.sql index 43025cc196b..1de7d7c82aa 100644 --- a/tests/cases/standalone/common/function/geo.sql +++ b/tests/cases/standalone/common/function/geo.sql @@ -180,3 +180,27 @@ FROM 'POLYGON ((-121.491698 38.653343, -121.582353 38.556757, -121.469721 38.449287, -121.315883 38.541721, -121.491698 38.653343))' AS polygon2, 'POLYGON ((-122.089628 37.450332, -122.20535 37.378342, -122.093062 37.36088, -122.044301 37.372886, -122.089628 37.450332))' AS polygon3, ); + +CREATE TABLE geo_path_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + lon DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO geo_path_partitioned VALUES + (1000, 1, 1, 11), + (2000, 11, 2, 12), + (3000, 21, 3, 13), + (4000, 2, 4, 14), + (5000, 12, 5, 15), + (6000, 22, 6, 16), + (7000, 3, 7, 17); + +SELECT lat > 3 AS g, geo_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +SELECT lat > 3 AS g, json_encode_path(lat, lon, ts) FROM geo_path_partitioned GROUP BY g ORDER BY g; + +DROP TABLE geo_path_partitioned; From b9f991502c98b28c98525d9a699e9acad5e42b58 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 05:38:54 +0000 Subject: [PATCH 029/133] refactor: remove experimental vector index (#9345) Signed-off-by: Dennis Zhuang --- .github/workflows/integration.yml | 2 +- .github/workflows/nightly-ci.yml | 2 +- .github/workflows/rust.yml | 4 +- Cargo.lock | 104 +- .../dashboards/metrics/cluster/dashboard.json | 6 +- .../dashboards/metrics/cluster/dashboard.md | 2 +- .../dashboards/metrics/cluster/dashboard.yaml | 6 +- .../metrics/standalone/dashboard.json | 6 +- .../metrics/standalone/dashboard.md | 2 +- .../metrics/standalone/dashboard.yaml | 6 +- src/api/src/v1/column_def.rs | 49 +- src/cmd/Cargo.toml | 1 - src/cmd/src/datanode/objbench.rs | 2 - src/datatypes/src/schema.rs | 10 +- src/datatypes/src/schema/column_schema.rs | 231 ---- src/index/Cargo.toml | 6 - src/index/src/lib.rs | 2 - src/index/src/vector.rs | 179 --- src/index/src/vector/apply.rs | 428 ------- src/index/src/vector/create.rs | 345 ------ src/index/src/vector/engine.rs | 63 - src/index/src/vector/engine/usearch.rs | 213 ---- src/index/src/vector/error.rs | 152 --- src/index/src/vector/format.rs | 200 ---- src/index/src/vector/util.rs | 108 -- src/mito2/Cargo.toml | 2 - src/mito2/src/access_layer.rs | 4 - src/mito2/src/cache.rs | 30 - src/mito2/src/cache/index.rs | 2 - src/mito2/src/cache/index/vector_index.rs | 134 --- src/mito2/src/cache/write_cache.rs | 8 - src/mito2/src/compaction/compactor.rs | 4 - src/mito2/src/config.rs | 50 - src/mito2/src/error.rs | 29 - src/mito2/src/flush.rs | 2 - src/mito2/src/read/scan_region.rs | 85 -- src/mito2/src/read/scan_util.rs | 32 - src/mito2/src/sst/file.rs | 34 +- src/mito2/src/sst/index.rs | 246 ---- src/mito2/src/sst/index/column_test.rs | 2 - src/mito2/src/sst/index/indexer/abort.rs | 24 - src/mito2/src/sst/index/indexer/finish.rs | 68 -- src/mito2/src/sst/index/indexer/update.rs | 60 - src/mito2/src/sst/index/sparse_test.rs | 2 - src/mito2/src/sst/index/vector_index.rs | 24 - .../src/sst/index/vector_index/applier.rs | 503 -------- .../src/sst/index/vector_index/creator.rs | 536 --------- src/mito2/src/sst/parquet.rs | 6 - src/mito2/src/sst/parquet/reader.rs | 180 --- src/mito2/src/worker/handle_rebuild_index.rs | 2 - src/query/Cargo.toml | 1 - src/query/src/dist_plan/analyzer/test.rs | 182 --- src/query/src/dist_plan/commutativity.rs | 45 - src/query/src/dummy_catalog.rs | 12 +- src/query/src/error.rs | 9 - src/query/src/optimizer/scan_hint.rs | 115 +- .../src/optimizer/scan_hint/vector_search.rs | 1066 ----------------- src/query/src/sql/show_create_table.rs | 48 +- src/sql/src/error.rs | 11 +- src/sql/src/parsers/create_parser.rs | 245 +--- src/sql/src/parsers/utils.rs | 14 - src/sql/src/statements.rs | 78 +- src/sql/src/statements/create.rs | 170 +-- src/store-api/src/metadata.rs | 18 +- src/store-api/src/storage.rs | 5 +- src/store-api/src/storage/requests.rs | 70 +- src/table/src/table/adapter.rs | 10 +- tests-integration/Cargo.toml | 1 - tests-integration/tests/http.rs | 16 +- .../function/vector/vector_index.result | 602 ---------- .../common/function/vector/vector_index.sql | 341 ------ .../vector/vector_index_explain.result | 272 ----- .../function/vector/vector_index_explain.sql | 163 --- .../function/vector/vector_index_join.result | 240 ---- .../function/vector/vector_index_join.sql | 159 --- .../common/types/vector/vector.result | 28 + .../standalone/common/types/vector/vector.sql | 6 + tests/runner/src/env/bare.rs | 2 +- 78 files changed, 125 insertions(+), 8002 deletions(-) delete mode 100644 src/index/src/vector.rs delete mode 100644 src/index/src/vector/apply.rs delete mode 100644 src/index/src/vector/create.rs delete mode 100644 src/index/src/vector/engine.rs delete mode 100644 src/index/src/vector/engine/usearch.rs delete mode 100644 src/index/src/vector/error.rs delete mode 100644 src/index/src/vector/format.rs delete mode 100644 src/index/src/vector/util.rs delete mode 100644 src/mito2/src/cache/index/vector_index.rs delete mode 100644 src/mito2/src/sst/index/vector_index.rs delete mode 100644 src/mito2/src/sst/index/vector_index/applier.rs delete mode 100644 src/mito2/src/sst/index/vector_index/creator.rs delete mode 100644 src/query/src/optimizer/scan_hint/vector_search.rs delete mode 100644 tests/cases/standalone/common/function/vector/vector_index.result delete mode 100644 tests/cases/standalone/common/function/vector/vector_index.sql delete mode 100644 tests/cases/standalone/common/function/vector/vector_index_explain.result delete mode 100644 tests/cases/standalone/common/function/vector/vector_index_explain.sql delete mode 100644 tests/cases/standalone/common/function/vector/vector_index_join.result delete mode 100644 tests/cases/standalone/common/function/vector/vector_index_join.sql diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index e33abac3330..db4f4bef65b 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -86,7 +86,7 @@ jobs: # `cargo gc` will invoke `cargo build` with specified args run: | started_epoch="$(date +%s)" - cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions" + cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,ai_functions" elapsed_secs=$(( $(date +%s) - started_epoch )) printf -- '- Binary build: **%ss**\n' "${elapsed_secs}" >>"${GITHUB_STEP_SUMMARY}" - name: Pack greptime binaries diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index 8441ed02312..a948dbcc20f 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -102,7 +102,7 @@ jobs: 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 + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 7302ce310b1..d4ab7e1d7b3 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -233,7 +233,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 @@ -290,7 +290,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions + run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/Cargo.lock b/Cargo.lock index 3ab47005352..fca6847c3db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2231,17 +2231,6 @@ dependencies = [ "unicode-width 0.1.14", ] -[[package]] -name = "codespan-reporting" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" -dependencies = [ - "serde", - "termcolor", - "unicode-width 0.1.14", -] - [[package]] name = "colorchoice" version = "1.0.4" @@ -3506,68 +3495,6 @@ dependencies = [ "cmov", ] -[[package]] -name = "cxx" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7620f6cfc4dcca21f2b085b7a890e16c60fd66f560cd69ee60594908dc72ab1" -dependencies = [ - "cc", - "cxx-build", - "cxxbridge-cmd", - "cxxbridge-flags", - "cxxbridge-macro", - "foldhash 0.2.0", - "link-cplusplus", -] - -[[package]] -name = "cxx-build" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a9bc1a22964ff6a355fbec24cf68266a0ed28f8b84c0864c386474ea3d0e479" -dependencies = [ - "cc", - "codespan-reporting 0.13.1", - "indexmap 2.14.2", - "proc-macro2", - "quote", - "scratch", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-cmd" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f29a879d35f7906e3c9b77d7a1005a6a0787d330c09dfe4ffb5f617728cb44" -dependencies = [ - "clap", - "codespan-reporting 0.13.1", - "indexmap 2.14.2", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-flags" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d67109015f93f683e364085aa6489a5b2118b4a40058482101d699936a7836d6" - -[[package]] -name = "cxxbridge-macro" -version = "1.0.190" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d187e019e7b05a1f3e69a8396b70800ee867aa9fc2ab972761173ccee03742df" -dependencies = [ - "indexmap 2.14.2", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "daemonize" version = "0.5.0" @@ -6897,7 +6824,6 @@ dependencies = [ "common-telemetry", "common-test-util", "criterion 0.7.0", - "datatypes", "fastbloom", "fst", "futures", @@ -6906,7 +6832,6 @@ dependencies = [ "jieba-rs", "lazy_static", "mockall", - "nalgebra", "pin-project", "prost 0.14.1", "puffin", @@ -6924,7 +6849,6 @@ dependencies = [ "tempfile", "tokio", "tokio-util", - "usearch", "uuid", ] @@ -7945,15 +7869,6 @@ dependencies = [ "vcpkg", ] -[[package]] -name = "link-cplusplus" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" -dependencies = [ - "cc", -] - [[package]] name = "link-section" version = "0.18.1" @@ -8694,7 +8609,6 @@ dependencies = [ "rand 0.9.4", "rayon", "regex", - "roaring", "rskafka", "rstest", "rstest_reuse", @@ -13160,12 +13074,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "scratch" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" - [[package]] name = "scrypt" version = "0.11.0" @@ -16100,16 +16008,6 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" -[[package]] -name = "usearch" -version = "2.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cc9fc5f872a3a4f9081d5f42624d788231b763e1846c829b9968a3755ac884d" -dependencies = [ - "cxx", - "cxx-build", -] - [[package]] name = "utf8-ranges" version = "1.0.5" @@ -16236,7 +16134,7 @@ dependencies = [ "ciborium", "cidr", "clap", - "codespan-reporting 0.12.0", + "codespan-reporting", "community-id", "convert_case 0.7.1", "crc", diff --git a/grafana/dashboards/metrics/cluster/dashboard.json b/grafana/dashboards/metrics/cluster/dashboard.json index bd22a0da597..ddbcc8e190c 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.json +++ b/grafana/dashboards/metrics/cluster/dashboard.json @@ -11583,7 +11583,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "hit-{{type}}", @@ -11596,7 +11596,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "miss-{{type}}", @@ -11609,7 +11609,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "evict-{{type}}-{{cause}}", diff --git a/grafana/dashboards/metrics/cluster/dashboard.md b/grafana/dashboards/metrics/cluster/dashboard.md index b8e050114c6..cc34ddad12a 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.md +++ b/grafana/dashboards/metrics/cluster/dashboard.md @@ -131,7 +131,7 @@ | Index Memory Usage | `greptime_index_apply_memory_usage`
`sum by (type) (greptime_index_create_memory_usage)` | `timeseries` | Memory used while applying and creating indexes. Growth here can explain memory pressure during indexed flush or compaction work. | `prometheus` | `bytes` | `apply` | | Index IO Bytes | `sum by (type, file_type) (rate(greptime_index_io_bytes_total[$__rate_interval]))` | `timeseries` | Index read and write byte rates by operation and file type for puffin and intermediate files. | `prometheus` | `Bps` | `{{type}}-{{file_type}}` | | Index IO Operations | `sum by (type, file_type) (rate(greptime_index_io_op_total[$__rate_interval]))` | `timeseries` | Index IO operation rates by operation and file type, including read, write, seek, and flush operations. | `prometheus` | `ops` | `{{type}}-{{file_type}}` | -| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | +| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | # Metasrv | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | diff --git a/grafana/dashboards/metrics/cluster/dashboard.yaml b/grafana/dashboards/metrics/cluster/dashboard.yaml index cd105edfcd2..9664db35f69 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.yaml +++ b/grafana/dashboards/metrics/cluster/dashboard.yaml @@ -1572,17 +1572,17 @@ groups: description: Index-related cache hits, misses, and evictions from Mito caches. unit: ops queries: - - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: hit-{{type}} - - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: miss-{{type}} - - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} diff --git a/grafana/dashboards/metrics/standalone/dashboard.json b/grafana/dashboards/metrics/standalone/dashboard.json index a90d05760e4..9d181b05ffe 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.json +++ b/grafana/dashboards/metrics/standalone/dashboard.json @@ -11583,7 +11583,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_hit{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "hit-{{type}}", @@ -11596,7 +11596,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type) (rate(greptime_mito_cache_miss{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "miss-{{type}}", @@ -11609,7 +11609,7 @@ "uid": "${metrics}" }, "editorMode": "code", - "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|vector_index|index_result\"}[$__rate_interval]))", + "expr": "sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~\"index.*|index_result\"}[$__rate_interval]))", "hide": false, "instant": false, "legendFormat": "evict-{{type}}-{{cause}}", diff --git a/grafana/dashboards/metrics/standalone/dashboard.md b/grafana/dashboards/metrics/standalone/dashboard.md index 2ed84c59213..dd00fb93d49 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.md +++ b/grafana/dashboards/metrics/standalone/dashboard.md @@ -131,7 +131,7 @@ | Index Memory Usage | `greptime_index_apply_memory_usage`
`sum by (type) (greptime_index_create_memory_usage)` | `timeseries` | Memory used while applying and creating indexes. Growth here can explain memory pressure during indexed flush or compaction work. | `prometheus` | `bytes` | `apply` | | Index IO Bytes | `sum by (type, file_type) (rate(greptime_index_io_bytes_total[$__rate_interval]))` | `timeseries` | Index read and write byte rates by operation and file type for puffin and intermediate files. | `prometheus` | `Bps` | `{{type}}-{{file_type}}` | | Index IO Operations | `sum by (type, file_type) (rate(greptime_index_io_op_total[$__rate_interval]))` | `timeseries` | Index IO operation rates by operation and file type, including read, write, seek, and flush operations. | `prometheus` | `ops` | `{{type}}-{{file_type}}` | -| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|vector_index\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | +| Index Cache | `sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*\|index_result"}[$__rate_interval]))`
`sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*\|index_result"}[$__rate_interval]))` | `timeseries` | Index-related cache hits, misses, and evictions from Mito caches. | `prometheus` | `ops` | `hit-{{type}}` | # Metasrv | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | diff --git a/grafana/dashboards/metrics/standalone/dashboard.yaml b/grafana/dashboards/metrics/standalone/dashboard.yaml index 8da2a8fa6f1..86076d3d33a 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.yaml +++ b/grafana/dashboards/metrics/standalone/dashboard.yaml @@ -1572,17 +1572,17 @@ groups: description: Index-related cache hits, misses, and evictions from Mito caches. unit: ops queries: - - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_hit{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: hit-{{type}} - - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type) (rate(greptime_mito_cache_miss{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} legendFormat: miss-{{type}} - - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|vector_index|index_result"}[$__rate_interval])) + - expr: sum by (type, cause) (rate(greptime_mito_cache_eviction{type=~"index.*|index_result"}[$__rate_interval])) datasource: type: prometheus uid: ${metrics} diff --git a/src/api/src/v1/column_def.rs b/src/api/src/v1/column_def.rs index 9b39d0f3296..b02b06feafd 100644 --- a/src/api/src/v1/column_def.rs +++ b/src/api/src/v1/column_def.rs @@ -18,7 +18,7 @@ use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_K use datatypes::schema::{ COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, - SkippingIndexOptions, SkippingIndexType, VECTOR_INDEX_KEY, + SkippingIndexOptions, SkippingIndexType, }; use greptime_proto::v1::{ Analyzer, FulltextBackend as PbFulltextBackend, SkippingIndexType as PbSkippingIndexType, @@ -35,14 +35,11 @@ const FULLTEXT_GRPC_KEY: &str = "fulltext"; const INVERTED_INDEX_GRPC_KEY: &str = "inverted_index"; /// Key used to store skip index options in gRPC column options. const SKIPPING_INDEX_GRPC_KEY: &str = "skipping_index"; -/// Key used to store vector index options in gRPC column options. -const VECTOR_INDEX_GRPC_KEY: &str = "vector_index"; -const COLUMN_OPTION_MAPPINGS: [(&str, &str); 6] = [ +const COLUMN_OPTION_MAPPINGS: [(&str, &str); 5] = [ (FULLTEXT_GRPC_KEY, FULLTEXT_KEY), (INVERTED_INDEX_GRPC_KEY, INVERTED_INDEX_KEY), (SKIPPING_INDEX_GRPC_KEY, SKIPPING_INDEX_KEY), - (VECTOR_INDEX_GRPC_KEY, VECTOR_INDEX_KEY), (EXTENSION_TYPE_NAME_KEY, EXTENSION_TYPE_NAME_KEY), (EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_METADATA_KEY), ]; @@ -80,9 +77,6 @@ pub fn try_as_column_schema(column_def: &ColumnDef) -> Result { if let Some(skipping_index) = options.options.get(SKIPPING_INDEX_GRPC_KEY) { metadata.insert(SKIPPING_INDEX_KEY.to_string(), skipping_index.to_owned()); } - if let Some(vector_index) = options.options.get(VECTOR_INDEX_GRPC_KEY) { - metadata.insert(VECTOR_INDEX_KEY.to_string(), vector_index.to_owned()); - } if let Some(extension_name) = options.options.get(EXTENSION_TYPE_NAME_KEY) { metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone()); } @@ -178,11 +172,6 @@ pub fn options_from_column_schema(column_schema: &ColumnSchema) -> Option Skipp mod tests { use datatypes::data_type::ConcreteDataType; - use datatypes::schema::{ - FulltextAnalyzer, FulltextBackend, VectorDistanceMetric, VectorIndexOptions, - }; - use serde_json::json; + use datatypes::schema::{FulltextAnalyzer, FulltextBackend}; use super::*; use crate::v1::ColumnDataType; @@ -297,10 +283,6 @@ mod tests { "{\"enable\":true}".to_string(), ), (INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string()), - ( - VECTOR_INDEX_GRPC_KEY.to_string(), - "{\"engine\":\"usearch\",\"metric\":\"l2sq\",\"connectivity\":16,\"expansion-add\":128,\"expansion-search\":64}".to_string(), - ), ]), }), }; @@ -323,8 +305,6 @@ mod tests { } ); assert!(schema.is_inverted_indexed()); - let vector_options = schema.vector_index_options().unwrap().unwrap(); - assert_eq!(vector_options.metric, VectorDistanceMetric::L2sq); } #[test] @@ -355,29 +335,6 @@ mod tests { ); } - #[test] - fn test_vector_index_options_roundtrip() { - let schema = ColumnSchema::new("test", ConcreteDataType::vector_datatype(4), true) - .with_vector_index_options(&VectorIndexOptions::default()) - .unwrap(); - let column_def = try_as_column_def(&schema, false).unwrap(); - let roundtrip = try_as_column_schema(&column_def).unwrap(); - let options = roundtrip.vector_index_options().unwrap().unwrap(); - assert_eq!(options.metric, VectorDistanceMetric::L2sq); - - let options = column_def.options.unwrap(); - let raw = options.options.get(VECTOR_INDEX_GRPC_KEY).unwrap(); - let json_value: serde_json::Value = serde_json::from_str(raw).unwrap(); - let expected = json!({ - "engine": "usearch", - "metric": "l2sq", - "connectivity": 16, - "expansion-add": 128, - "expansion-search": 64 - }); - assert_eq!(json_value, expected); - } - #[test] fn test_options_with_fulltext() { let fulltext = FulltextOptions::new_unchecked( diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 367b86809d8..a67dd895999 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -34,7 +34,6 @@ enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterpr dev-tools = [] mysql-object-store = ["object-store/mysql-object-store"] tokio-console = ["common-telemetry/tokio-console"] -vector_index = ["mito2/vector_index", "query/vector_index"] [lints] workspace = true diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index 14e289ea178..d9f7a00c8c4 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -209,8 +209,6 @@ impl ObjbenchCommand { fulltext_index_config, bloom_filter_index_config: MitoConfig::default().bloom_filter_index, preserve_row_sequence: false, - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; // Write SST diff --git a/src/datatypes/src/schema.rs b/src/datatypes/src/schema.rs index d53b42e3e7e..0f1a832ebd2 100644 --- a/src/datatypes/src/schema.rs +++ b/src/datatypes/src/schema.rs @@ -32,13 +32,9 @@ pub use crate::schema::column_schema::{ COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, - COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, - COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, COMMENT_KEY, - ColumnExtType, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, - INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, - TIME_INDEX_KEY, VECTOR_INDEX_KEY, VectorDistanceMetric, VectorIndexEngineType, - VectorIndexOptions, + COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, ColumnExtType, ColumnSchema, FULLTEXT_KEY, + FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, + SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY, }; pub use crate::schema::constraint::ColumnDefaultConstraint; diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index a71605afdb2..b3ce9cafedb 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/src/datatypes/src/schema/column_schema.rs @@ -48,8 +48,6 @@ pub const FULLTEXT_KEY: &str = "greptime:fulltext"; pub const INVERTED_INDEX_KEY: &str = "greptime:inverted_index"; /// Key used to store skip options in arrow field's metadata. pub const SKIPPING_INDEX_KEY: &str = "greptime:skipping_index"; -/// Key used to store vector index options in arrow field's metadata. -pub const VECTOR_INDEX_KEY: &str = "greptime:vector_index"; /// Keys used in fulltext options pub const COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE: &str = "enable"; @@ -64,13 +62,6 @@ pub const COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY: &str = "granularity"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE: &str = "false_positive_rate"; pub const COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE: &str = "type"; -/// Keys used in VECTOR index options -pub const COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE: &str = "engine"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_METRIC: &str = "metric"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY: &str = "connectivity"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD: &str = "expansion_add"; -pub const COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH: &str = "expansion_search"; - pub const DEFAULT_GRANULARITY: u32 = 10240; pub const DEFAULT_FALSE_POSITIVE_RATE: f64 = 0.01; @@ -255,53 +246,6 @@ impl ColumnSchema { self.metadata.contains_key(INVERTED_INDEX_KEY) } - /// Checks if this column has a vector index. - pub fn is_vector_indexed(&self) -> bool { - match self.vector_index_options() { - Ok(opts) => opts.is_some(), - Err(e) => { - common_telemetry::warn!( - "Failed to deserialize vector_index_options for column '{}': {}", - self.name, - e - ); - false - } - } - } - - /// Gets the vector index options. - pub fn vector_index_options(&self) -> Result> { - match self.metadata.get(VECTOR_INDEX_KEY) { - None => Ok(None), - Some(json) => { - let options = - serde_json::from_str(json).context(error::DeserializeSnafu { json })?; - Ok(Some(options)) - } - } - } - - /// Sets the vector index options. - pub fn set_vector_index_options(&mut self, options: &VectorIndexOptions) -> Result<()> { - self.metadata.insert( - VECTOR_INDEX_KEY.to_string(), - serde_json::to_string(options).context(error::SerializeSnafu)?, - ); - Ok(()) - } - - /// Removes the vector index options. - pub fn unset_vector_index_options(&mut self) { - self.metadata.remove(VECTOR_INDEX_KEY); - } - - /// Sets vector index options and returns self for chaining. - pub fn with_vector_index_options(mut self, options: &VectorIndexOptions) -> Result { - self.set_vector_index_options(options)?; - Ok(self) - } - /// Set default constraint. /// /// If a default constraint exists for the column, this method will @@ -1067,181 +1011,6 @@ impl TryFrom> for SkippingIndexOptions { } } -/// Distance metric for vector similarity search. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, Visit, VisitMut)] -#[serde(rename_all = "lowercase")] -pub enum VectorDistanceMetric { - /// Squared Euclidean distance (L2^2). - #[default] - L2sq, - /// Cosine distance (1 - cosine similarity). - Cosine, - /// Inner product (negative, for maximum inner product search). - #[serde(alias = "ip")] - InnerProduct, -} - -impl fmt::Display for VectorDistanceMetric { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - VectorDistanceMetric::L2sq => write!(f, "l2sq"), - VectorDistanceMetric::Cosine => write!(f, "cosine"), - VectorDistanceMetric::InnerProduct => write!(f, "ip"), - } - } -} - -impl std::str::FromStr for VectorDistanceMetric { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - match s.to_lowercase().as_str() { - "l2sq" | "l2" | "euclidean" => Ok(VectorDistanceMetric::L2sq), - "cosine" | "cos" => Ok(VectorDistanceMetric::Cosine), - "inner_product" | "ip" | "dot" => Ok(VectorDistanceMetric::InnerProduct), - _ => Err(format!( - "Unknown distance metric: {}. Expected: l2sq, cosine, or ip", - s - )), - } - } -} - -impl VectorDistanceMetric { - /// Returns the metric as u8 for blob serialization. - pub fn as_u8(&self) -> u8 { - match self { - Self::L2sq => 0, - Self::Cosine => 1, - Self::InnerProduct => 2, - } - } - - /// Parses metric from u8 (used when reading blob). - pub fn try_from_u8(v: u8) -> Option { - match v { - 0 => Some(Self::L2sq), - 1 => Some(Self::Cosine), - 2 => Some(Self::InnerProduct), - _ => None, - } - } -} - -/// Default HNSW connectivity parameter. -const DEFAULT_VECTOR_INDEX_CONNECTIVITY: u32 = 16; -/// Default expansion factor during index construction. -const DEFAULT_VECTOR_INDEX_EXPANSION_ADD: u32 = 128; -/// Default expansion factor during search. -const DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH: u32 = 64; - -fn default_vector_index_connectivity() -> u32 { - DEFAULT_VECTOR_INDEX_CONNECTIVITY -} - -fn default_vector_index_expansion_add() -> u32 { - DEFAULT_VECTOR_INDEX_EXPANSION_ADD -} - -fn default_vector_index_expansion_search() -> u32 { - DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH -} - -/// Supported vector index engine types. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, Visit, VisitMut)] -#[serde(rename_all = "lowercase")] -pub enum VectorIndexEngineType { - /// USearch HNSW implementation. - #[default] - Usearch, - // Future: Vsag, -} - -impl VectorIndexEngineType { - /// Returns the engine type as u8 for blob serialization. - pub fn as_u8(&self) -> u8 { - match self { - Self::Usearch => 0, - } - } - - /// Parses engine type from u8 (used when reading blob). - pub fn try_from_u8(v: u8) -> Option { - match v { - 0 => Some(Self::Usearch), - _ => None, - } - } -} - -impl fmt::Display for VectorIndexEngineType { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Usearch => write!(f, "usearch"), - } - } -} - -impl std::str::FromStr for VectorIndexEngineType { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - match s.to_lowercase().as_str() { - "usearch" => Ok(Self::Usearch), - _ => Err(format!( - "Unknown vector index engine: {}. Expected: usearch", - s - )), - } - } -} - -/// Options for vector index (HNSW). -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)] -#[serde(rename_all = "kebab-case")] -pub struct VectorIndexOptions { - /// Vector index engine type (default: usearch). - #[serde(default)] - pub engine: VectorIndexEngineType, - /// Distance metric for similarity search. - #[serde(default)] - pub metric: VectorDistanceMetric, - /// HNSW connectivity parameter (M in the paper). - /// Higher values improve recall but increase memory usage. - #[serde(default = "default_vector_index_connectivity")] - pub connectivity: u32, - /// Expansion factor during index construction (ef_construction). - /// Higher values improve index quality but slow down construction. - #[serde(default = "default_vector_index_expansion_add")] - pub expansion_add: u32, - /// Expansion factor during search (ef_search). - /// Higher values improve recall but slow down search. - #[serde(default = "default_vector_index_expansion_search")] - pub expansion_search: u32, -} - -impl Default for VectorIndexOptions { - fn default() -> Self { - Self { - engine: VectorIndexEngineType::default(), - metric: VectorDistanceMetric::default(), - connectivity: DEFAULT_VECTOR_INDEX_CONNECTIVITY, - expansion_add: DEFAULT_VECTOR_INDEX_EXPANSION_ADD, - expansion_search: DEFAULT_VECTOR_INDEX_EXPANSION_SEARCH, - } - } -} - -impl fmt::Display for VectorIndexOptions { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "engine={}, metric={}, connectivity={}, expansion_add={}, expansion_search={}", - self.engine, self.metric, self.connectivity, self.expansion_add, self.expansion_search - ) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index b8518112610..bb1225c12c8 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -7,9 +7,6 @@ license.workspace = true [lints] workspace = true -[features] -vector_index = ["dep:nalgebra", "dep:usearch"] - [dependencies] async-trait.workspace = true asynchronous-codec = "0.7.0" @@ -20,7 +17,6 @@ common-error.workspace = true common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -datatypes.workspace = true fastbloom = "0.8" fst.workspace = true futures.workspace = true @@ -29,7 +25,6 @@ itertools.workspace = true jieba-rs = "0.10" lazy_static.workspace = true mockall.workspace = true -nalgebra = { workspace = true, optional = true } pin-project.workspace = true prost.workspace = true puffin.workspace = true @@ -44,7 +39,6 @@ tantivy = { version = "0.26", features = ["zstd-compression"] } tantivy-jieba = "0.20" tokio.workspace = true tokio-util.workspace = true -usearch = { version = "2.21", default-features = false, features = ["fp16lib"], optional = true } uuid.workspace = true [dev-dependencies] diff --git a/src/index/src/lib.rs b/src/index/src/lib.rs index c3ae35e9072..2c1ceb4e3f6 100644 --- a/src/index/src/lib.rs +++ b/src/index/src/lib.rs @@ -22,8 +22,6 @@ pub mod external_provider; pub mod fulltext_index; pub mod inverted_index; pub mod target; -#[cfg(feature = "vector_index")] -pub mod vector; pub type Bytes = Vec; pub type BytesRef<'a> = &'a [u8]; diff --git a/src/index/src/vector.rs b/src/index/src/vector.rs deleted file mode 100644 index 9c7cfded330..00000000000 --- a/src/index/src/vector.rs +++ /dev/null @@ -1,179 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index implementation for KNN search. -//! -//! This module provides the core algorithms for vector index creation, reading, -//! and search operations. The architecture follows the same pattern as other indexes: -//! -//! - `error.rs`: Error types for vector index operations -//! - `format.rs`: Proto conversion helpers for blob format -//! - `util.rs`: Byte utilities for vector data -//! - `engine/`: Pluggable vector index engines (e.g., USearch HNSW) -//! - `create.rs`: `VectorIndexCreator` trait and implementations for building indices -//! - `apply.rs`: `VectorIndexApplier` trait and implementations for searching indices - -pub mod apply; -pub mod create; -pub mod engine; -pub mod error; -pub mod format; -pub mod util; - -pub use apply::{HnswVectorIndexApplier, IndexApplier, VectorSearchOutput}; -pub use create::{HnswVectorIndexCreator, IndexCreator}; -pub use datatypes::schema::{VectorDistanceMetric, VectorIndexOptions}; -use nalgebra::DVectorView; -pub use usearch::MetricKind; - -/// Converts a VectorDistanceMetric to a USearch MetricKind. -pub fn distance_metric_to_usearch(metric: VectorDistanceMetric) -> MetricKind { - match metric { - VectorDistanceMetric::L2sq => MetricKind::L2sq, - VectorDistanceMetric::Cosine => MetricKind::Cos, - VectorDistanceMetric::InnerProduct => MetricKind::IP, - } -} - -/// Computes distance between two vectors using the specified metric. -/// -/// Uses SIMD-optimized implementations via nalgebra. -/// -/// **Note:** The caller must ensure that the two vectors have the same length -/// and are non-empty. Empty vectors return 0.0 for all metrics. -pub fn compute_distance(v1: &[f32], v2: &[f32], metric: VectorDistanceMetric) -> f32 { - // Empty vectors are degenerate; return 0.0 uniformly across all metrics. - if v1.is_empty() || v2.is_empty() { - return 0.0; - } - - match metric { - VectorDistanceMetric::L2sq => l2sq(v1, v2), - VectorDistanceMetric::Cosine => cosine(v1, v2), - VectorDistanceMetric::InnerProduct => -dot(v1, v2), - } -} - -/// Calculates the squared L2 distance between two vectors. -fn l2sq(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs = DVectorView::from_slice(lhs, lhs.len()); - let rhs = DVectorView::from_slice(rhs, rhs.len()); - (lhs - rhs).norm_squared() -} - -/// Calculates the cosine distance between two vectors. -/// -/// Returns a value in `[0.0, 2.0]` where 0.0 means identical direction and 2.0 means -/// opposite direction. For degenerate cases (zero or near-zero magnitude vectors), -/// returns 1.0 (maximum uncertainty) to avoid NaN and ensure safe index operations. -fn cosine(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs_vec = DVectorView::from_slice(lhs, lhs.len()); - let rhs_vec = DVectorView::from_slice(rhs, rhs.len()); - - let dot_product = lhs_vec.dot(&rhs_vec); - let lhs_norm = lhs_vec.norm(); - let rhs_norm = rhs_vec.norm(); - - // Zero-magnitude vectors have undefined direction; return max distance as safe fallback. - if dot_product.abs() < f32::EPSILON - || lhs_norm.abs() < f32::EPSILON - || rhs_norm.abs() < f32::EPSILON - { - return 1.0; - } - - let cos_similar = dot_product / (lhs_norm * rhs_norm); - let res = 1.0 - cos_similar; - // Clamp near-zero results to exactly 0.0 to avoid floating-point artifacts. - if res.abs() < f32::EPSILON { 0.0 } else { res } -} - -/// Calculates the dot product between two vectors. -fn dot(lhs: &[f32], rhs: &[f32]) -> f32 { - let lhs = DVectorView::from_slice(lhs, lhs.len()); - let rhs = DVectorView::from_slice(rhs, rhs.len()); - lhs.dot(&rhs) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_distance_metric_to_usearch() { - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::L2sq), - MetricKind::L2sq - ); - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::Cosine), - MetricKind::Cos - ); - assert_eq!( - distance_metric_to_usearch(VectorDistanceMetric::InnerProduct), - MetricKind::IP - ); - } - - #[test] - fn test_vector_index_options_default() { - let options = VectorIndexOptions::default(); - assert_eq!(options.metric, VectorDistanceMetric::L2sq); - assert_eq!(options.connectivity, 16); - assert_eq!(options.expansion_add, 128); - assert_eq!(options.expansion_search, 64); - } - - #[test] - fn test_compute_distance_l2sq() { - let v1 = vec![1.0, 2.0, 3.0]; - let v2 = vec![4.0, 5.0, 6.0]; - // L2sq = (4-1)^2 + (5-2)^2 + (6-3)^2 = 9 + 9 + 9 = 27 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::L2sq); - assert!((dist - 27.0).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_cosine() { - let v1 = vec![1.0, 0.0, 0.0]; - let v2 = vec![0.0, 1.0, 0.0]; - // Orthogonal vectors have cosine similarity of 0, distance of 1 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::Cosine); - assert!((dist - 1.0).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_inner_product() { - let v1 = vec![1.0, 2.0, 3.0]; - let v2 = vec![4.0, 5.0, 6.0]; - // Inner product = 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32 - // Distance is negated: -32 - let dist = compute_distance(&v1, &v2, VectorDistanceMetric::InnerProduct); - assert!((dist - (-32.0)).abs() < 1e-6); - } - - #[test] - fn test_compute_distance_empty_vectors() { - // Empty vectors should return 0.0 uniformly for all metrics - assert_eq!(compute_distance(&[], &[], VectorDistanceMetric::L2sq), 0.0); - assert_eq!( - compute_distance(&[], &[], VectorDistanceMetric::Cosine), - 0.0 - ); - assert_eq!( - compute_distance(&[], &[], VectorDistanceMetric::InnerProduct), - 0.0 - ); - } -} diff --git a/src/index/src/vector/apply.rs b/src/index/src/vector/apply.rs deleted file mode 100644 index 82db03b9efd..00000000000 --- a/src/index/src/vector/apply.rs +++ /dev/null @@ -1,428 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index application trait and implementation. - -use greptime_proto::v1::index::VectorIndexMeta; -use prost::Message; -use roaring::RoaringBitmap; -use snafu::ResultExt; -use store_api::storage::VectorIndexEngine; - -use crate::vector::VectorDistanceMetric; -use crate::vector::engine::{self, VectorIndexConfig}; -use crate::vector::error::{ - BlobTooSmallSnafu, BlobTruncatedSnafu, DecodeProtoSnafu, DeserializeBitmapSnafu, EngineSnafu, - InvalidBlobSnafu, KeyMappingSnafu, Result, UnknownDistanceMetricSnafu, UnknownEngineTypeSnafu, -}; -use crate::vector::format::{META_SIZE_LEN, distance_metric_from_proto, engine_type_from_proto}; - -/// Output of a vector index search. -#[derive(Debug, Clone)] -pub struct VectorSearchOutput { - /// Row offsets in the original data. - pub row_offsets: Vec, - /// Distances to query vector. - pub distances: Vec, -} - -/// `IndexApplier` is for searching a vector index. -pub trait IndexApplier: Send + Sync { - /// Searches for k nearest neighbors. - fn search(&self, query: &[f32], k: usize) -> Result; - - /// Returns the memory usage. - fn memory_usage(&self) -> usize; - - /// Returns the dimension of vectors in this index. - fn dimensions(&self) -> u32; - - /// Returns the distance metric used for this index. - fn metric(&self) -> VectorDistanceMetric; - - /// Returns the total number of rows (including NULLs). - fn total_rows(&self) -> u64; - - /// Returns the number of indexed rows (excluding NULLs). - fn indexed_rows(&self) -> u64; -} - -/// HNSW-based vector index applier. -pub struct HnswVectorIndexApplier { - /// The loaded vector index engine. - engine: Box, - /// Bitmap tracking which row offsets have NULL vectors. - null_bitmap: RoaringBitmap, - /// The dimension of vectors in this index. - dimensions: u32, - /// The distance metric used for this index. - metric: VectorDistanceMetric, - /// Total number of rows (including NULLs). - total_rows: u64, - /// Number of indexed rows (excluding NULLs). - indexed_rows: u64, -} - -impl HnswVectorIndexApplier { - /// Creates a new HNSW vector index applier from blob data. - /// - /// Blob layout: - /// ```text - /// +----------------------+ - /// | Null Bitmap | <- RoaringBitmap serialized - /// +----------------------+ - /// | Index Data | <- engine-specific serialized - /// +----------------------+ - /// | Meta (protobuf) | <- VectorIndexMeta encoded - /// +----------------------+ - /// | Meta Size (4B) | <- u32 LE - /// +----------------------+ - /// ``` - pub fn from_blob(data: &[u8]) -> Result { - if data.len() < META_SIZE_LEN { - return BlobTooSmallSnafu { - min_size: META_SIZE_LEN, - actual_size: data.len(), - } - .fail(); - } - - // Read meta size from the end - let blob_len = data.len(); - let meta_size = - u32::from_le_bytes(data[blob_len - META_SIZE_LEN..].try_into().unwrap()) as usize; - - if blob_len < META_SIZE_LEN + meta_size { - return BlobTruncatedSnafu { - reason: "blob truncated while reading meta", - } - .fail(); - } - - // Read and decode meta - let meta_start = blob_len - META_SIZE_LEN - meta_size; - let meta = VectorIndexMeta::decode(&data[meta_start..blob_len - META_SIZE_LEN]) - .context(DecodeProtoSnafu)?; - - let null_bitmap_size = meta.null_bitmap_size as usize; - let index_size = meta.index_size as usize; - - if null_bitmap_size + index_size > meta_start { - return InvalidBlobSnafu { - reason: "invalid size fields", - } - .fail(); - } - - // Read null bitmap - let null_bitmap = RoaringBitmap::deserialize_from(&data[..null_bitmap_size]) - .context(DeserializeBitmapSnafu)?; - - // Read index data - let index_start = null_bitmap_size; - let index_end = index_start + index_size; - let index_bytes = &data[index_start..index_end]; - - let engine_proto = meta.engine.try_into().map_err(|_| { - UnknownEngineTypeSnafu { - engine_type: meta.engine, - } - .build() - })?; - let engine_type = engine_type_from_proto(engine_proto).ok_or_else(|| { - UnknownEngineTypeSnafu { - engine_type: meta.engine, - } - .build() - })?; - - let metric_proto = meta.metric.try_into().map_err(|_| { - UnknownDistanceMetricSnafu { - metric: meta.metric, - } - .build() - })?; - let distance_metric = distance_metric_from_proto(metric_proto).ok_or_else(|| { - UnknownDistanceMetricSnafu { - metric: meta.metric, - } - .build() - })?; - - let config = VectorIndexConfig { - engine: engine_type, - dim: meta.dim as usize, - distance_metric, - connectivity: meta.connectivity as usize, - expansion_add: meta.expansion_add as usize, - expansion_search: meta.expansion_search as usize, - }; - - let engine = engine::load_engine(engine_type, &config, index_bytes)?; - - let stats = meta.stats.ok_or_else(|| { - InvalidBlobSnafu { - reason: "missing stats in vector index metadata", - } - .build() - })?; - - Ok(Self { - engine, - null_bitmap, - dimensions: meta.dim, - metric: distance_metric, - total_rows: stats.total_row_count, - indexed_rows: stats.indexed_row_count, - }) - } - - /// Maps HNSW keys from search results to row offsets in the original data. - fn map_keys_to_row_offsets(&self, keys: Vec) -> Result> { - if self.total_rows == 0 { - return Ok(Vec::new()); - } - let total_rows_u32 = u32::try_from(self.total_rows).map_err(|_| { - KeyMappingSnafu { - reason: format!("Total rows {} exceeds u32::MAX", self.total_rows), - } - .build() - })?; - - let mut row_offsets = Vec::with_capacity(keys.len()); - for key in keys { - let offset = self.hnsw_key_to_row_offset(total_rows_u32, key)?; - row_offsets.push(offset as u64); - } - Ok(row_offsets) - } - - /// Converts a single HNSW key to its corresponding row offset. - fn hnsw_key_to_row_offset(&self, total_rows: u32, key: u64) -> Result { - if total_rows == 0 { - return KeyMappingSnafu { - reason: "Total rows is zero", - } - .fail(); - } - if key >= self.indexed_rows { - return KeyMappingSnafu { - reason: format!( - "HNSW key {} exceeds indexed rows {}", - key, self.indexed_rows - ), - } - .fail(); - } - - // Fast path: no NULLs means key == row offset - if self.null_bitmap.is_empty() { - return Ok(key as u32); - } - - // Binary search to find the row offset - let mut left: u32 = 0; - let mut right: u32 = total_rows - 1; - while left <= right { - let mid = left + (right - left) / 2; - let nulls_before = self.null_bitmap.rank(mid); - let non_nulls = (mid as u64 + 1).saturating_sub(nulls_before); - if non_nulls > key { - if mid == 0 { - break; - } - right = mid - 1; - } else { - left = mid + 1; - } - } - - if left >= total_rows { - return KeyMappingSnafu { - reason: "Failed to map HNSW key to row offset", - } - .fail(); - } - - Ok(left) - } -} - -impl IndexApplier for HnswVectorIndexApplier { - fn search(&self, query: &[f32], k: usize) -> Result { - if self.indexed_rows == 0 || k == 0 { - return Ok(VectorSearchOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - } - - let actual_k = k.min(self.indexed_rows as usize); - let matches = self.engine.search(query, actual_k).map_err(|e| { - EngineSnafu { - reason: e.to_string(), - } - .build() - })?; - - let row_offsets = self.map_keys_to_row_offsets(matches.keys)?; - - Ok(VectorSearchOutput { - row_offsets, - distances: matches.distances, - }) - } - - fn memory_usage(&self) -> usize { - self.engine.memory_usage() - + self.null_bitmap.serialized_size() - + std::mem::size_of::() - } - - fn dimensions(&self) -> u32 { - self.dimensions - } - - fn metric(&self) -> VectorDistanceMetric { - self.metric - } - - fn total_rows(&self) -> u64 { - self.total_rows - } - - fn indexed_rows(&self) -> u64 { - self.indexed_rows - } -} - -#[cfg(test)] -mod tests { - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use prost::Message; - use roaring::RoaringBitmap; - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::engine::VectorIndexConfig; - use crate::vector::format::{distance_metric_to_proto, engine_type_to_proto}; - - fn build_test_blob( - config: &VectorIndexConfig, - vectors: Vec<(u64, Vec)>, - null_bitmap: &RoaringBitmap, - total_rows: u64, - indexed_rows: u64, - ) -> Vec { - let mut engine = engine::create_engine(config.engine, config).unwrap(); - for (key, vector) in vectors { - engine.add(key, &vector).unwrap(); - } - let index_size = engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - engine.save_to_buffer(&mut index_bytes).unwrap(); - - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let null_count = total_rows - indexed_rows; - let meta = VectorIndexMeta { - engine: engine_type_to_proto(config.engine).into(), - dim: config.dim as u32, - metric: distance_metric_to_proto(config.distance_metric).into(), - connectivity: config.connectivity as u32, - expansion_add: config.expansion_add as u32, - expansion_search: config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(&index_bytes); - blob.extend_from_slice(&meta_bytes); - blob.extend_from_slice(&meta_size.to_le_bytes()); - blob - } - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 2, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_hnsw_vector_index_applier_search() { - let config = test_config(); - let mut null_bitmap = RoaringBitmap::new(); - null_bitmap.insert(1); - - let blob = build_test_blob( - &config, - vec![(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])], - &null_bitmap, - 3, - 2, - ); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - assert_eq!(applier.dimensions(), 2); - assert_eq!(applier.metric(), VectorDistanceMetric::L2sq); - assert_eq!(applier.total_rows(), 3); - assert_eq!(applier.indexed_rows(), 2); - - let output = applier.search(&[1.0, 0.0], 2).unwrap(); - assert_eq!(output.row_offsets.len(), 2); - // First result should be row 0 (closest to [1.0, 0.0]) - assert_eq!(output.row_offsets[0], 0); - // Second result should be row 2 (maps from HNSW key 1) - assert_eq!(output.row_offsets[1], 2); - } - - #[test] - fn test_hnsw_vector_index_applier_empty() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - - let blob = build_test_blob(&config, vec![], &null_bitmap, 0, 0); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - let output = applier.search(&[1.0, 0.0], 1).unwrap(); - assert!(output.row_offsets.is_empty()); - } - - #[test] - fn test_hnsw_vector_index_applier_k_zero() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - - let blob = build_test_blob(&config, vec![(0, vec![1.0, 0.0])], &null_bitmap, 1, 1); - - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - let output = applier.search(&[1.0, 0.0], 0).unwrap(); - assert!(output.row_offsets.is_empty()); - } -} diff --git a/src/index/src/vector/create.rs b/src/index/src/vector/create.rs deleted file mode 100644 index 006bd786908..00000000000 --- a/src/index/src/vector/create.rs +++ /dev/null @@ -1,345 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index creation trait and implementation. - -use async_trait::async_trait; -use common_error::ext::BoxedError; -use futures::io::Cursor; -use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; -use prost::Message; -use puffin::puffin_manager::{PuffinWriter, PutOptions}; -use roaring::RoaringBitmap; -use snafu::ResultExt; -use store_api::storage::{VectorIndexEngine, VectorIndexEngineType}; - -use crate::vector::engine::{self, VectorIndexConfig}; -use crate::vector::error::{ - EngineSnafu, ExternalSnafu, FinishSnafu, Result, RowCountOverflowSnafu, SerializeBitmapSnafu, -}; -use crate::vector::format::{META_SIZE_LEN, distance_metric_to_proto, engine_type_to_proto}; - -/// `IndexCreator` is for creating a vector index. -#[async_trait] -pub trait IndexCreator: Send { - /// Pushes a vector to the index. - fn push_vector(&mut self, vector: &[f32]) -> Result<()>; - - /// Records a NULL vector. - fn push_null(&mut self) -> Result<()>; - - /// Records multiple NULL vectors. - fn push_nulls(&mut self, n: usize) -> Result<()>; - - /// Finalizes the index and writes to puffin. - /// Returns bytes written. - async fn finish( - &mut self, - puffin_writer: &mut (impl PuffinWriter + Send), - blob_key: &str, - put_options: PutOptions, - ) -> Result; - - /// Aborts the creation. - async fn abort(&mut self) -> Result<()>; - - /// Returns the memory usage. - fn memory_usage(&self) -> usize; - - /// Returns the number of indexed vectors. - fn size(&self) -> usize; -} - -/// HNSW-based vector index creator. -pub struct HnswVectorIndexCreator { - /// The vector index engine (e.g., USearch HNSW). - engine: Box, - /// Configuration for this index. - config: VectorIndexConfig, - /// Bitmap tracking which row offsets have NULL vectors. - null_bitmap: RoaringBitmap, - /// Current row offset (including NULLs). - current_row_offset: u64, - /// Next HNSW key to assign (only for non-NULL vectors). - next_hnsw_key: u64, - /// Memory usage estimation. - memory_usage_bytes: usize, - /// Whether the index has been finished. - finished: bool, -} - -impl HnswVectorIndexCreator { - /// Creates a new HNSW vector index creator. - pub fn new(config: VectorIndexConfig) -> Result { - let engine_instance = engine::create_engine(config.engine, &config)?; - - Ok(Self { - engine: engine_instance, - config, - null_bitmap: RoaringBitmap::new(), - current_row_offset: 0, - next_hnsw_key: 0, - memory_usage_bytes: 0, - finished: false, - }) - } - - /// Reserves capacity for the expected number of vectors. - #[allow(dead_code)] - pub fn reserve(&mut self, capacity: usize) -> Result<()> { - self.engine.reserve(capacity).map_err(|e| { - EngineSnafu { - reason: format!("Failed to reserve capacity: {}", e), - } - .build() - }) - } - - /// Returns the engine type. - pub fn engine_type(&self) -> VectorIndexEngineType { - self.config.engine - } - - /// Returns the configuration. - pub fn config(&self) -> &VectorIndexConfig { - &self.config - } - - /// Returns the null bitmap. - pub fn null_bitmap(&self) -> &RoaringBitmap { - &self.null_bitmap - } - - /// Returns the current row offset (total rows including NULLs). - pub fn current_row_offset(&self) -> u64 { - self.current_row_offset - } - - /// Returns the next HNSW key (total indexed rows excluding NULLs). - pub fn next_hnsw_key(&self) -> u64 { - self.next_hnsw_key - } - - /// Builds the blob data for the vector index. - fn build_blob(&self) -> Result> { - // Serialize null bitmap - let mut null_bitmap_bytes = Vec::new(); - self.null_bitmap - .serialize_into(&mut null_bitmap_bytes) - .context(SerializeBitmapSnafu)?; - - // Serialize the vector index - let index_size = self.engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - self.engine.save_to_buffer(&mut index_bytes).map_err(|e| { - EngineSnafu { - reason: format!("Failed to serialize index: {}", e), - } - .build() - })?; - - // Build VectorIndexMeta - let total_rows = self.current_row_offset; - let indexed_rows = self.next_hnsw_key; - let null_count = total_rows - indexed_rows; - - let meta = VectorIndexMeta { - engine: engine_type_to_proto(self.config.engine).into(), - dim: self.config.dim as u32, - metric: distance_metric_to_proto(self.config.distance_metric).into(), - connectivity: self.config.connectivity as u32, - expansion_add: self.config.expansion_add as u32, - expansion_search: self.config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - // Build blob: null_bitmap | index_data | meta | meta_size - let total_size = - null_bitmap_bytes.len() + index_bytes.len() + meta_bytes.len() + META_SIZE_LEN; - let mut blob_data = Vec::with_capacity(total_size); - blob_data.extend_from_slice(&null_bitmap_bytes); - blob_data.extend_from_slice(&index_bytes); - blob_data.extend_from_slice(&meta_bytes); - blob_data.extend_from_slice(&meta_size.to_le_bytes()); - - Ok(blob_data) - } -} - -#[async_trait] -impl IndexCreator for HnswVectorIndexCreator { - fn push_vector(&mut self, vector: &[f32]) -> Result<()> { - if self.current_row_offset >= u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment: 1u64, - } - .fail(); - } - - let key = self.next_hnsw_key; - self.engine.add(key, vector).map_err(|e| { - EngineSnafu { - reason: e.to_string(), - } - .build() - })?; - self.next_hnsw_key += 1; - self.current_row_offset += 1; - self.memory_usage_bytes = self.engine.memory_usage(); - Ok(()) - } - - fn push_null(&mut self) -> Result<()> { - if self.current_row_offset >= u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment: 1u64, - } - .fail(); - } - - self.null_bitmap.insert(self.current_row_offset as u32); - self.current_row_offset += 1; - Ok(()) - } - - fn push_nulls(&mut self, n: usize) -> Result<()> { - let increment = n as u64; - if self.current_row_offset.saturating_add(increment) > u32::MAX as u64 { - return RowCountOverflowSnafu { - current: self.current_row_offset, - increment, - } - .fail(); - } - - let start = self.current_row_offset as u32; - let end = (self.current_row_offset + increment) as u32; - self.null_bitmap.insert_range(start..end); - self.current_row_offset += increment; - Ok(()) - } - - async fn finish( - &mut self, - puffin_writer: &mut (impl PuffinWriter + Send), - blob_key: &str, - put_options: PutOptions, - ) -> Result { - if self.finished { - return FinishSnafu { - reason: "Index already finished", - } - .fail(); - } - self.finished = true; - - let blob_data = self.build_blob()?; - - puffin_writer - .put_blob( - blob_key, - Cursor::new(blob_data), - put_options, - Default::default(), - ) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - - async fn abort(&mut self) -> Result<()> { - self.finished = true; - Ok(()) - } - - fn memory_usage(&self) -> usize { - self.memory_usage_bytes + self.null_bitmap.serialized_size() - } - - fn size(&self) -> usize { - self.engine.size() - } -} - -#[cfg(test)] -mod tests { - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::VectorDistanceMetric; - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 4, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_hnsw_vector_index_creator() { - let config = test_config(); - let mut creator = HnswVectorIndexCreator::new(config).unwrap(); - - creator.reserve(10).unwrap(); - - // Add some vectors - let v1 = vec![1.0f32, 0.0, 0.0, 0.0]; - let v2 = vec![0.0f32, 1.0, 0.0, 0.0]; - - creator.push_vector(&v1).unwrap(); - creator.push_null().unwrap(); - creator.push_vector(&v2).unwrap(); - - assert_eq!(creator.size(), 2); // 2 vectors (excluding NULL) - assert_eq!(creator.current_row_offset(), 3); // 3 rows total - assert!(creator.null_bitmap().contains(1)); // Row 1 is NULL - } - - #[test] - fn test_hnsw_vector_index_creator_build_blob() { - let config = test_config(); - let mut creator = HnswVectorIndexCreator::new(config).unwrap(); - - creator.reserve(10).unwrap(); - - // Add vectors - creator.push_vector(&[1.0, 0.0, 0.0, 0.0]).unwrap(); - creator.push_null().unwrap(); - creator.push_vector(&[0.0, 1.0, 0.0, 0.0]).unwrap(); - - let blob = creator.build_blob().unwrap(); - assert!(!blob.is_empty()); - - // Verify blob can be parsed using the applier - use crate::vector::apply::{HnswVectorIndexApplier, IndexApplier}; - let applier = HnswVectorIndexApplier::from_blob(&blob).unwrap(); - assert_eq!(applier.dimensions(), 4); - assert_eq!(applier.total_rows(), 3); - assert_eq!(applier.indexed_rows(), 2); - } -} diff --git a/src/index/src/vector/engine.rs b/src/index/src/vector/engine.rs deleted file mode 100644 index b99ba89b96d..00000000000 --- a/src/index/src/vector/engine.rs +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Pluggable vector index engine implementations. - -mod usearch; - -use store_api::storage::{VectorIndexEngine, VectorIndexEngineType}; -pub use usearch::UsearchEngine; - -use crate::vector::VectorDistanceMetric; -use crate::vector::error::Result; - -/// Configuration for creating a vector index. -#[derive(Debug, Clone)] -pub struct VectorIndexConfig { - /// The vector index engine type. - pub engine: VectorIndexEngineType, - /// The dimension of vectors in this column. - pub dim: usize, - /// The distance metric (e.g., L2sq, Cosine, InnerProduct). - pub distance_metric: VectorDistanceMetric, - /// HNSW connectivity parameter (M in the paper). - /// Higher values give better recall but use more memory. - pub connectivity: usize, - /// Expansion factor during index construction (ef_construction). - pub expansion_add: usize, - /// Expansion factor during search (ef_search). - pub expansion_search: usize, -} - -/// Creates a new vector index engine based on the engine type. -pub fn create_engine( - engine_type: VectorIndexEngineType, - config: &VectorIndexConfig, -) -> Result> { - match engine_type { - VectorIndexEngineType::Usearch => Ok(Box::new(UsearchEngine::create(config)?)), - } -} - -/// Loads a vector index engine from serialized data. -#[allow(unused)] -pub fn load_engine( - engine_type: VectorIndexEngineType, - config: &VectorIndexConfig, - data: &[u8], -) -> Result> { - match engine_type { - VectorIndexEngineType::Usearch => Ok(Box::new(UsearchEngine::load(config, data)?)), - } -} diff --git a/src/index/src/vector/engine/usearch.rs b/src/index/src/vector/engine/usearch.rs deleted file mode 100644 index 6c4352bb032..00000000000 --- a/src/index/src/vector/engine/usearch.rs +++ /dev/null @@ -1,213 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! USearch HNSW implementation of VectorIndexEngine. - -use common_error::ext::BoxedError; -use store_api::storage::{VectorIndexEngine, VectorSearchMatches}; -use usearch::{Index, IndexOptions, ScalarKind}; - -use crate::vector::distance_metric_to_usearch; -use crate::vector::engine::VectorIndexConfig; -use crate::vector::error::{EngineSnafu, Result}; - -type EngineResult = std::result::Result; - -/// USearch-based vector index engine using HNSW algorithm. -pub struct UsearchEngine { - index: Index, -} - -impl UsearchEngine { - /// Creates a new USearch engine with the given configuration. - pub fn create(config: &VectorIndexConfig) -> Result { - let options = IndexOptions { - dimensions: config.dim, - metric: distance_metric_to_usearch(config.distance_metric), - quantization: ScalarKind::F32, - connectivity: config.connectivity, - expansion_add: config.expansion_add, - expansion_search: config.expansion_search, - multi: false, - }; - - let index = Index::new(&options).map_err(|e| { - EngineSnafu { - reason: format!("Failed to create USearch index: {}", e), - } - .build() - })?; - - Ok(Self { index }) - } - - /// Loads a USearch engine from serialized data. - #[allow(unused)] - pub fn load(config: &VectorIndexConfig, data: &[u8]) -> Result { - let options = IndexOptions { - dimensions: config.dim, - metric: distance_metric_to_usearch(config.distance_metric), - quantization: ScalarKind::F32, - // These will be loaded from serialized data - connectivity: 0, - expansion_add: 0, - expansion_search: 0, - multi: false, - }; - - let index = Index::new(&options).map_err(|e| { - EngineSnafu { - reason: format!("Failed to create USearch index for loading: {}", e), - } - .build() - })?; - - index.load_from_buffer(data).map_err(|e| { - EngineSnafu { - reason: format!("Failed to load USearch index from buffer: {}", e), - } - .build() - })?; - - Ok(Self { index }) - } - - /// Helper to create a BoxedError from a reason string. - fn boxed_engine_error(reason: String) -> BoxedError { - BoxedError::new(EngineSnafu { reason }.build()) - } -} - -impl VectorIndexEngine for UsearchEngine { - fn add(&mut self, key: u64, vector: &[f32]) -> EngineResult<()> { - // Reserve capacity if needed - if self.index.size() >= self.index.capacity() { - let new_capacity = std::cmp::max(1, self.index.capacity() * 2); - self.index.reserve(new_capacity).map_err(|e| { - Self::boxed_engine_error(format!("Failed to reserve capacity: {}", e)) - })?; - } - - self.index - .add(key, vector) - .map_err(|e| Self::boxed_engine_error(format!("Failed to add vector: {}", e))) - } - - fn search(&self, query: &[f32], k: usize) -> EngineResult { - let matches = self - .index - .search(query, k) - .map_err(|e| Self::boxed_engine_error(format!("Failed to search: {}", e)))?; - - Ok(VectorSearchMatches { - keys: matches.keys, - distances: matches.distances, - }) - } - - fn serialized_length(&self) -> usize { - self.index.serialized_length() - } - - fn save_to_buffer(&self, buffer: &mut [u8]) -> EngineResult<()> { - self.index - .save_to_buffer(buffer) - .map_err(|e| Self::boxed_engine_error(format!("Failed to save to buffer: {}", e))) - } - - fn reserve(&mut self, capacity: usize) -> EngineResult<()> { - self.index - .reserve(capacity) - .map_err(|e| Self::boxed_engine_error(format!("Failed to reserve: {}", e))) - } - - fn size(&self) -> usize { - self.index.size() - } - - fn capacity(&self) -> usize { - self.index.capacity() - } - - fn memory_usage(&self) -> usize { - self.index.memory_usage() - } -} - -#[cfg(test)] -mod tests { - use store_api::storage::VectorIndexEngineType; - - use super::*; - use crate::vector::VectorDistanceMetric; - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 4, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[test] - fn test_usearch_engine_create() { - let config = test_config(); - let engine = UsearchEngine::create(&config).unwrap(); - assert_eq!(engine.size(), 0); - } - - #[test] - fn test_usearch_engine_add_and_search() { - let config = test_config(); - let mut engine = UsearchEngine::create(&config).unwrap(); - - // Add some vectors - engine.add(0, &[1.0, 0.0, 0.0, 0.0]).unwrap(); - engine.add(1, &[0.0, 1.0, 0.0, 0.0]).unwrap(); - engine.add(2, &[0.0, 0.0, 1.0, 0.0]).unwrap(); - - assert_eq!(engine.size(), 3); - - // Search - let matches = engine.search(&[1.0, 0.0, 0.0, 0.0], 2).unwrap(); - assert_eq!(matches.keys.len(), 2); - // First result should be the exact match (key 0) - assert_eq!(matches.keys[0], 0); - } - - #[test] - fn test_usearch_engine_serialization() { - let config = test_config(); - let mut engine = UsearchEngine::create(&config).unwrap(); - - engine.add(0, &[1.0, 0.0, 0.0, 0.0]).unwrap(); - engine.add(1, &[0.0, 1.0, 0.0, 0.0]).unwrap(); - - // Serialize - let len = engine.serialized_length(); - let mut buffer = vec![0u8; len]; - engine.save_to_buffer(&mut buffer).unwrap(); - - // Load - let loaded = UsearchEngine::load(&config, &buffer).unwrap(); - assert_eq!(loaded.size(), 2); - - // Verify search works on loaded index - let matches = loaded.search(&[1.0, 0.0, 0.0, 0.0], 1).unwrap(); - assert_eq!(matches.keys[0], 0); - } -} diff --git a/src/index/src/vector/error.rs b/src/index/src/vector/error.rs deleted file mode 100644 index 10839db2068..00000000000 --- a/src/index/src/vector/error.rs +++ /dev/null @@ -1,152 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::any::Any; - -use common_error::ext::{BoxedError, ErrorExt}; -use common_error::status_code::StatusCode; -use common_macro::stack_trace_debug; -use snafu::{Location, Snafu}; - -#[derive(Snafu)] -#[snafu(visibility(pub))] -#[stack_trace_debug] -pub enum Error { - #[snafu(display("Failed to decode protobuf"))] - DecodeProto { - #[snafu(source)] - error: prost::DecodeError, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to deserialize bitmap"))] - DeserializeBitmap { - #[snafu(source)] - error: std::io::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to serialize bitmap"))] - SerializeBitmap { - #[snafu(source)] - error: std::io::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index blob too small"))] - BlobTooSmall { - min_size: usize, - actual_size: usize, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index blob truncated"))] - BlobTruncated { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Invalid vector index blob"))] - InvalidBlob { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Unknown vector index engine type: {engine_type}"))] - UnknownEngineType { - engine_type: i32, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Unknown vector distance metric: {metric}"))] - UnknownDistanceMetric { - metric: i32, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Vector index engine error: {reason}"))] - Engine { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to map HNSW key to row offset: {reason}"))] - KeyMapping { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("External error"))] - External { - #[snafu(source)] - error: BoxedError, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to finish vector index: {reason}"))] - Finish { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Row count overflow: {current} + {increment} exceeds u32::MAX"))] - RowCountOverflow { - current: u64, - increment: u64, - #[snafu(implicit)] - location: Location, - }, -} - -impl ErrorExt for Error { - fn status_code(&self) -> StatusCode { - use Error::*; - - match self { - DecodeProto { .. } - | DeserializeBitmap { .. } - | SerializeBitmap { .. } - | BlobTooSmall { .. } - | BlobTruncated { .. } - | InvalidBlob { .. } - | Engine { .. } - | KeyMapping { .. } - | Finish { .. } => StatusCode::Unexpected, - - UnknownEngineType { .. } | UnknownDistanceMetric { .. } | RowCountOverflow { .. } => { - StatusCode::InvalidArguments - } - - External { error, .. } => error.status_code(), - } - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -pub type Result = std::result::Result; diff --git a/src/index/src/vector/format.rs b/src/index/src/vector/format.rs deleted file mode 100644 index 28a7f9744e9..00000000000 --- a/src/index/src/vector/format.rs +++ /dev/null @@ -1,200 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index blob format helpers. -//! -//! Blob layout (footer-last): -//! ```text -//! +----------------------+ -//! | Null Bitmap | <- RoaringBitmap serialized, size = meta.null_bitmap_size -//! +----------------------+ -//! | Index Data | <- engine-specific serialized, size = meta.index_size -//! +----------------------+ -//! | Meta (protobuf) | <- VectorIndexMeta encoded -//! +----------------------+ -//! | Meta Size (4B) | <- u32 LE -//! +----------------------+ -//! ``` - -use greptime_proto::v1::index::{ - VectorDistanceMetric as ProtoVectorDistanceMetric, VectorIndexEngine as ProtoVectorIndexEngine, -}; -use store_api::storage::{VectorDistanceMetric, VectorIndexEngineType}; - -/// Size of the meta length field in bytes (u32 LE). -pub const META_SIZE_LEN: usize = 4; - -/// Converts Rust `VectorIndexEngineType` to proto `VectorIndexEngine`. -pub fn engine_type_to_proto(engine: VectorIndexEngineType) -> ProtoVectorIndexEngine { - match engine { - VectorIndexEngineType::Usearch => ProtoVectorIndexEngine::Usearch, - } -} - -/// Converts proto `VectorIndexEngine` to Rust `VectorIndexEngineType`. -pub fn engine_type_from_proto(engine: ProtoVectorIndexEngine) -> Option { - match engine { - ProtoVectorIndexEngine::Usearch => Some(VectorIndexEngineType::Usearch), - } -} - -/// Converts Rust `VectorDistanceMetric` to proto `VectorDistanceMetric`. -pub fn distance_metric_to_proto(metric: VectorDistanceMetric) -> ProtoVectorDistanceMetric { - match metric { - VectorDistanceMetric::L2sq => ProtoVectorDistanceMetric::L2sq, - VectorDistanceMetric::Cosine => ProtoVectorDistanceMetric::Cosine, - VectorDistanceMetric::InnerProduct => ProtoVectorDistanceMetric::InnerProduct, - } -} - -/// Converts proto `VectorDistanceMetric` to Rust `VectorDistanceMetric`. -pub fn distance_metric_from_proto( - metric: ProtoVectorDistanceMetric, -) -> Option { - match metric { - ProtoVectorDistanceMetric::L2sq => Some(VectorDistanceMetric::L2sq), - ProtoVectorDistanceMetric::Cosine => Some(VectorDistanceMetric::Cosine), - ProtoVectorDistanceMetric::InnerProduct => Some(VectorDistanceMetric::InnerProduct), - } -} - -#[cfg(test)] -mod tests { - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use prost::Message; - use roaring::RoaringBitmap; - - use super::*; - - #[test] - fn test_engine_type_roundtrip() { - let engine = VectorIndexEngineType::Usearch; - let proto = engine_type_to_proto(engine); - let back = engine_type_from_proto(proto).unwrap(); - assert_eq!(engine, back); - } - - #[test] - fn test_distance_metric_roundtrip() { - for metric in [ - VectorDistanceMetric::L2sq, - VectorDistanceMetric::Cosine, - VectorDistanceMetric::InnerProduct, - ] { - let proto = distance_metric_to_proto(metric); - let back = distance_metric_from_proto(proto).unwrap(); - assert_eq!(metric, back); - } - } - - #[test] - fn test_vector_index_meta_encode_decode() { - let meta = VectorIndexMeta { - engine: engine_type_to_proto(VectorIndexEngineType::Usearch).into(), - dim: 128, - metric: distance_metric_to_proto(VectorDistanceMetric::Cosine).into(), - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - null_bitmap_size: 10, - index_size: 1024, - stats: Some(VectorIndexStats { - total_row_count: 100, - indexed_row_count: 95, - null_count: 5, - }), - }; - - let encoded = meta.encode_to_vec(); - let decoded = VectorIndexMeta::decode(encoded.as_slice()).unwrap(); - - assert_eq!(decoded.engine, meta.engine); - assert_eq!(decoded.dim, meta.dim); - assert_eq!(decoded.metric, meta.metric); - assert_eq!(decoded.connectivity, meta.connectivity); - assert_eq!(decoded.expansion_add, meta.expansion_add); - assert_eq!(decoded.expansion_search, meta.expansion_search); - assert_eq!(decoded.null_bitmap_size, meta.null_bitmap_size); - assert_eq!(decoded.index_size, meta.index_size); - - let stats = decoded.stats.unwrap(); - assert_eq!(stats.total_row_count, 100); - assert_eq!(stats.indexed_row_count, 95); - assert_eq!(stats.null_count, 5); - } - - #[test] - fn test_blob_format_roundtrip() { - // Create test data - let null_bitmap = RoaringBitmap::from_iter([1u32, 3, 5]); - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let index_data = b"fake_index_data_here"; - - let meta = VectorIndexMeta { - engine: engine_type_to_proto(VectorIndexEngineType::Usearch).into(), - dim: 4, - metric: distance_metric_to_proto(VectorDistanceMetric::L2sq).into(), - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_data.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: 10, - indexed_row_count: 7, - null_count: 3, - }), - }; - - // Build blob: null_bitmap | index_data | meta | meta_size - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(index_data); - let meta_bytes = meta.encode_to_vec(); - blob.extend_from_slice(&meta_bytes); - let meta_size = meta_bytes.len() as u32; - blob.extend_from_slice(&meta_size.to_le_bytes()); - - // Parse blob - let blob_len = blob.len(); - let meta_size_read = - u32::from_le_bytes(blob[blob_len - META_SIZE_LEN..].try_into().unwrap()) as usize; - assert_eq!(meta_size_read, meta_bytes.len()); - - let meta_start = blob_len - META_SIZE_LEN - meta_size_read; - let decoded_meta = - VectorIndexMeta::decode(&blob[meta_start..blob_len - META_SIZE_LEN]).unwrap(); - - assert_eq!( - decoded_meta.null_bitmap_size, - null_bitmap_bytes.len() as u32 - ); - assert_eq!(decoded_meta.index_size, index_data.len() as u64); - - // Extract null bitmap - let null_bitmap_end = decoded_meta.null_bitmap_size as usize; - let restored_bitmap = RoaringBitmap::deserialize_from(&blob[..null_bitmap_end]).unwrap(); - assert_eq!(restored_bitmap.len(), 3); - assert!(restored_bitmap.contains(1)); - assert!(restored_bitmap.contains(3)); - assert!(restored_bitmap.contains(5)); - - // Extract index data - let index_start = null_bitmap_end; - let index_end = index_start + decoded_meta.index_size as usize; - assert_eq!(&blob[index_start..index_end], index_data); - } -} diff --git a/src/index/src/vector/util.rs b/src/index/src/vector/util.rs deleted file mode 100644 index 7610a07e7cc..00000000000 --- a/src/index/src/vector/util.rs +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Utility functions for vector index operations. - -use std::borrow::Cow; - -/// Converts a byte slice (little-endian format) to f32 slice, handling unaligned data gracefully. -/// Returns `Cow::Borrowed` for aligned data on little-endian systems (zero-copy) -/// or `Cow::Owned` for unaligned data or big-endian systems. -/// -/// # Panics -/// -/// Panics if the byte slice length is not a multiple of 4. -pub fn bytes_to_f32_slice(bytes: &[u8]) -> Cow<'_, [f32]> { - assert!( - bytes.len().is_multiple_of(4), - "Vector bytes length {} is not a multiple of 4", - bytes.len() - ); - - if bytes.is_empty() { - return Cow::Borrowed(&[]); - } - - let ptr = bytes.as_ptr(); - // Fast path: zero-copy only when data is aligned AND we're on little-endian system - // (since vector data is stored in little-endian format) - #[cfg(target_endian = "little")] - if (ptr as usize).is_multiple_of(std::mem::align_of::()) { - // Safety: We've verified alignment and length requirements, - // and on little-endian systems the byte representation matches f32 layout - return Cow::Borrowed(unsafe { - std::slice::from_raw_parts(ptr as *const f32, bytes.len() / 4) - }); - } - - // Slow path: data is not aligned or we're on big-endian system - let floats: Vec = bytes - .chunks_exact(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(); - Cow::Owned(floats) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_bytes_to_f32_slice() { - let floats = [1.0f32, 2.0, 3.0, 4.0]; - let bytes: Vec = floats.iter().flat_map(|f| f.to_le_bytes()).collect(); - - let result = bytes_to_f32_slice(&bytes); - assert_eq!(result.len(), 4); - assert_eq!(result[0], 1.0); - assert_eq!(result[1], 2.0); - assert_eq!(result[2], 3.0); - assert_eq!(result[3], 4.0); - } - - #[test] - fn test_bytes_to_f32_slice_unaligned() { - // Create a buffer with an extra byte at the start to force misalignment - let floats = [1.0f32, 2.0, 3.0, 4.0]; - let mut bytes: Vec = vec![0u8]; // padding byte - bytes.extend(floats.iter().flat_map(|f| f.to_le_bytes())); - - // Take a slice starting at offset 1 (unaligned) - let unaligned_bytes = &bytes[1..]; - - // Verify it's actually unaligned - let ptr = unaligned_bytes.as_ptr(); - let is_aligned = (ptr as usize).is_multiple_of(std::mem::align_of::()); - - // The function should work regardless of alignment - let result = bytes_to_f32_slice(unaligned_bytes); - assert_eq!(result.len(), 4); - assert_eq!(result[0], 1.0); - assert_eq!(result[1], 2.0); - assert_eq!(result[2], 3.0); - assert_eq!(result[3], 4.0); - - // If it was unaligned, it should return an owned Vec (Cow::Owned) - if !is_aligned { - assert!(matches!(result, Cow::Owned(_))); - } - } - - #[test] - fn test_bytes_to_f32_slice_empty() { - let bytes: &[u8] = &[]; - let result = bytes_to_f32_slice(bytes); - assert!(result.is_empty()); - } -} diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 82a2884549a..9ab2dee10df 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -10,7 +10,6 @@ test = ["common-test-util", "rstest", "rstest_reuse", "rskafka"] testing = ["test"] test-shared-fs-region-migration = [] enterprise = [] -vector_index = ["dep:roaring", "index/vector_index"] [lints] workspace = true @@ -72,7 +71,6 @@ puffin.workspace = true rand.workspace = true rayon = "1.10" regex.workspace = true -roaring = { version = "0.10", optional = true } rskafka = { workspace = true, optional = true } rstest = { workspace = true, optional = true } rstest_reuse = { workspace = true, optional = true } diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 685cc0d2fc7..06fdfd79240 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -391,8 +391,6 @@ impl AccessLayer { inverted_index_config: request.inverted_index_config, fulltext_index_config: request.fulltext_index_config, bloom_filter_index_config: request.bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config: request.vector_index_config, }; // We disable write cache on file system but we still use atomic write. // TODO(yingwen): If we support other non-fs stores without the write cache, then @@ -595,8 +593,6 @@ pub struct SstWriteRequest { pub inverted_index_config: InvertedIndexConfig, pub fulltext_index_config: FulltextIndexConfig, pub bloom_filter_index_config: BloomFilterConfig, - #[cfg(feature = "vector_index")] - pub vector_index_config: crate::config::VectorIndexConfig, } /// Cleaner to remove temp files on the atomic write dir. diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index ef150dfb070..70dee2f7dd8 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -58,8 +58,6 @@ pub use write_cache::{WriteCacheUploadStoreWrapper, WriteCacheUploadStoreWrapper use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; use crate::cache::index::inverted_index::{InvertedIndexCache, InvertedIndexCacheRef}; -#[cfg(feature = "vector_index")] -use crate::cache::index::vector_index::{VectorIndexCache, VectorIndexCacheRef}; use crate::cache::write_cache::WriteCacheRef; use crate::error::{ CompressObjectSnafu, DecompressObjectSnafu, InvalidMetadataSnafu, InvalidParquetSnafu, @@ -977,16 +975,6 @@ impl CacheStrategy { } } - /// Calls [CacheManager::vector_index_cache()]. - /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. - #[cfg(feature = "vector_index")] - pub fn vector_index_cache(&self) -> Option<&VectorIndexCacheRef> { - match self { - CacheStrategy::EnableAll(cache_manager) => cache_manager.vector_index_cache(), - CacheStrategy::Compaction(_) | CacheStrategy::Disabled => None, - } - } - /// Calls [CacheManager::puffin_metadata_cache()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. pub fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> { @@ -1045,9 +1033,6 @@ pub struct CacheManager { inverted_index_cache: Option, /// Cache for bloom filter index. bloom_filter_index_cache: Option, - /// Cache for vector index. - #[cfg(feature = "vector_index")] - vector_index_cache: Option, /// Puffin metadata cache. puffin_metadata_cache: Option, /// Cache for time series selectors. @@ -1366,11 +1351,6 @@ impl CacheManager { cache.invalidate_file(file_id.file_id()); } - #[cfg(feature = "vector_index")] - if let Some(cache) = &self.vector_index_cache { - cache.invalidate_file(file_id.file_id()); - } - if let Some(cache) = &self.puffin_metadata_cache { cache.remove(&file_id.to_string()); } @@ -1461,11 +1441,6 @@ impl CacheManager { self.bloom_filter_index_cache.as_ref() } - #[cfg(feature = "vector_index")] - pub(crate) fn vector_index_cache(&self) -> Option<&VectorIndexCacheRef> { - self.vector_index_cache.as_ref() - } - pub(crate) fn puffin_metadata_cache(&self) -> Option<&PuffinMetadataCacheRef> { self.puffin_metadata_cache.as_ref() } @@ -1655,9 +1630,6 @@ impl CacheManagerBuilder { self.index_content_size, self.index_content_page_size, ); - #[cfg(feature = "vector_index")] - let vector_index_cache = (self.index_content_size != 0) - .then(|| Arc::new(VectorIndexCache::new(self.index_content_size))); let index_result_cache = (self.index_result_cache_size != 0) .then(|| IndexResultCache::new(self.index_result_cache_size)); let prefilter_result_cache = (self.prefilter_result_cache_size != 0) @@ -1702,8 +1674,6 @@ impl CacheManagerBuilder { write_cache: self.write_cache, inverted_index_cache: Some(Arc::new(inverted_index_cache)), bloom_filter_index_cache: Some(Arc::new(bloom_filter_index_cache)), - #[cfg(feature = "vector_index")] - vector_index_cache, puffin_metadata_cache: Some(Arc::new(puffin_metadata_cache)), selector_result_cache, range_result_cache, diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index a4049787149..0cab159a7db 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -15,8 +15,6 @@ pub mod bloom_filter_index; pub mod inverted_index; pub mod result_cache; -#[cfg(feature = "vector_index")] -pub mod vector_index; use std::future::Future; use std::hash::Hash; diff --git a/src/mito2/src/cache/index/vector_index.rs b/src/mito2/src/cache/index/vector_index.rs deleted file mode 100644 index 222c8bb0690..00000000000 --- a/src/mito2/src/cache/index/vector_index.rs +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; - -use index::vector::apply::IndexApplier; -use moka::notification::RemovalCause; -use moka::sync::Cache; -use store_api::storage::{ColumnId, FileId, IndexVersion}; - -use crate::metrics::{CACHE_BYTES, CACHE_EVICTION}; - -const VECTOR_INDEX_CACHE_TYPE: &str = "vector_index"; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct VectorIndexCacheKey { - file_id: FileId, - index_version: IndexVersion, - column_id: ColumnId, -} - -impl VectorIndexCacheKey { - pub fn new(file_id: FileId, index_version: IndexVersion, column_id: ColumnId) -> Self { - Self { - file_id, - index_version, - column_id, - } - } -} - -/// Cached vector index wrapping an `IndexApplier` for efficient searches. -pub struct CachedVectorIndex { - /// The underlying vector index applier. - applier: Box, - /// Size in bytes for cache weighing. - size_bytes: usize, -} - -impl CachedVectorIndex { - /// Creates a new `CachedVectorIndex` from an `IndexApplier`. - pub fn new(applier: Box) -> Self { - let size_bytes = applier.memory_usage(); - Self { - applier, - size_bytes, - } - } - - /// Returns the underlying applier for searching. - pub fn applier(&self) -> &dyn IndexApplier { - self.applier.as_ref() - } - - /// Returns the cache size in bytes. - pub fn size_bytes(&self) -> usize { - self.size_bytes - } -} - -impl std::fmt::Debug for CachedVectorIndex { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let null_count = self.applier.total_rows() - self.applier.indexed_rows(); - f.debug_struct("CachedVectorIndex") - .field("size_bytes", &self.size_bytes) - .field("dimensions", &self.applier.dimensions()) - .field("metric", &self.applier.metric()) - .field("total_rows", &self.applier.total_rows()) - .field("indexed_rows", &self.applier.indexed_rows()) - .field("null_count", &null_count) - .finish() - } -} - -pub struct VectorIndexCache { - inner: Cache>, -} - -pub type VectorIndexCacheRef = Arc; - -impl VectorIndexCache { - pub fn new(capacity: u64) -> Self { - fn to_str(cause: RemovalCause) -> &'static str { - match cause { - RemovalCause::Expired => "expired", - RemovalCause::Explicit => "explicit", - RemovalCause::Replaced => "replaced", - RemovalCause::Size => "size", - } - } - - let inner = Cache::builder() - .max_capacity(capacity) - .weigher(|_k, v: &Arc| v.size_bytes() as u32) - .eviction_listener(|_k, v, cause| { - CACHE_BYTES - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE]) - .sub(v.size_bytes() as i64); - CACHE_EVICTION - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE, to_str(cause)]) - .inc(); - }) - .build(); - Self { inner } - } - - pub fn get(&self, key: &VectorIndexCacheKey) -> Option> { - self.inner.get(key) - } - - pub fn insert(&self, key: VectorIndexCacheKey, value: Arc) { - CACHE_BYTES - .with_label_values(&[VECTOR_INDEX_CACHE_TYPE]) - .add(value.size_bytes() as i64); - self.inner.insert(key, value); - } - - pub fn invalidate_file(&self, file_id: FileId) { - let _ = self - .inner - .invalidate_entries_if(move |k, _| k.file_id == file_id); - } -} diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 87ab4f0081a..e625556e087 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -266,8 +266,6 @@ impl WriteCache { inverted_index_config: write_request.inverted_index_config, fulltext_index_config: write_request.fulltext_index_config, bloom_filter_index_config: write_request.bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config: write_request.vector_index_config, }; let cleaner = TempFileCleaner::new(region_id, store.clone()); @@ -713,8 +711,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let upload_request = SstUploadRequest { @@ -833,8 +829,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let write_opts = WriteOptions { row_group_size: 512, @@ -928,8 +922,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let write_opts = WriteOptions { row_group_size: 512, diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index c7d55fb8a28..c6a9999ee41 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -479,8 +479,6 @@ impl SstMerger for DefaultSstMerger { let inverted_index_config = compaction_region.engine_config.inverted_index.clone(); let fulltext_index_config = compaction_region.engine_config.fulltext_index.clone(); let bloom_filter_index_config = compaction_region.engine_config.bloom_filter_index.clone(); - #[cfg(feature = "vector_index")] - let vector_index_config = compaction_region.engine_config.vector_index.clone(); let input_file_names = output .inputs @@ -526,8 +524,6 @@ impl SstMerger for DefaultSstMerger { inverted_index_config, fulltext_index_config, bloom_filter_index_config, - #[cfg(feature = "vector_index")] - vector_index_config, }, &write_opts, &mut metrics, diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 035feb75a15..78b79aec81a 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -192,9 +192,6 @@ pub struct MitoConfig { pub fulltext_index: FulltextIndexConfig, /// Bloom filter index configs. pub bloom_filter_index: BloomFilterConfig, - /// Vector index configs (HNSW). - #[cfg(feature = "vector_index")] - pub vector_index: VectorIndexConfig, /// Minimum time interval between two compactions. /// To align with the old behavior, the default value is 0 (no restrictions). @@ -262,8 +259,6 @@ impl Default for MitoConfig { inverted_index: InvertedIndexConfig::default(), fulltext_index: FulltextIndexConfig::default(), bloom_filter_index: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index: VectorIndexConfig::default(), min_compaction_interval: Duration::from_secs(0), schedule_compaction_after_edit: true, default_flat_format: true, @@ -805,51 +800,6 @@ impl BloomFilterConfig { } } -/// Configuration options for the vector index (HNSW). -#[cfg(feature = "vector_index")] -#[serde_as] -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] -#[serde(default)] -pub struct VectorIndexConfig { - /// Whether to create the index on flush: automatically or never. - pub create_on_flush: Mode, - /// Whether to create the index on compaction: automatically or never. - pub create_on_compaction: Mode, - /// Whether to apply the index on query: automatically or never. - pub apply_on_query: Mode, - /// Memory threshold for creating the index. - pub mem_threshold_on_create: MemoryThreshold, -} - -#[cfg(feature = "vector_index")] -impl Default for VectorIndexConfig { - fn default() -> Self { - Self { - create_on_flush: Mode::Auto, - create_on_compaction: Mode::Auto, - apply_on_query: Mode::Auto, - mem_threshold_on_create: MemoryThreshold::Auto, - } - } -} - -#[cfg(feature = "vector_index")] -impl VectorIndexConfig { - pub fn mem_threshold_on_create(&self) -> Option { - match self.mem_threshold_on_create { - MemoryThreshold::Auto => { - if let Some(sys_memory) = get_total_memory_readable() { - Some((sys_memory / INDEX_CREATE_MEM_THRESHOLD_FACTOR).as_bytes() as usize) - } else { - Some(ReadableSize::mb(64).as_bytes() as usize) - } - } - MemoryThreshold::Unlimited => None, - MemoryThreshold::Size(size) => Some(size.as_bytes() as usize), - } - } -} - /// Divide cpu num by a non-zero `divisor` and returns at least 1. fn divide_num_cpus(divisor: usize) -> usize { debug_assert!(divisor > 0); diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 1592bfebd60..7d43548dd2b 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -773,14 +773,6 @@ pub enum Error { location: Location, }, - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to apply vector index: {}", reason))] - ApplyVectorIndex { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Failed to push index value"))] PushIndexValue { source: index::inverted_index::error::Error, @@ -1197,22 +1189,6 @@ pub enum Error { location: Location, }, - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to build vector index: {}", reason))] - VectorIndexBuild { - reason: String, - #[snafu(implicit)] - location: Location, - }, - - #[cfg(feature = "vector_index")] - #[snafu(display("Failed to finish vector index: {}", reason))] - VectorIndexFinish { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Manual compaction is override by following operations."))] ManualCompactionOverride {}, @@ -1571,8 +1547,6 @@ impl ErrorExt for Error { | PushIndexValue { source, .. } | ApplyInvertedIndex { source, .. } | IndexFinish { source, .. } => source.status_code(), - #[cfg(feature = "vector_index")] - ApplyVectorIndex { .. } => StatusCode::Internal, PuffinReadBlob { source, .. } | PuffinAddBlob { source, .. } | PuffinInitStager { source, .. } @@ -1615,9 +1589,6 @@ impl ErrorExt for Error { source.status_code() } - #[cfg(feature = "vector_index")] - VectorIndexBuild { .. } | VectorIndexFinish { .. } => StatusCode::Internal, - ManualCompactionOverride {} | CompactionCancelled {} | FlushCancelled {} => { StatusCode::Cancelled } diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 8a96276e5f7..66182693f78 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -849,8 +849,6 @@ impl RegionFlushTask { inverted_index_config: self.engine_config.inverted_index.clone(), fulltext_index_config: self.engine_config.fulltext_index.clone(), bloom_filter_index_config: self.engine_config.bloom_filter_index.clone(), - #[cfg(feature = "vector_index")] - vector_index_config: self.engine_config.vector_index.clone(), } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 1d0991a6bfb..8edca2a0a9c 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -85,16 +85,11 @@ use crate::sst::index::fulltext_index::applier::FulltextIndexApplierRef; use crate::sst::index::fulltext_index::applier::builder::FulltextIndexApplierBuilder; use crate::sst::index::inverted_index::applier::InvertedIndexApplierRef; use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; -#[cfg(feature = "vector_index")] -use crate::sst::index::vector_index::applier::{VectorIndexApplier, VectorIndexApplierRef}; use crate::sst::parquet::Json2RewriteTargets; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::reader::ReaderMetrics; use crate::sst::primary_key::PrimaryKeyRangeMapper; -#[cfg(feature = "vector_index")] -const VECTOR_INDEX_OVERFETCH_MULTIPLIER: usize = 2; - /// A scanner scans a region and returns a [SendableRecordBatchStream]. pub(crate) enum Scanner { /// Sequential scan. @@ -572,16 +567,6 @@ impl ScanRegion { self.build_fulltext_index_applier(&non_field_filters), self.build_fulltext_index_applier(&field_filters), ]; - #[cfg(feature = "vector_index")] - let vector_index_applier = self.build_vector_index_applier(); - #[cfg(feature = "vector_index")] - let vector_index_k = self.request.vector_search.as_ref().map(|search| { - if self.request.filters.is_empty() { - search.k - } else { - search.k.saturating_mul(VECTOR_INDEX_OVERFETCH_MULTIPLIER) - } - }); let input = ScanInput::builder(self.access_layer, mapper) .with_series_index(self.series_index) @@ -614,11 +599,6 @@ impl ScanRegion { ) .with_sequence_range(sequence_range) .with_query_stat_counters(self.query_stat_counters); - #[cfg(feature = "vector_index")] - let input = input - .with_vector_index_applier(vector_index_applier) - .with_vector_index_k(vector_index_k); - #[cfg(feature = "enterprise")] let input = if !self.request.skip_sst_files && let Some(provider) = self.extension_range_provider @@ -906,31 +886,6 @@ impl ScanRegion { .flatten() .map(Arc::new) } - - /// Build the vector index applier from vector search request. - #[cfg(feature = "vector_index")] - fn build_vector_index_applier(&self) -> Option { - let vector_search = self.request.vector_search.as_ref()?; - - let file_cache = self.cache_strategy.write_cache().map(|w| w.file_cache()); - let puffin_metadata_cache = self.cache_strategy.puffin_metadata_cache().cloned(); - let vector_index_cache = self.cache_strategy.vector_index_cache().cloned(); - - let applier = VectorIndexApplier::new( - self.access_layer.table_dir().to_string(), - self.access_layer.path_type(), - self.access_layer.object_store().clone(), - self.access_layer.puffin_manager_factory().clone(), - vector_search.column_id, - vector_search.query_vector.clone(), - vector_search.metric, - ) - .with_file_cache(file_cache) - .with_puffin_metadata_cache(puffin_metadata_cache) - .with_vector_index_cache(vector_index_cache); - - Some(Arc::new(applier)) - } } /// Returns true if the time range of a SST `file` matches the `predicate`. @@ -995,12 +950,6 @@ pub struct ScanInput { inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], fulltext_index_appliers: [Option; 2], - /// Vector index applier for KNN search. - #[cfg(feature = "vector_index")] - pub(crate) vector_index_applier: Option, - /// Over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - pub(crate) vector_index_k: Option, /// Start time of the query. pub(crate) query_start: Option, /// The region is using append mode. @@ -1079,10 +1028,6 @@ impl ScanInput { inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], - #[cfg(feature = "vector_index")] - vector_index_applier: None, - #[cfg(feature = "vector_index")] - vector_index_k: None, query_start: None, append_mode: false, filter_deleted: true, @@ -1365,25 +1310,6 @@ impl ScanInputBuilder { self } - /// Sets vector index applier for KNN search. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn with_vector_index_applier( - mut self, - applier: Option, - ) -> Self { - self.input.vector_index_applier = applier; - self - } - - /// Sets over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn with_vector_index_k(mut self, k: Option) -> Self { - self.input.vector_index_k = k; - self - } - /// Sets start time of the query. #[must_use] pub(crate) fn with_start_time(mut self, now: Option) -> Self { @@ -1632,13 +1558,6 @@ impl ScanInput { } else { reader }; - #[cfg(feature = "vector_index")] - let reader = { - let mut reader = reader; - reader = - reader.vector_index_applier(self.vector_index_applier.clone(), self.vector_index_k); - reader - }; let res = reader .expected_metadata(Some(self.mapper.metadata().clone())) .compaction(self.compaction) @@ -2197,10 +2116,6 @@ impl StreamContext { write!(f, ", \"dyn_filters\": {:?}", dyn_filters)?; } } - #[cfg(feature = "vector_index")] - if let Some(vector_index_k) = self.input.vector_index_k { - write!(f, ", \"vector_index_k\": {}", vector_index_k)?; - } if !self.input.files.is_empty() { write!(f, ", \"files\": ")?; f.debug_list() diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index 2d93cbbbe03..b58757930eb 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -158,8 +158,6 @@ pub(crate) struct ScanMetricsSet { rg_minmax_filtered: usize, /// Number of row groups filtered by bloom filter index. rg_bloom_filtered: usize, - /// Number of row groups filtered by vector index. - rg_vector_filtered: usize, /// Number of rows in row group before filtering. rows_before_filter: usize, /// Number of rows in row group filtered by fulltext index. @@ -168,10 +166,6 @@ pub(crate) struct ScanMetricsSet { rows_inverted_filtered: usize, /// Number of rows in row group filtered by bloom filter index. rows_bloom_filtered: usize, - /// Number of rows filtered by vector index. - rows_vector_filtered: usize, - /// Number of rows selected by vector index. - rows_vector_selected: usize, /// Number of rows filtered by precise filter. rows_precise_filtered: usize, /// Number of index result cache hits for fulltext index. @@ -310,13 +304,10 @@ impl fmt::Debug for ScanMetricsSet { rg_inverted_filtered, rg_minmax_filtered, rg_bloom_filtered, - rg_vector_filtered, rows_before_filter, rows_fulltext_filtered, rows_inverted_filtered, rows_bloom_filtered, - rows_vector_filtered, - rows_vector_selected, rows_precise_filtered, fulltext_index_cache_hit, fulltext_index_cache_miss, @@ -409,9 +400,6 @@ impl fmt::Debug for ScanMetricsSet { if *rg_bloom_filtered > 0 { write!(f, ", \"rg_bloom_filtered\":{rg_bloom_filtered}")?; } - if *rg_vector_filtered > 0 { - write!(f, ", \"rg_vector_filtered\":{rg_vector_filtered}")?; - } if *rows_fulltext_filtered > 0 { write!(f, ", \"rows_fulltext_filtered\":{rows_fulltext_filtered}")?; } @@ -421,12 +409,6 @@ impl fmt::Debug for ScanMetricsSet { if *rows_bloom_filtered > 0 { write!(f, ", \"rows_bloom_filtered\":{rows_bloom_filtered}")?; } - if *rows_vector_filtered > 0 { - write!(f, ", \"rows_vector_filtered\":{rows_vector_filtered}")?; - } - if *rows_vector_selected > 0 { - write!(f, ", \"rows_vector_selected\":{rows_vector_selected}")?; - } if *rows_precise_filtered > 0 { write!(f, ", \"rows_precise_filtered\":{rows_precise_filtered}")?; } @@ -677,13 +659,10 @@ impl ScanMetricsSet { rg_inverted_filtered, rg_minmax_filtered, rg_bloom_filtered, - rg_vector_filtered, rows_total, rows_fulltext_filtered, rows_inverted_filtered, rows_bloom_filtered, - rows_vector_filtered, - rows_vector_selected, rows_precise_filtered, fulltext_index_cache_hit, fulltext_index_cache_miss, @@ -721,14 +700,11 @@ impl ScanMetricsSet { self.rg_inverted_filtered += *rg_inverted_filtered; self.rg_minmax_filtered += *rg_minmax_filtered; self.rg_bloom_filtered += *rg_bloom_filtered; - self.rg_vector_filtered += *rg_vector_filtered; self.rows_before_filter += *rows_total; self.rows_fulltext_filtered += *rows_fulltext_filtered; self.rows_inverted_filtered += *rows_inverted_filtered; self.rows_bloom_filtered += *rows_bloom_filtered; - self.rows_vector_filtered += *rows_vector_filtered; - self.rows_vector_selected += *rows_vector_selected; self.rows_precise_filtered += *rows_precise_filtered; self.fulltext_index_cache_hit += *fulltext_index_cache_hit; @@ -861,10 +837,6 @@ impl ScanMetricsSet { READ_ROW_GROUPS_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rg_bloom_filtered as u64); - #[cfg(feature = "vector_index")] - READ_ROW_GROUPS_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rg_vector_filtered as u64); PRECISE_FILTER_ROWS_TOTAL .with_label_values(&["parquet"]) @@ -881,10 +853,6 @@ impl ScanMetricsSet { READ_ROWS_IN_ROW_GROUP_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rows_bloom_filtered as u64); - #[cfg(feature = "vector_index")] - READ_ROWS_IN_ROW_GROUP_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rows_vector_filtered as u64); } } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index f160bbb3785..c18d6195f19 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -361,8 +361,7 @@ pub enum IndexType { FulltextIndex, /// Bloom Filter index BloomFilterIndex, - /// Vector index (HNSW). - #[cfg(feature = "vector_index")] + /// Retained for manifests written with the removed experimental vector index. VectorIndex, } @@ -426,12 +425,6 @@ impl FileMeta { .contains(&IndexType::BloomFilterIndex) } - /// Returns true if the file has a vector index. - #[cfg(feature = "vector_index")] - pub fn vector_index_available(&self) -> bool { - self.available_indexes.contains(&IndexType::VectorIndex) - } - pub fn index_file_size(&self) -> u64 { self.index_file_size } @@ -1062,6 +1055,31 @@ mod tests { assert_eq!(file_meta, deserialized_file_meta); } + #[test] + fn test_deserialize_legacy_vector_index() { + let json = r#"{ + "region_id": 0, + "file_id": "bc5896ec-e4d8-4017-a80d-f2de73188d55", + "time_range": [{"value":0,"unit":"Millisecond"},{"value":0,"unit":"Millisecond"}], + "available_indexes": ["VectorIndex", "InvertedIndex"], + "indexes": [{"column_id": 1, "created_indexes": ["VectorIndex"]}], + "level": 0 + }"#; + let meta: FileMeta = serde_json::from_str(json).unwrap(); + assert!(meta.inverted_index_available()); + assert!(!meta.fulltext_index_available()); + assert!(!meta.bloom_filter_index_available()); + let encoded = serde_json::to_value(&meta).unwrap(); + assert_eq!( + encoded["available_indexes"], + serde_json::json!(["VectorIndex", "InvertedIndex"]) + ); + assert_eq!( + encoded["indexes"][0]["created_indexes"], + serde_json::json!(["VectorIndex"]) + ); + } + #[test] fn test_file_meta_with_partition_expr() { let file_id = FileId::random(); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 2dd29e9803e..32d0419a33c 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -26,8 +26,6 @@ pub mod puffin_manager; mod sparse_test; mod statistics; pub(crate) mod store; -#[cfg(feature = "vector_index")] -pub(crate) mod vector_index; use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -48,15 +46,11 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId}; use strum::IntoStaticStr; use tokio::sync::mpsc::Sender; -#[cfg(feature = "vector_index")] -use vector_index::creator::VectorIndexer; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; -#[cfg(feature = "vector_index")] -use crate::config::VectorIndexConfig; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; use crate::error::{ BuildIndexAsyncSnafu, Error, JoinSnafu, RegionClosedSnafu, RegionDroppedSnafu, @@ -89,8 +83,6 @@ use crate::worker::WorkerListener; pub(crate) const TYPE_INVERTED_INDEX: &str = "inverted_index"; pub(crate) const TYPE_FULLTEXT_INDEX: &str = "fulltext_index"; pub(crate) const TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index"; -#[cfg(feature = "vector_index")] -pub(crate) const TYPE_VECTOR_INDEX: &str = "vector_index"; /// Evicts local state for a stale index artifact. /// @@ -170,9 +162,6 @@ pub struct IndexOutput { pub fulltext_index: FulltextIndexOutput, /// Bloom filter output. pub bloom_filter: BloomFilterOutput, - /// Vector index output. - #[cfg(feature = "vector_index")] - pub vector_index: VectorIndexOutput, } impl IndexOutput { @@ -187,10 +176,6 @@ impl IndexOutput { if self.bloom_filter.is_available() { indexes.push(IndexType::BloomFilterIndex); } - #[cfg(feature = "vector_index")] - if self.vector_index.is_available() { - indexes.push(IndexType::VectorIndex); - } indexes } @@ -214,12 +199,6 @@ impl IndexOutput { .push(IndexType::BloomFilterIndex); } } - #[cfg(feature = "vector_index")] - if self.vector_index.is_available() { - for &col in &self.vector_index.columns { - map.entry(col).or_default().push(IndexType::VectorIndex); - } - } map.into_iter() .map(|(column_id, created_indexes)| ColumnIndexMetadata { @@ -253,9 +232,6 @@ pub type InvertedIndexOutput = IndexBaseOutput; pub type FulltextIndexOutput = IndexBaseOutput; /// Output of the bloom filter creation. pub type BloomFilterOutput = IndexBaseOutput; -/// Output of the vector index creation. -#[cfg(feature = "vector_index")] -pub type VectorIndexOutput = IndexBaseOutput; /// The index creator that hides the error handling details. #[derive(Default)] @@ -276,10 +252,6 @@ pub struct Indexer { 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")] - last_mem_vector_index: usize, intermediate_manager: Option, } @@ -388,18 +360,6 @@ impl Indexer { .with_label_values(&[TYPE_BLOOM_FILTER_INDEX]) .add(bloom_filter_mem as i64 - self.last_mem_bloom_filter as i64); self.last_mem_bloom_filter = bloom_filter_mem; - - #[cfg(feature = "vector_index")] - { - let vector_mem = self - .vector_indexer - .as_ref() - .map_or(0, |creator| creator.memory_usage()); - INDEX_CREATE_MEMORY_USAGE - .with_label_values(&[TYPE_VECTOR_INDEX]) - .add(vector_mem as i64 - self.last_mem_vector_index as i64); - self.last_mem_vector_index = vector_mem; - } } } @@ -424,8 +384,6 @@ pub(crate) struct IndexerBuilderImpl { pub(crate) inverted_index_config: InvertedIndexConfig, pub(crate) fulltext_index_config: FulltextIndexConfig, pub(crate) bloom_filter_index_config: BloomFilterConfig, - #[cfg(feature = "vector_index")] - pub(crate) vector_index_config: VectorIndexConfig, } #[async_trait::async_trait] @@ -451,18 +409,8 @@ impl IndexerBuilder for IndexerBuilderImpl { 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()); - } indexer.intermediate_manager = Some(self.intermediate_manager.clone()); - #[cfg(feature = "vector_index")] - let has_any_indexer = indexer.inverted_indexer.is_some() - || indexer.fulltext_indexer.is_some() - || indexer.bloom_filter_indexer.is_some() - || indexer.vector_indexer.is_some(); - #[cfg(not(feature = "vector_index"))] let has_any_indexer = indexer.inverted_indexer.is_some() || indexer.fulltext_indexer.is_some() || indexer.bloom_filter_indexer.is_some(); @@ -641,69 +589,6 @@ impl IndexerBuilderImpl { None } - - #[cfg(feature = "vector_index")] - fn build_vector_indexer(&self, file_id: FileId) -> Option { - let create = match self.build_type { - IndexBuildType::Flush => self.vector_index_config.create_on_flush.auto(), - IndexBuildType::Compact => self.vector_index_config.create_on_compaction.auto(), - _ => true, - }; - - if !create { - debug!( - "Skip creating vector index due to config, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - return None; - } - - // Get vector index column IDs and options from metadata - let vector_index_options = self.metadata.vector_indexed_column_ids(); - if vector_index_options.is_empty() { - debug!( - "No vector columns to index, skip creating vector index, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - return None; - } - - let mem_limit = self.vector_index_config.mem_threshold_on_create(); - let indexer = VectorIndexer::new( - file_id, - &self.metadata, - self.intermediate_manager.clone(), - mem_limit, - &vector_index_options, - ); - - let err = match indexer { - Ok(indexer) => { - if indexer.is_none() { - debug!( - "Skip creating vector index due to no columns require indexing, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - } - return indexer; - } - Err(err) => err, - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to create vector index, region_id: {}, file_id: {}, err: {:?}", - self.metadata.region_id, file_id, err - ); - } else { - warn!( - err; "Failed to create vector index, region_id: {}, file_id: {}", - self.metadata.region_id, file_id, - ); - } - - None - } } /// Type of an index build task. @@ -1565,8 +1450,6 @@ mod tests { with_inverted: bool, with_fulltext: bool, with_skipping_bloom: bool, - #[cfg(feature = "vector_index")] - with_vector: bool, } async fn seed_manifest_file(manifest_ctx: &ManifestContextRef, file_meta: &FileMeta) { @@ -1593,8 +1476,6 @@ mod tests { with_inverted, with_fulltext, with_skipping_bloom, - #[cfg(feature = "vector_index")] - with_vector, }: MetaConfig, ) -> RegionMetadataRef { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2)); @@ -1660,24 +1541,6 @@ mod tests { builder.push_column_metadata(column); } - #[cfg(feature = "vector_index")] - if with_vector { - use index::vector::VectorIndexOptions; - - let options = VectorIndexOptions::default(); - let column_schema = - ColumnSchema::new("vec", ConcreteDataType::vector_datatype(4), true) - .with_vector_index_options(&options) - .unwrap(); - let column = ColumnMetadata { - column_schema, - semantic_type: SemanticType::Field, - column_id: 6, - }; - - builder.push_column_metadata(column); - } - Arc::new(builder.build().unwrap()) } @@ -1735,8 +1598,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); env.access_layer @@ -1786,8 +1647,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }) } @@ -1801,8 +1660,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1814,8 +1671,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1835,8 +1690,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1851,8 +1704,6 @@ mod tests { }, fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1874,8 +1725,6 @@ mod tests { ..Default::default() }, bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1897,8 +1746,6 @@ mod tests { create_on_compaction: Mode::Disable, ..Default::default() }, - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1918,8 +1765,6 @@ mod tests { with_inverted: false, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1931,8 +1776,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1945,8 +1788,6 @@ mod tests { with_inverted: true, with_fulltext: false, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1958,8 +1799,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -1972,8 +1811,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: false, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -1985,8 +1822,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(1024)) .await; @@ -2006,8 +1841,6 @@ mod tests { with_inverted: true, with_fulltext: true, with_skipping_bloom: true, - #[cfg(feature = "vector_index")] - with_vector: false, }); let indexer = IndexerBuilderImpl { build_type: IndexBuildType::Flush, @@ -2019,8 +1852,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } .build(random_region_file_id(), 0, Some(0)) .await; @@ -2028,81 +1859,6 @@ mod tests { assert!(indexer.inverted_indexer.is_some()); } - #[cfg(feature = "vector_index")] - #[tokio::test] - async fn test_update_flat_builds_vector_index() { - use datatypes::arrow::array::BinaryBuilder; - use datatypes::arrow::datatypes::{DataType, Field, Schema}; - - struct TestPathProvider; - - impl FilePathProvider for TestPathProvider { - fn build_index_file_path(&self, file_id: RegionFileId) -> String { - format!("index/{}.puffin", file_id) - } - - fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String { - format!("index/{}.puffin", index_id) - } - - fn build_sst_file_path(&self, file_id: RegionFileId) -> String { - format!("sst/{}.parquet", file_id) - } - } - - fn f32s_to_bytes(values: &[f32]) -> Vec { - let mut bytes = Vec::with_capacity(values.len() * 4); - for v in values { - bytes.extend_from_slice(&v.to_le_bytes()); - } - bytes - } - - let (dir, factory) = - PuffinManagerFactory::new_for_test_async("test_update_flat_builds_vector_index_").await; - let intm_manager = mock_intm_mgr(dir.path().to_string_lossy()).await; - - let metadata = mock_region_metadata(MetaConfig { - with_inverted: false, - with_fulltext: false, - with_skipping_bloom: false, - with_vector: true, - }); - - let mut indexer = IndexerBuilderImpl { - build_type: IndexBuildType::Flush, - metadata, - puffin_manager: factory.build(mock_object_store(), TestPathProvider), - write_cache_enabled: false, - intermediate_manager: intm_manager, - index_options: IndexOptions::default(), - inverted_index_config: InvertedIndexConfig::default(), - fulltext_index_config: FulltextIndexConfig::default(), - bloom_filter_index_config: BloomFilterConfig::default(), - vector_index_config: Default::default(), - } - .build(random_region_file_id(), 0, Some(1024)) - .await; - - assert!(indexer.vector_indexer.is_some()); - - let vec1 = f32s_to_bytes(&[1.0, 0.0, 0.0, 0.0]); - let vec2 = f32s_to_bytes(&[0.0, 1.0, 0.0, 0.0]); - - let mut builder = BinaryBuilder::with_capacity(2, vec1.len() + vec2.len()); - builder.append_value(&vec1); - builder.append_value(&vec2); - - let schema = Arc::new(Schema::new(vec![Field::new("vec", DataType::Binary, true)])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(builder.finish())]).unwrap(); - - indexer.update_flat(&batch).await; - let output = indexer.finish().await; - - assert!(output.vector_index.is_available()); - assert!(output.vector_index.columns.contains(&6)); - } - #[tokio::test] async fn test_index_build_task_sst_not_exist() { let env = SchedulerEnv::new().await; @@ -2513,8 +2269,6 @@ mod tests { inverted_index_config: InvertedIndexConfig::default(), fulltext_index_config: FulltextIndexConfig::default(), bloom_filter_index_config: BloomFilterConfig::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }); let sst_info = mock_sst_file(metadata.clone(), &env, IndexBuildMode::Async).await; diff --git a/src/mito2/src/sst/index/column_test.rs b/src/mito2/src/sst/index/column_test.rs index 34ff1c19fde..fb7677eff58 100644 --- a/src/mito2/src/sst/index/column_test.rs +++ b/src/mito2/src/sst/index/column_test.rs @@ -199,8 +199,6 @@ async fn materialized_runs_match_legacy_index_bytes_and_bloom_lookups() { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut expected = None; // Demand is the union of active creators' PK columns, not their counts or diff --git a/src/mito2/src/sst/index/indexer/abort.rs b/src/mito2/src/sst/index/indexer/abort.rs index 24cbb797cbd..36c002f13f0 100644 --- a/src/mito2/src/sst/index/indexer/abort.rs +++ b/src/mito2/src/sst/index/indexer/abort.rs @@ -24,8 +24,6 @@ impl Indexer { self.do_abort_inverted_index().await; self.do_abort_fulltext_index().await; self.do_abort_bloom_filter().await; - #[cfg(feature = "vector_index")] - self.do_abort_vector_index().await; self.do_prune_intm_sst_dir().await; if self.write_cache_enabled { self.do_abort_clean_fs_temp_dir().await; @@ -109,26 +107,4 @@ impl Indexer { .to_string(); TempFileCleaner::clean_atomic_dir_files(fs_accessor.store().store(), &[&fs_handle]).await; } - - #[cfg(feature = "vector_index")] - async fn do_abort_vector_index(&mut self) { - let Some(mut indexer) = self.vector_indexer.take() else { - return; - }; - let Err(err) = indexer.abort().await else { - return; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to abort vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to abort vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - } } diff --git a/src/mito2/src/sst/index/indexer/finish.rs b/src/mito2/src/sst/index/indexer/finish.rs index 823cd5c2b8a..6a6c3f34d0b 100644 --- a/src/mito2/src/sst/index/indexer/finish.rs +++ b/src/mito2/src/sst/index/indexer/finish.rs @@ -17,8 +17,6 @@ use puffin::puffin_manager::{PuffinManager, PuffinWriter}; use store_api::storage::ColumnId; use crate::sst::file::{RegionFileId, RegionIndexId}; -#[cfg(feature = "vector_index")] -use crate::sst::index::VectorIndexOutput; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::statistics::{ByteCount, RowCount}; use crate::sst::index::{ @@ -57,15 +55,6 @@ impl Indexer { return IndexOutput::default(); } - #[cfg(feature = "vector_index")] - { - let success = self.do_finish_vector_index(&mut writer, &mut output).await; - if !success { - self.do_abort().await; - return IndexOutput::default(); - } - } - self.do_prune_intm_sst_dir().await; output.file_size = self.do_finish_puffin_writer(writer).await; output.version = self.index_version; @@ -288,63 +277,6 @@ impl Indexer { output.columns = column_ids; } - #[cfg(feature = "vector_index")] - async fn do_finish_vector_index( - &mut self, - puffin_writer: &mut SstPuffinWriter, - index_output: &mut IndexOutput, - ) -> bool { - let Some(mut indexer) = self.vector_indexer.take() else { - return true; - }; - - let column_ids = indexer.column_ids().collect(); - let err = match indexer.finish(puffin_writer).await { - Ok((row_count, byte_count)) => { - self.fill_vector_index_output( - &mut index_output.vector_index, - row_count, - byte_count, - column_ids, - ); - return true; - } - Err(err) => err, - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to finish vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to finish vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } - - #[cfg(feature = "vector_index")] - fn fill_vector_index_output( - &mut self, - output: &mut VectorIndexOutput, - row_count: RowCount, - byte_count: ByteCount, - column_ids: Vec, - ) { - debug!( - "Vector index created, region_id: {}, file_id: {}, written_bytes: {}, written_rows: {}, columns: {:?}", - self.region_id, self.file_id, byte_count, row_count, column_ids - ); - - output.index_size = byte_count; - output.row_count = row_count; - output.columns = column_ids; - } - pub(crate) async fn do_prune_intm_sst_dir(&mut self) { if let Some(manager) = self.intermediate_manager.take() && let Err(e) = manager.prune_sst_dir(&self.region_id, &self.file_id).await diff --git a/src/mito2/src/sst/index/indexer/update.rs b/src/mito2/src/sst/index/indexer/update.rs index 0168bee876e..7e2e1e35022 100644 --- a/src/mito2/src/sst/index/indexer/update.rs +++ b/src/mito2/src/sst/index/indexer/update.rs @@ -38,10 +38,6 @@ impl Indexer { if !self.do_update_bloom_filter(batch).await { self.do_abort().await; } - #[cfg(feature = "vector_index")] - if !self.do_update_vector_index(batch).await { - self.do_abort().await; - } } /// Handles decode errors before entering asynchronous cleanup, following the @@ -136,32 +132,6 @@ impl Indexer { false } - /// Returns false if the update failed. - #[cfg(feature = "vector_index")] - async fn do_update_vector_index(&mut self, batch: &mut Batch) -> bool { - let Some(creator) = self.vector_indexer.as_mut() else { - return true; - }; - - let Err(err) = creator.update(batch).await else { - return true; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to update vector index, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to update vector index, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } - pub(crate) async fn do_update_flat(&mut self, batch: &RecordBatch) { if batch.num_rows() == 0 { return; @@ -176,10 +146,6 @@ impl Indexer { if !self.do_update_flat_bloom_filter(batch).await { self.do_abort().await; } - #[cfg(feature = "vector_index")] - if !self.do_update_flat_vector_index(batch).await { - self.do_abort().await; - } } /// Returns false if the update failed. @@ -256,30 +222,4 @@ impl Indexer { false } - - /// Returns false if the update failed. - #[cfg(feature = "vector_index")] - async fn do_update_flat_vector_index(&mut self, batch: &RecordBatch) -> bool { - let Some(creator) = self.vector_indexer.as_mut() else { - return true; - }; - - let Err(err) = creator.update_flat(batch).await else { - return true; - }; - - if cfg!(any(test, feature = "test")) { - panic!( - "Failed to update vector index with flat format, region_id: {}, file_id: {}, err: {:?}", - self.region_id, self.file_id, err - ); - } else { - warn!( - err; "Failed to update vector index with flat format, region_id: {}, file_id: {}", - self.region_id, self.file_id, - ); - } - - false - } } diff --git a/src/mito2/src/sst/index/sparse_test.rs b/src/mito2/src/sst/index/sparse_test.rs index 15d26bd1bc7..c6031d50ad6 100644 --- a/src/mito2/src/sst/index/sparse_test.rs +++ b/src/mito2/src/sst/index/sparse_test.rs @@ -141,8 +141,6 @@ async fn sparse_and_materialized_tags_produce_identical_indexes() { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut expected = None; for batch in [materialized, sparse, mixed] { diff --git a/src/mito2/src/sst/index/vector_index.rs b/src/mito2/src/sst/index/vector_index.rs deleted file mode 100644 index f1b9aa3899e..00000000000 --- a/src/mito2/src/sst/index/vector_index.rs +++ /dev/null @@ -1,24 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index module for HNSW-based approximate nearest neighbor search. -//! -//! This module provides the mito2 integration layer for vector indexes. -//! The core algorithms (creator, reader, applier) are in the `index` crate. - -pub(crate) mod applier; -pub(crate) mod creator; - -/// The blob type identifier for vector index in puffin files. -pub(crate) const INDEX_BLOB_TYPE: &str = "greptime-vector-index-v1"; diff --git a/src/mito2/src/sst/index/vector_index/applier.rs b/src/mito2/src/sst/index/vector_index/applier.rs deleted file mode 100644 index 40c8c257277..00000000000 --- a/src/mito2/src/sst/index/vector_index/applier.rs +++ /dev/null @@ -1,503 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector index applier for KNN search. - -use std::sync::Arc; - -use common_base::range_read::RangeReader; -use common_telemetry::warn; -use index::vector::VectorDistanceMetric; -use index::vector::apply::HnswVectorIndexApplier; -use puffin::puffin_manager::cache::PuffinMetadataCacheRef; -use puffin::puffin_manager::{PuffinManager, PuffinReader}; -use snafu::ResultExt; -use store_api::storage::ColumnId; - -use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider}; -use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; -use crate::cache::index::vector_index::{ - CachedVectorIndex, VectorIndexCacheKey, VectorIndexCacheRef, -}; -use crate::error::{ApplyVectorIndexSnafu, PuffinBuildReaderSnafu, PuffinReadBlobSnafu, Result}; -use crate::sst::file::RegionIndexId; -use crate::sst::index::puffin_manager::{BlobReader, PuffinManagerFactory}; -use crate::sst::index::trigger_index_background_download; -use crate::sst::index::vector_index::INDEX_BLOB_TYPE; - -/// Result of applying vector index. -#[derive(Debug)] -pub struct VectorIndexApplyOutput { - /// Row offsets in the SST file. - pub row_offsets: Vec, - /// Distances to the query vector. - /// TODO(dennis): Not used yet, reserved for future use. - #[allow(dead_code)] - pub distances: Vec, -} - -/// Vector index applier for KNN search against SST blobs. -pub struct VectorIndexApplier { - table_dir: String, - path_type: store_api::region_request::PathType, - object_store: object_store::ObjectStore, - puffin_manager_factory: PuffinManagerFactory, - file_cache: Option, - puffin_metadata_cache: Option, - vector_index_cache: Option, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, -} - -pub type VectorIndexApplierRef = Arc; - -impl VectorIndexApplier { - pub fn new( - table_dir: String, - path_type: store_api::region_request::PathType, - object_store: object_store::ObjectStore, - puffin_manager_factory: PuffinManagerFactory, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, - ) -> Self { - Self { - table_dir, - path_type, - object_store, - puffin_manager_factory, - file_cache: None, - puffin_metadata_cache: None, - vector_index_cache: None, - column_id, - query_vector, - metric, - } - } - - pub fn with_file_cache(mut self, file_cache: Option) -> Self { - self.file_cache = file_cache; - self - } - - pub fn with_puffin_metadata_cache( - mut self, - puffin_metadata_cache: Option, - ) -> Self { - self.puffin_metadata_cache = puffin_metadata_cache; - self - } - - pub fn with_vector_index_cache(mut self, cache: Option) -> Self { - self.vector_index_cache = cache; - self - } - - /// Applies vector index to the file and returns candidates. - /// - /// This method loads the vector index blob (from cache or remote), runs - /// a KNN search against the indexed vectors, and maps the HNSW keys back - /// to row offsets in the SST file. It returns only row offsets; callers - /// are responsible for any higher-level ordering or limit enforcement. - pub async fn apply_with_k( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - k: usize, - ) -> Result { - if k == 0 { - return Ok(VectorIndexApplyOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - } - - let cached = self.load_or_read_index(file_id, file_size_hint).await?; - let Some(cached) = cached else { - return Ok(VectorIndexApplyOutput { - row_offsets: Vec::new(), - distances: Vec::new(), - }); - }; - - let applier = cached.applier(); - if self.query_vector.len() != applier.dimensions() as usize { - return ApplyVectorIndexSnafu { - reason: format!( - "Query vector dimension {} does not match index dimension {}", - self.query_vector.len(), - applier.dimensions() - ), - } - .fail(); - } - if self.metric != applier.metric() { - return ApplyVectorIndexSnafu { - reason: format!( - "Query metric {} does not match index metric {}", - self.metric, - applier.metric() - ), - } - .fail(); - } - - let output = applier.search(&self.query_vector, k).map_err(|e| { - ApplyVectorIndexSnafu { - reason: e.to_string(), - } - .build() - })?; - - Ok(VectorIndexApplyOutput { - row_offsets: output.row_offsets, - distances: output.distances, - }) - } - - async fn load_or_read_index( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result>> { - let cache_key = - VectorIndexCacheKey::new(file_id.file_id(), file_id.version, self.column_id); - if let Some(cache) = &self.vector_index_cache - && let Some(cached) = cache.get(&cache_key) - { - return Ok(Some(cached)); - } - - let reader = match self.cached_blob_reader(file_id, file_size_hint).await { - Ok(Some(reader)) => reader, - Ok(None) => self.remote_blob_reader(file_id, file_size_hint).await?, - Err(err) => { - if is_blob_not_found(&err) { - self.remote_blob_reader(file_id, file_size_hint).await? - } else { - warn!(err; "Failed to read cached vector index blob, fallback to remote"); - self.remote_blob_reader(file_id, file_size_hint).await? - } - } - }; - - let blob_data = read_all_blob(reader, file_size_hint).await?; - if blob_data.is_empty() { - return Ok(None); - } - - let applier = HnswVectorIndexApplier::from_blob(&blob_data).map_err(|e| { - ApplyVectorIndexSnafu { - reason: e.to_string(), - } - .build() - })?; - - let cached = Arc::new(CachedVectorIndex::new(Box::new(applier))); - if let Some(cache) = &self.vector_index_cache { - cache.insert(cache_key, cached.clone()); - } - - Ok(Some(cached)) - } - - async fn cached_blob_reader( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result> { - let Some(file_cache) = &self.file_cache else { - return Ok(None); - }; - - let index_key = IndexKey::new( - file_id.region_id(), - file_id.file_id(), - FileType::Puffin(file_id.version), - ); - if file_cache.get(index_key).await.is_none() { - return Ok(None); - } - - let puffin_manager = self.puffin_manager_factory.build( - file_cache.local_store(), - WriteCachePathProvider::new(file_cache.clone()), - ); - let blob_name = column_blob_name(self.column_id); - - let reader = puffin_manager - .reader(&file_id) - .await - .context(PuffinBuildReaderSnafu)? - .with_file_size_hint(file_size_hint) - .blob(&blob_name) - .await - .context(PuffinReadBlobSnafu)? - .reader() - .await - .context(PuffinBuildReaderSnafu)?; - Ok(Some(reader)) - } - - async fn remote_blob_reader( - &self, - file_id: RegionIndexId, - file_size_hint: Option, - ) -> Result { - let path_factory = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type); - - trigger_index_background_download( - self.file_cache.as_ref(), - &file_id, - file_size_hint, - &path_factory, - &self.object_store, - ); - - let puffin_manager = self - .puffin_manager_factory - .build(self.object_store.clone(), path_factory) - .with_puffin_metadata_cache(self.puffin_metadata_cache.clone()); - - let blob_name = column_blob_name(self.column_id); - - puffin_manager - .reader(&file_id) - .await - .context(PuffinBuildReaderSnafu)? - .with_file_size_hint(file_size_hint) - .blob(&blob_name) - .await - .context(PuffinReadBlobSnafu)? - .reader() - .await - .context(PuffinBuildReaderSnafu) - } -} - -fn column_blob_name(column_id: ColumnId) -> String { - format!("{INDEX_BLOB_TYPE}-{}", column_id) -} - -fn is_blob_not_found(err: &crate::error::Error) -> bool { - matches!( - err, - crate::error::Error::PuffinReadBlob { - source: puffin::error::Error::BlobNotFound { .. }, - .. - } - ) -} - -async fn read_all_blob(reader: BlobReader, file_size_hint: Option) -> Result> { - let metadata = reader.metadata().await.map_err(|e| { - ApplyVectorIndexSnafu { - reason: format!("Failed to read vector index metadata: {}", e), - } - .build() - })?; - if let Some(limit) = file_size_hint - && metadata.content_length > limit - { - return ApplyVectorIndexSnafu { - reason: format!( - "Vector index blob size {} exceeds file size hint {}", - metadata.content_length, limit - ), - } - .fail(); - } - let bytes = reader.read(0..metadata.content_length).await.map_err(|e| { - ApplyVectorIndexSnafu { - reason: format!("Failed to read vector index data: {}", e), - } - .build() - })?; - Ok(bytes.to_vec()) -} - -#[cfg(test)] -mod tests { - use common_test_util::temp_dir::TempDir; - use futures::io::Cursor; - use greptime_proto::v1::index::{VectorIndexMeta, VectorIndexStats}; - use index::vector::engine::{self, VectorIndexConfig}; - use index::vector::format::{distance_metric_to_proto, engine_type_to_proto}; - use object_store::ObjectStore; - use object_store::services::Memory; - use prost::Message; - use puffin::puffin_manager::PuffinWriter; - use roaring::RoaringBitmap; - use store_api::region_request::PathType; - use store_api::storage::{ColumnId, FileId, VectorIndexEngineType}; - - use super::*; - use crate::access_layer::RegionFilePathFactory; - use crate::sst::file::RegionFileId; - use crate::sst::index::puffin_manager::PuffinManagerFactory; - - async fn build_applier_with_blob( - blob: Vec, - column_id: ColumnId, - query_vector: Vec, - metric: VectorDistanceMetric, - ) -> (TempDir, VectorIndexApplier, RegionIndexId, u64) { - let (dir, puffin_manager_factory) = - PuffinManagerFactory::new_for_test_async("test_vector_index_applier_").await; - let object_store = ObjectStore::new(Memory::default()).unwrap(); - let file_id = RegionFileId::new(0.into(), FileId::random()); - let index_id = RegionIndexId::new(file_id, 0); - let table_dir = "table_dir".to_string(); - - let puffin_manager = puffin_manager_factory.build( - object_store.clone(), - RegionFilePathFactory::new(table_dir.clone(), PathType::Bare), - ); - let mut writer = puffin_manager.writer(&index_id).await.unwrap(); - let blob_name = column_blob_name(column_id); - let _bytes_written = writer - .put_blob( - blob_name.as_str(), - Cursor::new(blob), - Default::default(), - Default::default(), - ) - .await - .unwrap(); - let file_size = writer.finish().await.unwrap(); - - let applier = VectorIndexApplier::new( - table_dir, - PathType::Bare, - object_store, - puffin_manager_factory, - column_id, - query_vector, - metric, - ); - - (dir, applier, index_id, file_size) - } - - fn build_blob_with_vectors( - config: &VectorIndexConfig, - vectors: Vec<(u64, Vec)>, - null_bitmap: &RoaringBitmap, - total_rows: u64, - indexed_rows: u64, - ) -> Vec { - let mut engine = engine::create_engine(config.engine, config).unwrap(); - for (key, vector) in vectors { - engine.add(key, &vector).unwrap(); - } - let index_size = engine.serialized_length(); - let mut index_bytes = vec![0u8; index_size]; - engine.save_to_buffer(&mut index_bytes).unwrap(); - - let mut null_bitmap_bytes = Vec::new(); - null_bitmap.serialize_into(&mut null_bitmap_bytes).unwrap(); - - let null_count = total_rows - indexed_rows; - let meta = VectorIndexMeta { - engine: engine_type_to_proto(config.engine).into(), - dim: config.dim as u32, - metric: distance_metric_to_proto(config.distance_metric).into(), - connectivity: config.connectivity as u32, - expansion_add: config.expansion_add as u32, - expansion_search: config.expansion_search as u32, - null_bitmap_size: null_bitmap_bytes.len() as u32, - index_size: index_bytes.len() as u64, - stats: Some(VectorIndexStats { - total_row_count: total_rows, - indexed_row_count: indexed_rows, - null_count, - }), - }; - let meta_bytes = meta.encode_to_vec(); - let meta_size = meta_bytes.len() as u32; - - // Build blob: null_bitmap | index_data | meta | meta_size - let mut blob = Vec::new(); - blob.extend_from_slice(&null_bitmap_bytes); - blob.extend_from_slice(&index_bytes); - blob.extend_from_slice(&meta_bytes); - blob.extend_from_slice(&meta_size.to_le_bytes()); - blob - } - - fn test_config() -> VectorIndexConfig { - VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 2, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - } - } - - #[tokio::test] - async fn test_apply_with_k_returns_offsets() { - let config = test_config(); - let mut null_bitmap = RoaringBitmap::new(); - null_bitmap.insert(1); - let blob = build_blob_with_vectors( - &config, - vec![(0, vec![1.0, 0.0]), (1, vec![0.0, 1.0])], - &null_bitmap, - 3, - 2, - ); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0, 0.0], VectorDistanceMetric::L2sq).await; - let output = applier - .apply_with_k(index_id, Some(size_bytes), 2) - .await - .unwrap(); - assert_eq!(output.row_offsets, vec![0, 2]); - } - - #[tokio::test] - async fn test_apply_with_k_dimension_mismatch() { - let config = test_config(); - let null_bitmap = RoaringBitmap::new(); - let blob = build_blob_with_vectors(&config, vec![(0, vec![1.0, 0.0])], &null_bitmap, 1, 1); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0, 0.0, 0.0], VectorDistanceMetric::L2sq).await; - let res = applier.apply_with_k(index_id, Some(size_bytes), 1).await; - assert!(res.is_err()); - } - - #[tokio::test] - async fn test_apply_with_k_empty_blob() { - let config = VectorIndexConfig { - engine: VectorIndexEngineType::Usearch, - dim: 1, - distance_metric: VectorDistanceMetric::L2sq, - connectivity: 16, - expansion_add: 128, - expansion_search: 64, - }; - let null_bitmap = RoaringBitmap::new(); - let blob = build_blob_with_vectors(&config, Vec::new(), &null_bitmap, 0, 0); - let (_dir, applier, index_id, size_bytes) = - build_applier_with_blob(blob, 1, vec![1.0], VectorDistanceMetric::L2sq).await; - let output = applier - .apply_with_k(index_id, Some(size_bytes), 1) - .await - .unwrap(); - assert!(output.row_offsets.is_empty()); - } -} diff --git a/src/mito2/src/sst/index/vector_index/creator.rs b/src/mito2/src/sst/index/vector_index/creator.rs deleted file mode 100644 index 2971c12cc66..00000000000 --- a/src/mito2/src/sst/index/vector_index/creator.rs +++ /dev/null @@ -1,536 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Vector indexer for managing vector indexes across multiple columns. - -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; - -use common_telemetry::warn; -use datatypes::arrow::array::{Array, BinaryArray}; -use datatypes::arrow::record_batch::RecordBatch; -use datatypes::data_type::ConcreteDataType; -use datatypes::prelude::ValueRef; -use index::vector::VectorIndexOptions; -use index::vector::create::{HnswVectorIndexCreator, IndexCreator}; -use index::vector::engine::VectorIndexConfig; -use index::vector::util::bytes_to_f32_slice; -use puffin::puffin_manager::PutOptions; -use snafu::ensure; -use store_api::metadata::RegionMetadataRef; -use store_api::storage::{ColumnId, FileId}; - -use crate::error::{ - OperateAbortedIndexSnafu, Result, VectorIndexBuildSnafu, VectorIndexFinishSnafu, -}; -use crate::metrics::{INDEX_CREATE_BYTES_TOTAL, INDEX_CREATE_ROWS_TOTAL}; -use crate::read::Batch; -use crate::sst::index::TYPE_VECTOR_INDEX; -use crate::sst::index::intermediate::{ - IntermediateLocation, IntermediateManager, TempFileProvider, -}; -use crate::sst::index::puffin_manager::SstPuffinWriter; -use crate::sst::index::statistics::{ByteCount, RowCount, Statistics}; -use crate::sst::index::vector_index::INDEX_BLOB_TYPE; - -/// Helper to create VectorIndexConfig from VectorIndexOptions. -fn create_config(dim: usize, options: &VectorIndexOptions) -> VectorIndexConfig { - VectorIndexConfig { - engine: options.engine, - dim, - distance_metric: options.metric, - connectivity: options.connectivity as usize, - expansion_add: options.expansion_add as usize, - expansion_search: options.expansion_search as usize, - } -} - -fn build_err(e: impl std::fmt::Display) -> crate::error::Error { - VectorIndexBuildSnafu { - reason: e.to_string(), - } - .build() -} - -/// The indexer for vector indexes across multiple columns. -pub struct VectorIndexer { - /// Per-column vector index creators. - creators: HashMap, - /// Provider for intermediate files. - temp_file_provider: Arc, - /// Whether the indexing process has been aborted. - aborted: bool, - /// Statistics for this indexer. - stats: Statistics, - /// Global memory usage tracker. - #[allow(dead_code)] - global_memory_usage: Arc, - /// Region metadata for column lookups. - metadata: RegionMetadataRef, - /// Memory usage threshold. - memory_usage_threshold: Option, -} - -impl VectorIndexer { - /// Creates a new vector indexer. - /// - /// Returns `None` if there are no vector columns that need indexing. - pub fn new( - sst_file_id: FileId, - metadata: &RegionMetadataRef, - intermediate_manager: IntermediateManager, - memory_usage_threshold: Option, - vector_index_options: &HashMap, - ) -> Result> { - let mut creators = HashMap::new(); - - let temp_file_provider = Arc::new(TempFileProvider::new( - IntermediateLocation::new(&metadata.region_id, &sst_file_id), - intermediate_manager, - )); - let global_memory_usage = Arc::new(AtomicUsize::new(0)); - - // Find all vector columns that have vector index enabled - for column in &metadata.column_metadatas { - // Check if this column has vector index options configured - let Some(options) = vector_index_options.get(&column.column_id) else { - continue; - }; - - // Verify the column is a vector type - let ConcreteDataType::Vector(vector_type) = &column.column_schema.data_type else { - continue; - }; - - let config = create_config(vector_type.dim as usize, options); - let creator = HnswVectorIndexCreator::new(config).map_err(build_err)?; - creators.insert(column.column_id, creator); - } - - if creators.is_empty() { - return Ok(None); - } - - let indexer = Self { - creators, - temp_file_provider, - aborted: false, - stats: Statistics::new(TYPE_VECTOR_INDEX), - global_memory_usage, - metadata: metadata.clone(), - memory_usage_threshold, - }; - - Ok(Some(indexer)) - } - - /// Updates index with a batch of rows. - /// Garbage will be cleaned up if failed to update. - pub async fn update(&mut self, batch: &mut Batch) -> Result<()> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.creators.is_empty() { - return Ok(()); - } - - if let Err(update_err) = self.do_update(batch).await { - // Clean up garbage if failed to update - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - return Err(update_err); - } - - Ok(()) - } - - /// Updates index with a flat format `RecordBatch`. - /// Garbage will be cleaned up if failed to update. - pub async fn update_flat(&mut self, batch: &RecordBatch) -> Result<()> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.creators.is_empty() || batch.num_rows() == 0 { - return Ok(()); - } - - if let Err(update_err) = self.do_update_flat(batch).await { - // Clean up garbage if failed to update - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - return Err(update_err); - } - - Ok(()) - } - - /// Internal update implementation. - async fn do_update(&mut self, batch: &mut Batch) -> Result<()> { - let mut guard = self.stats.record_update(); - let n = batch.num_rows(); - guard.inc_row_count(n); - - for (col_id, creator) in &mut self.creators { - let Some(values) = batch.field_col_value(*col_id) else { - creator.push_nulls(n).map_err(build_err)?; - continue; - }; - - for i in 0..n { - let value = values.data.get_ref(i); - if value.is_null() { - creator.push_null().map_err(build_err)?; - } else if let ValueRef::Binary(bytes) = value { - let floats = bytes_to_f32_slice(bytes); - if floats.len() != creator.config().dim { - return VectorIndexBuildSnafu { - reason: format!( - "Vector dimension mismatch: expected {}, got {}", - creator.config().dim, - floats.len() - ), - } - .fail(); - } - creator.push_vector(&floats).map_err(build_err)?; - } else { - common_telemetry::debug!( - "Unexpected non-binary value for vector index column, treating as null, data_type: {:?}", - value.data_type() - ); - creator.push_null().map_err(build_err)?; - } - } - - // Check memory limit - abort index creation if exceeded - if let Some(threshold) = self.memory_usage_threshold { - let current_usage = creator.memory_usage(); - if current_usage > threshold { - warn!( - "Vector index memory usage {} exceeds threshold {}, aborting index creation, region_id: {}", - current_usage, threshold, self.metadata.region_id - ); - return VectorIndexBuildSnafu { - reason: format!( - "Memory usage {} exceeds threshold {}", - current_usage, threshold - ), - } - .fail(); - } - } - } - - Ok(()) - } - - /// Internal flat update implementation. - async fn do_update_flat(&mut self, batch: &RecordBatch) -> Result<()> { - let mut guard = self.stats.record_update(); - let n = batch.num_rows(); - guard.inc_row_count(n); - - for (col_id, creator) in &mut self.creators { - // This should never happen: creator exists but column not in metadata - let column_meta = self.metadata.column_by_id(*col_id).ok_or_else(|| { - VectorIndexBuildSnafu { - reason: format!( - "Column {} not found in region metadata, this is a bug", - col_id - ), - } - .build() - })?; - - let column_name = &column_meta.column_schema.name; - // Column not in batch is normal for flat format - treat as NULLs - let Some(column_array) = batch.column_by_name(column_name) else { - creator.push_nulls(n).map_err(build_err)?; - continue; - }; - - let binary_array = column_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - VectorIndexBuildSnafu { - reason: format!( - "Column {} is not a binary array, got {:?}", - column_name, - column_array.data_type() - ), - } - .build() - })?; - - for i in 0..n { - if !binary_array.is_valid(i) { - creator.push_null().map_err(build_err)?; - } else { - let bytes = binary_array.value(i); - let floats = bytes_to_f32_slice(bytes); - if floats.len() != creator.config().dim { - return VectorIndexBuildSnafu { - reason: format!( - "Vector dimension mismatch: expected {}, got {}", - creator.config().dim, - floats.len() - ), - } - .fail(); - } - creator.push_vector(&floats).map_err(build_err)?; - } - } - - if let Some(threshold) = self.memory_usage_threshold { - let current_usage = creator.memory_usage(); - if current_usage > threshold { - warn!( - "Vector index memory usage {} exceeds threshold {}, aborting index creation, region_id: {}", - current_usage, threshold, self.metadata.region_id - ); - return VectorIndexBuildSnafu { - reason: format!( - "Memory usage {} exceeds threshold {}", - current_usage, threshold - ), - } - .fail(); - } - } - } - - Ok(()) - } - - /// Finishes index creation and writes to puffin. - /// Returns the number of rows and bytes written. - pub async fn finish( - &mut self, - puffin_writer: &mut SstPuffinWriter, - ) -> Result<(RowCount, ByteCount)> { - ensure!(!self.aborted, OperateAbortedIndexSnafu); - - if self.stats.row_count() == 0 { - // No IO is performed, no garbage to clean up - return Ok((0, 0)); - } - - let finish_res = self.do_finish(puffin_writer).await; - // Clean up garbage no matter finish successfully or not - if let Err(err) = self.do_cleanup().await { - if cfg!(any(test, feature = "test")) { - panic!("Failed to clean up vector index creator, err: {err:?}"); - } else { - warn!(err; "Failed to clean up vector index creator"); - } - } - - if finish_res.is_ok() { - INDEX_CREATE_ROWS_TOTAL - .with_label_values(&[TYPE_VECTOR_INDEX]) - .inc_by(self.stats.row_count() as u64); - INDEX_CREATE_BYTES_TOTAL - .with_label_values(&[TYPE_VECTOR_INDEX]) - .inc_by(self.stats.byte_count()); - } - - finish_res.map(|_| (self.stats.row_count(), self.stats.byte_count())) - } - - /// Internal finish implementation. - async fn do_finish(&mut self, puffin_writer: &mut SstPuffinWriter) -> Result<()> { - let mut guard = self.stats.record_finish(); - - for (id, creator) in &mut self.creators { - if creator.size() == 0 { - // No vectors to index - continue; - } - - let written_bytes = Self::do_finish_single_creator(*id, creator, puffin_writer).await?; - guard.inc_byte_count(written_bytes); - } - - Ok(()) - } - - /// Finishes a single column's vector index using the VectorIndexCreator trait. - async fn do_finish_single_creator( - col_id: ColumnId, - creator: &mut HnswVectorIndexCreator, - puffin_writer: &mut SstPuffinWriter, - ) -> Result { - let blob_name = format!("{}-{}", INDEX_BLOB_TYPE, col_id); - - let written_bytes = creator - .finish(puffin_writer, &blob_name, PutOptions::default()) - .await - .map_err(|e| { - VectorIndexFinishSnafu { - reason: e.to_string(), - } - .build() - })?; - - Ok(written_bytes) - } - - /// Aborts index creation and cleans up garbage. - pub async fn abort(&mut self) -> Result<()> { - if self.aborted { - return Ok(()); - } - self.aborted = true; - self.do_cleanup().await - } - - /// Cleans up temporary files. - async fn do_cleanup(&mut self) -> Result<()> { - let mut _guard = self.stats.record_cleanup(); - self.creators.clear(); - self.temp_file_provider.cleanup().await - } - - /// Returns the memory usage of the indexer. - pub fn memory_usage(&self) -> usize { - self.creators.values().map(|c| c.memory_usage()).sum() - } - - /// Returns the column IDs being indexed. - pub fn column_ids(&self) -> impl Iterator + '_ { - self.creators.keys().copied() - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - - use api::v1::SemanticType; - use datatypes::data_type::ConcreteDataType; - use datatypes::schema::ColumnSchema; - use datatypes::value::ValueRef; - use datatypes::vectors::{TimestampMillisecondVector, UInt8Vector, UInt64Vector}; - use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; - use store_api::storage::{ColumnId, FileId, RegionId}; - - use super::*; - use crate::read::BatchColumn; - - fn mock_region_metadata_with_vector() -> RegionMetadataRef { - let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("tag", ConcreteDataType::int64_datatype(), true), - semantic_type: SemanticType::Tag, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("vec", ConcreteDataType::vector_datatype(2), true), - semantic_type: SemanticType::Field, - column_id: 3, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "field_u64", - ConcreteDataType::uint64_datatype(), - true, - ), - semantic_type: SemanticType::Field, - column_id: 4, - }) - .primary_key(vec![2]); - - Arc::new(builder.build().unwrap()) - } - - fn new_batch_missing_vector_column(column_id: ColumnId, rows: usize) -> Batch { - let fields = vec![(0, SortField::new(ConcreteDataType::int64_datatype()))]; - let codec = DensePrimaryKeyCodec::with_fields(fields); - let primary_key = codec.encode([ValueRef::Int64(1)].into_iter()).unwrap(); - - let field = BatchColumn { - column_id, - data: Arc::new(UInt64Vector::from_iter_values(0..rows as u64)), - }; - - Batch::new( - primary_key, - Arc::new(TimestampMillisecondVector::from_values( - (0..rows).map(|i| i as i64).collect::>(), - )), - Arc::new(UInt64Vector::from_iter_values(std::iter::repeat_n(0, rows))), - Arc::new(UInt8Vector::from_iter_values(std::iter::repeat_n(1, rows))), - vec![field], - ) - .unwrap() - } - - #[tokio::test] - async fn test_vector_index_missing_column_as_nulls() { - let tempdir = common_test_util::temp_dir::create_temp_dir( - "test_vector_index_missing_column_as_nulls_", - ); - let intm_mgr = IntermediateManager::init_fs(tempdir.path().to_string_lossy()) - .await - .unwrap(); - let region_metadata = mock_region_metadata_with_vector(); - - let mut vector_index_options = HashMap::new(); - vector_index_options.insert(3, VectorIndexOptions::default()); - - let mut indexer = VectorIndexer::new( - FileId::random(), - ®ion_metadata, - intm_mgr, - None, - &vector_index_options, - ) - .unwrap() - .unwrap(); - - let mut batch = new_batch_missing_vector_column(4, 3); - indexer.update(&mut batch).await.unwrap(); - - let creator = indexer.creators.get(&3).unwrap(); - assert_eq!(creator.size(), 0); - assert_eq!(creator.current_row_offset(), 3); - assert_eq!(creator.null_bitmap().len(), 3); - for idx in 0..3 { - assert!(creator.null_bitmap().contains(idx as u32)); - } - } -} diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index b49d2ba1e18..301842bde72 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -1175,8 +1175,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); @@ -1695,8 +1693,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), } } @@ -1816,8 +1812,6 @@ mod tests { inverted_index_config: Default::default(), fulltext_index_config: Default::default(), bloom_filter_index_config: Default::default(), - #[cfg(feature = "vector_index")] - vector_index_config: Default::default(), }; let mut metrics = Metrics::new(WriteType::Flush); diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 714e62ae4af..90c2728e35d 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -14,8 +14,6 @@ //! Parquet reader. -#[cfg(feature = "vector_index")] -use std::collections::BTreeSet; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -55,8 +53,6 @@ use table::predicate::Predicate; use crate::cache::index::result_cache::PredicateKey; use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; -#[cfg(feature = "vector_index")] -use crate::error::ApplyVectorIndexSnafu; use crate::error::{ ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu, UnexpectedSnafu, @@ -79,8 +75,6 @@ use crate::sst::index::fulltext_index::applier::{ use crate::sst::index::inverted_index::applier::{ InvertedIndexApplierRef, InvertedIndexApplyMetrics, }; -#[cfg(feature = "vector_index")] -use crate::sst::index::vector_index::applier::VectorIndexApplierRef; use crate::sst::parquet::file_range::{ FileRangeContext, FileRangeContextRef, PartitionFilterContext, PreFilterMode, RangeBase, }; @@ -167,7 +161,6 @@ fn should_read_pk_as_binary_with_limit( } const INDEX_TYPE_INVERTED: &str = "inverted"; const INDEX_TYPE_BLOOM: &str = "bloom filter"; -const INDEX_TYPE_VECTOR: &str = "vector"; macro_rules! handle_index_error { ($err:expr, $file_handle:expr, $index_type:expr) => { @@ -215,12 +208,6 @@ pub struct ParquetReaderBuilder { inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], fulltext_index_appliers: [Option; 2], - /// Vector index applier for KNN search. - #[cfg(feature = "vector_index")] - vector_index_applier: Option, - /// Over-fetched k for vector index scan. - #[cfg(feature = "vector_index")] - vector_index_k: Option, /// Expected metadata of the region while reading the SST. /// This is usually the latest metadata of the region. The reader use /// it get the correct column id of a column by name. @@ -262,10 +249,6 @@ impl ParquetReaderBuilder { inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], - #[cfg(feature = "vector_index")] - vector_index_applier: None, - #[cfg(feature = "vector_index")] - vector_index_k: None, expected_metadata: None, compaction: false, pre_filter_mode: PreFilterMode::All, @@ -352,19 +335,6 @@ impl ParquetReaderBuilder { self } - /// Attaches the vector index applier to the builder. - #[cfg(feature = "vector_index")] - #[must_use] - pub(crate) fn vector_index_applier( - mut self, - applier: Option, - k: Option, - ) -> Self { - self.vector_index_applier = applier; - self.vector_index_k = k; - self - } - /// Attaches the expected metadata to the builder. #[must_use] pub fn expected_metadata(mut self, expected_metadata: Option) -> Self { @@ -981,19 +951,6 @@ impl ParquetReaderBuilder { ) .await; } - #[cfg(feature = "vector_index")] - { - self.prune_row_groups_by_vector_index( - row_group_size, - num_row_groups, - &mut output, - metrics, - ) - .await; - if output.is_empty() { - return output; - } - } output } @@ -1259,49 +1216,6 @@ impl ParquetReaderBuilder { pruned } - /// Prunes row groups by vector index results. - #[cfg(feature = "vector_index")] - async fn prune_row_groups_by_vector_index( - &self, - row_group_size: usize, - num_row_groups: usize, - output: &mut RowGroupSelection, - metrics: &mut ReaderFilterMetrics, - ) { - let Some(applier) = &self.vector_index_applier else { - return; - }; - let Some(k) = self.vector_index_k else { - return; - }; - if !self.file_handle.meta_ref().vector_index_available() { - return; - } - - let file_size_hint = self.file_handle.meta_ref().index_file_size(); - let apply_res = applier - .apply_with_k(self.file_handle.index_id(), Some(file_size_hint), k) - .await; - let row_ids = match apply_res { - Ok(res) => res.row_offsets, - Err(err) => { - handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR); - return; - } - }; - - let selection = match vector_selection_from_offsets(row_ids, row_group_size, num_row_groups) - { - Ok(selection) => selection, - Err(err) => { - handle_index_error!(err, self.file_handle, INDEX_TYPE_VECTOR); - return; - } - }; - metrics.rows_vector_selected += selection.row_count(); - apply_selection_and_update_metrics(output, &selection, metrics, INDEX_TYPE_VECTOR); - } - async fn prune_row_groups_by_fulltext_bloom( &self, row_group_size: usize, @@ -1525,29 +1439,6 @@ fn apply_selection_and_update_metrics( *output = intersection; } -#[cfg(feature = "vector_index")] -fn vector_selection_from_offsets( - row_offsets: Vec, - row_group_size: usize, - num_row_groups: usize, -) -> Result { - let mut row_ids = BTreeSet::new(); - for offset in row_offsets { - let row_id = u32::try_from(offset).map_err(|_| { - ApplyVectorIndexSnafu { - reason: format!("Row offset {} exceeds u32::MAX", offset), - } - .build() - })?; - row_ids.insert(row_id); - } - Ok(RowGroupSelection::from_row_ids( - row_ids, - row_group_size, - num_row_groups, - )) -} - fn all_required_row_groups_searched( required_row_groups: &RowGroupSelection, cached_row_groups: &RowGroupSelection, @@ -1573,8 +1464,6 @@ pub(crate) struct ReaderFilterMetrics { pub(crate) rg_minmax_filtered: usize, /// Number of row groups filtered by bloom filter index. pub(crate) rg_bloom_filtered: usize, - /// Number of row groups filtered by vector index. - pub(crate) rg_vector_filtered: usize, /// Number of rows in row group before filtering. pub(crate) rows_total: usize, @@ -1584,10 +1473,6 @@ pub(crate) struct ReaderFilterMetrics { pub(crate) rows_inverted_filtered: usize, /// Number of rows in row group filtered by bloom filter index. pub(crate) rows_bloom_filtered: usize, - /// Number of rows filtered by vector index. - pub(crate) rows_vector_filtered: usize, - /// Number of rows selected by vector index. - pub(crate) rows_vector_selected: usize, /// Number of rows filtered by precise filter. pub(crate) rows_precise_filtered: usize, @@ -1633,14 +1518,11 @@ impl ReaderFilterMetrics { self.rg_inverted_filtered += other.rg_inverted_filtered; self.rg_minmax_filtered += other.rg_minmax_filtered; self.rg_bloom_filtered += other.rg_bloom_filtered; - self.rg_vector_filtered += other.rg_vector_filtered; self.rows_total += other.rows_total; self.rows_fulltext_filtered += other.rows_fulltext_filtered; self.rows_inverted_filtered += other.rows_inverted_filtered; self.rows_bloom_filtered += other.rows_bloom_filtered; - self.rows_vector_filtered += other.rows_vector_filtered; - self.rows_vector_selected += other.rows_vector_selected; self.rows_precise_filtered += other.rows_precise_filtered; self.fulltext_index_cache_hit += other.fulltext_index_cache_hit; @@ -1692,9 +1574,6 @@ impl ReaderFilterMetrics { READ_ROW_GROUPS_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rg_bloom_filtered as u64); - READ_ROW_GROUPS_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rg_vector_filtered as u64); PRECISE_FILTER_ROWS_TOTAL .with_label_values(&["parquet"]) @@ -1711,9 +1590,6 @@ impl ReaderFilterMetrics { READ_ROWS_IN_ROW_GROUP_TOTAL .with_label_values(&["bloom_filter_index_filtered"]) .inc_by(self.rows_bloom_filtered as u64); - READ_ROWS_IN_ROW_GROUP_TOTAL - .with_label_values(&["vector_index_filtered"]) - .inc_by(self.rows_vector_filtered as u64); } fn update_index_metrics(&mut self, index_type: &str, row_group_count: usize, row_count: usize) { @@ -1730,67 +1606,11 @@ impl ReaderFilterMetrics { self.rg_bloom_filtered += row_group_count; self.rows_bloom_filtered += row_count; } - INDEX_TYPE_VECTOR => { - self.rg_vector_filtered += row_group_count; - self.rows_vector_filtered += row_count; - } _ => {} } } } -#[cfg(all(test, feature = "vector_index"))] -mod vector_index_tests { - use super::*; - - #[test] - fn test_vector_selection_from_offsets() { - let row_group_size = 4; - let num_row_groups = 3; - let selection = - vector_selection_from_offsets(vec![0, 1, 5, 9], row_group_size, num_row_groups) - .unwrap(); - - assert_eq!(selection.row_group_count(), 3); - assert_eq!(selection.row_count(), 4); - assert!(selection.contains_non_empty_row_group(0)); - assert!(selection.contains_non_empty_row_group(1)); - assert!(selection.contains_non_empty_row_group(2)); - } - - #[test] - fn test_vector_selection_from_offsets_out_of_range() { - let row_group_size = 4; - let num_row_groups = 2; - let selection = vector_selection_from_offsets( - vec![0, 7, u64::from(u32::MAX) + 1], - row_group_size, - num_row_groups, - ); - assert!(selection.is_err()); - } - - #[test] - fn test_vector_selection_updates_metrics() { - let row_group_size = 4; - let total_rows = 8; - let mut output = RowGroupSelection::new(row_group_size, total_rows); - let selection = vector_selection_from_offsets(vec![1], row_group_size, 2).unwrap(); - let mut metrics = ReaderFilterMetrics::default(); - - apply_selection_and_update_metrics( - &mut output, - &selection, - &mut metrics, - INDEX_TYPE_VECTOR, - ); - - assert_eq!(metrics.rg_vector_filtered, 1); - assert_eq!(metrics.rows_vector_filtered, 7); - assert_eq!(output.row_count(), 1); - } -} - /// Metrics for parquet metadata cache operations. #[derive(Default, Clone, Copy)] pub struct MetadataCacheMetrics { diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 3eb45b1ab21..7b203e9af04 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -68,8 +68,6 @@ impl RegionWorkerLoop { inverted_index_config: self.config.inverted_index.clone(), fulltext_index_config: self.config.fulltext_index.clone(), bloom_filter_index_config: self.config.bloom_filter_index.clone(), - #[cfg(feature = "vector_index")] - vector_index_config: self.config.vector_index.clone(), index_options: version.options.index_options.clone(), intermediate_manager, puffin_manager, diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index e81d93b7781..c91bb477c78 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -11,7 +11,6 @@ workspace = true default = ["ai_functions"] ai_functions = ["common-function/ai_functions"] enterprise = ["sql/enterprise"] -vector_index = [] [dependencies] ahash.workspace = true diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 554628b3de2..5171688993c 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -528,188 +528,6 @@ fn dictionary_literals_in_values_keep_their_schema() { ); } -#[cfg(feature = "vector_index")] -mod vector_search_tests { - use std::sync::Arc; - - use common_function::function::Function; - use common_function::scalars::udf::create_udf; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, Signature, Volatility, col, lit}; - use datatypes::schema::{ColumnSchema, SchemaBuilder}; - use store_api::storage::ConcreteDataType; - use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMeta, TableType}; - use table::table::adapter::DfTableProviderAdapter; - use table::{Table, TableRef}; - - use super::*; - use crate::dist_plan::MergeScanLogicalPlan; - - struct TestVectorFunction { - name: &'static str, - signature: Signature, - } - - impl TestVectorFunction { - fn new(name: &'static str) -> Self { - Self { - name, - signature: Signature::any(2, Volatility::Immutable), - } - } - } - - impl std::fmt::Display for TestVectorFunction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl Function for TestVectorFunction { - fn name(&self) -> &str { - self.name - } - - fn return_type( - &self, - _input_types: &[datatypes::arrow::datatypes::DataType], - ) -> datafusion_common::Result { - Ok(datatypes::arrow::datatypes::DataType::Float32) - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn invoke_with_args( - &self, - _args: datafusion_expr::ScalarFunctionArgs, - ) -> datafusion_common::Result { - Err(datafusion_common::DataFusionError::Execution( - "test udf should not be invoked".to_string(), - )) - } - } - - fn build_vector_table(table_id: TableId) -> TableRef { - let schema = { - let columns = vec![ - ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true), - ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("v", ConcreteDataType::vector_datatype(2), false), - ]; - Arc::new( - SchemaBuilder::try_from_columns(columns) - .unwrap() - .build() - .unwrap(), - ) - }; - - let table_meta = TableMeta { - schema: schema.clone(), - primary_key_indices: vec![0], - value_indices: vec![2], - engine: "test_engine".to_string(), - next_column_id: 3, - options: Default::default(), - created_on: Default::default(), - updated_on: Default::default(), - partition_key_indices: vec![0], - column_ids: vec![0, 1, 2], - }; - - let table_info = TableInfoBuilder::default() - .table_id(table_id) - .name("t".to_string()) - .catalog_name(DEFAULT_CATALOG_NAME) - .schema_name(DEFAULT_SCHEMA_NAME) - .table_version(0) - .table_type(TableType::Base) - .meta(table_meta) - .build() - .unwrap(); - - let data_source = Arc::new(TestDataSource::new(schema)); - Arc::new(Table::new( - Arc::new(table_info), - FilterPushDownType::Unsupported, - data_source, - )) - } - - fn vector_distance_expr() -> Expr { - let udf = create_udf(Arc::new(TestVectorFunction::new("vec_l2sq_distance"))); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - col("v"), - lit(ScalarValue::Utf8(Some("[1.0, 2.0]".to_string()))), - ], - )) - } - - #[test] - fn vector_search_rewrite_keeps_sort_in_child_plan() { - init_default_ut_logging(); - let table = build_vector_table(0); - let table_source = Arc::new(DefaultTableSource::new(Arc::new( - DfTableProviderAdapter::new(table), - ))); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![vector_distance_expr().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); - - let plan_str = result.to_string(); - assert!(plan_str.contains("MergeSort: vec_l2sq_distance")); - assert!(plan_str.contains("Sort: vec_l2sq_distance")); - assert!(plan_str.contains(MergeScanLogicalPlan::name())); - } - - #[test] - fn vector_search_rewrite_with_filter_keeps_sort_in_child_plan() { - init_default_ut_logging(); - let table = build_vector_table(0); - let table_source = Arc::new(DefaultTableSource::new(Arc::new( - DfTableProviderAdapter::new(table), - ))); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .filter(col("k0").eq(lit("hello"))) - .unwrap() - .sort(vec![vector_distance_expr().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); - - let plan_str = result.to_string(); - assert!(plan_str.contains("MergeSort: vec_l2sq_distance")); - assert!(plan_str.contains("Sort: vec_l2sq_distance")); - assert!(plan_str.contains("Filter: t.k0 = Utf8(\"hello\")")); - assert!(plan_str.contains(MergeScanLogicalPlan::name())); - } -} - fn try_encode_decode_substrait(plan: &LogicalPlan, state: SessionState) { let sub_plan_bytes = substrait::DFLogicalSubstraitConvertor .encode(plan, crate::query_engine::DefaultSerializer) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index bc41516ba3c..a236118822b 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -16,17 +16,9 @@ use std::collections::HashSet; use std::sync::Arc; use common_function::aggrs::aggr_wrapper::{StateMergeHelper, is_all_aggr_exprs_steppable}; -#[cfg(feature = "vector_index")] -use common_function::scalars::vector::distance::{ - VEC_COS_DISTANCE, VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE, -}; use common_telemetry::debug; use datafusion::error::Result as DfResult; -#[cfg(feature = "vector_index")] -use datafusion_common::DataFusionError; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -#[cfg(feature = "vector_index")] -use datafusion_expr::Sort; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, @@ -37,37 +29,6 @@ use crate::dist_plan::MergeScanLogicalPlan; use crate::dist_plan::analyzer::AliasMapping; use crate::dist_plan::merge_sort::{MergeSortLogicalPlan, merge_sort_transformer}; -#[cfg(feature = "vector_index")] -fn is_vector_sort(sort: &Sort) -> bool { - if sort.expr.len() != 1 { - return false; - } - let sort_expr = &sort.expr[0].expr; - let Expr::ScalarFunction(func) = sort_expr else { - return false; - }; - matches!( - func.name().to_lowercase().as_str(), - VEC_L2SQ_DISTANCE | VEC_COS_DISTANCE | VEC_DOT_PRODUCT - ) -} - -#[cfg(feature = "vector_index")] -fn vector_sort_transformer(plan: &LogicalPlan) -> DfResult { - let LogicalPlan::Sort(sort) = plan else { - return Err(DataFusionError::Internal(format!( - "vector_sort_transformer expects Sort, got {plan}" - ))); - }; - Ok(TransformerAction { - extra_parent_plans: vec![ - MergeSortLogicalPlan::new(sort.input.clone(), sort.expr.clone(), sort.fetch) - .into_logical_plan(), - ], - new_child_plan: Some(LogicalPlan::Sort(sort.clone())), - }) -} - pub struct StepTransformAction { extra_parent_plans: Vec, new_child_plan: Option, @@ -196,12 +157,6 @@ impl Categorizer { // sort plan needs to consider column priority // Change Sort to MergeSort which assumes the input streams are already sorted hence can be more efficient. - #[cfg(feature = "vector_index")] - if is_vector_sort(_sort) { - return Ok(Commutativity::TransformedCommutative { - transformer: Some(Arc::new(vector_sort_transformer)), - }); - } Commutativity::ConditionalCommutative(Some(Arc::new(merge_sort_transformer))) } LogicalPlan::Join(_) => Commutativity::NonCommutative, diff --git a/src/query/src/dummy_catalog.rs b/src/query/src/dummy_catalog.rs index 2131423e3fe..5644194f1f9 100644 --- a/src/query/src/dummy_catalog.rs +++ b/src/query/src/dummy_catalog.rs @@ -37,9 +37,7 @@ use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; use store_api::metadata::RegionMetadataRef; use store_api::region_engine::RegionEngineRef; -use store_api::storage::{ - RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, VectorSearchRequest, -}; +use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use table::TableRef; use table::metadata::{TableId, TableInfoRef}; use table::table::adapter::{dictionary_encode_string_columns, supports_pk_dictionary_encoding}; @@ -281,14 +279,6 @@ impl DummyTableProvider { } } - pub fn with_vector_search_hint(&self, hint: VectorSearchRequest) { - self.scan_request.lock().unwrap().vector_search = Some(hint); - } - - pub fn get_vector_search_hint(&self) -> Option { - self.scan_request.lock().unwrap().vector_search.clone() - } - pub fn with_sequence(&self, sequence: u64) { self.scan_request.lock().unwrap().memtable_max_sequence = Some(sequence); } diff --git a/src/query/src/error.rs b/src/query/src/error.rs index a2946415570..919ef67b94f 100644 --- a/src/query/src/error.rs +++ b/src/query/src/error.rs @@ -331,13 +331,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to get VECTOR index options"))] - GetVectorIndexOptions { - source: datatypes::error::Error, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display( "Column schema mismatch in CTE {}, original: {:?}, expected: {:?}", cte_name, @@ -455,7 +448,6 @@ impl ErrorExt for Error { GetFulltextOptions { source, .. } | GetSkippingIndexOptions { source, .. } - | GetVectorIndexOptions { source, .. } | Datatypes { source, .. } => source.status_code(), } } @@ -481,7 +473,6 @@ impl ErrorExt for Error { TableMutation { source, .. } => source.retry_hint(), GetFulltextOptions { source, .. } | GetSkippingIndexOptions { source, .. } - | GetVectorIndexOptions { source, .. } | Datatypes { source, .. } => source.retry_hint(), _ => RetryHint::NonRetryable, } diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index edb73559139..4685889c998 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -31,10 +31,6 @@ use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector}; use crate::dummy_catalog::DummyTableProvider; -#[cfg(feature = "vector_index")] -mod vector_search; -#[cfg(feature = "vector_index")] -use vector_search::VectorSearchState; /// This rule will traverse the plan to collect necessary hints for leaf /// table scan node and set them in [`ScanRequest`]. Hints include: @@ -93,24 +89,18 @@ impl ScanHintRule { }; let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row; - #[cfg(feature = "vector_index")] - let has_vector_hint = rewriter.vector_search.need_rewrite(); - #[cfg(not(feature = "vector_index"))] - let has_vector_hint = false; - let has_hint = rewriter.order_expr.is_some() - || rewriter.ts_row_selector.is_some() - || use_last_row - || has_vector_hint; + let has_hint = + rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row; if !has_hint { return Ok(Transformed::no(LogicalPlan::TableScan(table_scan))); } // A provider can be used by several TableScan nodes. Fork its request // for every hinted use-site before applying hints, rather than mutating - // the shared catalog provider. This keeps order/vector/legacy hints + // the shared catalog provider. This keeps order and aggregate hints // local as well as the new LastRow hint. let adapter = original.clone_for_scan(); - Self::apply_hints(&adapter, rewriter, &table_scan, use_last_row); + Self::apply_hints(&adapter, rewriter); if use_last_row { // Apply the instant-derived hint after the aggregate hint. Both // select LastRow today, and this ordering preserves the existing @@ -156,35 +146,13 @@ impl ScanHintRule { true } - fn apply_hints( - adapter: &DummyTableProvider, - rewriter: &mut ScanHintRewriter, - table_scan: &datafusion_expr::logical_plan::TableScan, - use_last_row: bool, - ) { - #[cfg(not(feature = "vector_index"))] - let _ = (table_scan, use_last_row); + fn apply_hints(adapter: &DummyTableProvider, rewriter: &mut ScanHintRewriter) { if let Some(order_expr) = &rewriter.order_expr { Self::set_order_hint(adapter, order_expr); } if let Some((group_by_cols, order_by_col)) = &rewriter.ts_row_selector { Self::set_time_series_row_selector_hint(adapter, group_by_cols, order_by_col); } - #[cfg(feature = "vector_index")] - if use_last_row { - // LastRow and vector search are mutually exclusive for one scan: - // vector search would bypass the ordinary sort/limit path needed by - // the single-evaluation semantics. Still consume the queued hint so - // it cannot be applied to a later scan of the same table. - let _ = rewriter - .vector_search - .take_vector_request_from_dummy(adapter, &table_scan.table_name); - } else if let Some(vector_request) = rewriter - .vector_search - .take_vector_request_from_dummy(adapter, &table_scan.table_name) - { - adapter.with_vector_search_hint(vector_request); - } } fn set_order_hint(adapter: &DummyTableProvider, order_expr: &Vec) { @@ -307,8 +275,6 @@ struct ScanHintRewriter { ts_stack: Vec, Column)>>, inside_single_evaluation: bool, single_evaluation_stack: Vec, - #[cfg(feature = "vector_index")] - vector_search: VectorSearchState, } impl TreeNodeRewriter for ScanHintRewriter { @@ -355,44 +321,10 @@ impl TreeNodeRewriter for ScanHintRewriter { } } - #[cfg(feature = "vector_index")] - { - if let LogicalPlan::Limit(limit) = &node { - self.vector_search.on_limit_enter(limit); - } - if let LogicalPlan::Sort(sort) = &node { - self.vector_search.on_sort_enter(sort); - } - if is_branching_for_vector(&node) { - self.vector_search.on_branching_enter(); - } - if let LogicalPlan::Filter(filter) = &node { - self.vector_search.on_filter_enter(&filter.predicate); - } - if let LogicalPlan::TableScan(table_scan) = &node { - self.vector_search.on_table_scan(table_scan); - } - } - ScanHintRule::set_hints(node, self) } fn f_up(&mut self, node: LogicalPlan) -> Result> { - #[cfg(feature = "vector_index")] - { - match &node { - LogicalPlan::Limit(_) => self.vector_search.on_limit_exit(), - LogicalPlan::Sort(_) => self.vector_search.on_sort_exit(), - LogicalPlan::Filter(_) => self.vector_search.on_filter_exit(), - LogicalPlan::Subquery(_) | LogicalPlan::SubqueryAlias(_) - if is_branching_for_vector(&node) => - { - self.vector_search.on_branching_exit() - } - _ if node.inputs().len() > 1 => self.vector_search.on_branching_exit(), - _ => {} - } - } if let Some(previous) = self.order_stack.pop() { self.order_expr = previous; } @@ -510,43 +442,6 @@ impl ScanHintRewriter { } } -#[cfg(feature = "vector_index")] -fn is_branching_for_vector(node: &LogicalPlan) -> bool { - if node.inputs().len() > 1 { - return true; - } - - match node { - LogicalPlan::Subquery(subquery) => has_non_inlineable_ops(subquery.subquery.as_ref()), - LogicalPlan::SubqueryAlias(alias) => has_non_inlineable_ops(alias.input.as_ref()), - _ => false, - } -} - -#[cfg(feature = "vector_index")] -fn has_non_inlineable_ops(plan: &LogicalPlan) -> bool { - if matches!( - plan, - LogicalPlan::Limit(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Distinct(_) - | LogicalPlan::Aggregate(_) - | LogicalPlan::Window(_) - | LogicalPlan::Union(_) - | LogicalPlan::Join(_) - ) { - return true; - } - - for input in plan.inputs() { - if has_non_inlineable_ops(input) { - return true; - } - } - - false -} - #[cfg(test)] mod test { use std::collections::HashMap; diff --git a/src/query/src/optimizer/scan_hint/vector_search.rs b/src/query/src/optimizer/scan_hint/vector_search.rs deleted file mode 100644 index 90f45378ecd..00000000000 --- a/src/query/src/optimizer/scan_hint/vector_search.rs +++ /dev/null @@ -1,1066 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::collections::{HashMap, HashSet, VecDeque}; - -use common_function::scalars::vector::distance::{ - VEC_COS_DISTANCE, VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE, -}; -use common_telemetry::debug; -use datafusion_common::{ScalarValue, TableReference}; -use datafusion_expr::logical_plan::FetchType; -use datafusion_expr::utils::split_conjunction; -use datafusion_expr::{Expr, SortExpr}; -use datatypes::types::parse_string_to_vector_type_value; -use store_api::storage::{VectorDistanceMetric, VectorSearchRequest}; - -use crate::dummy_catalog::DummyTableProvider; - -/// Tracks vector search hints while traversing the logical plan. -/// -/// Vector search requests are emitted only when: -/// - ORDER BY uses a supported vector distance function and its direction matches the metric. -/// - A LIMIT (or Sort.fetch) is present to derive k. -/// - The hint stays within a single input chain (not across join/subquery branches). -/// - The target column is non-nullable, or an explicit IS NOT NULL filter exists. -/// -/// Known limitations: -/// - Dynamic overfetching is not implemented yet. When filters exist or ORDER BY includes -/// additional tie-breaker columns (e.g., ORDER BY distance, id), the current fixed k may -/// return incorrect results. A future improvement should dynamically adjust k based on -/// filter selectivity and secondary sort requirements. -/// - Hints only block at subquery boundaries when the subquery contains non-inlineable -/// operators (Limit/Sort/Distinct/Aggregate/Window). Simple subqueries without these -/// operators allow hints to propagate through. In distributed mode, the dist analyzer -/// may inline subqueries before this rule runs, further reducing isolation. -#[derive(Default)] -pub(crate) struct VectorSearchState { - current_distance: Option, - current_limit: Option, - distance_stack: Vec>, - limit_stack: Vec>, - non_null_columns: HashSet, - non_null_stack: Vec>, - vector_hints: HashMap>, -} - -#[derive(Clone)] -struct VectorDistanceInfo { - table_reference: Option, - column_name: String, - query_vector: Vec, - metric: VectorDistanceMetric, -} - -#[derive(Clone)] -struct VectorLimitInfo { - fetch: usize, - skip: usize, -} - -impl VectorLimitInfo { - fn k(&self) -> Option { - self.fetch.checked_add(self.skip) - } -} - -#[derive(Clone)] -struct VectorHintEntry { - distance: VectorDistanceInfo, - limit: VectorLimitInfo, - non_null_constraint: bool, -} - -#[derive(Clone, Hash, PartialEq, Eq)] -struct ColumnKey { - table_reference: Option, - column_name: String, -} - -impl VectorSearchState { - pub(crate) fn need_rewrite(&self) -> bool { - !self.vector_hints.is_empty() - } - - pub(crate) fn on_branching_enter(&mut self) { - // Clear per-branch state so hints are only derived within a single input chain. - self.distance_stack.push(self.current_distance.take()); - self.limit_stack.push(self.current_limit.take()); - self.non_null_stack - .push(std::mem::take(&mut self.non_null_columns)); - } - - pub(crate) fn on_branching_exit(&mut self) { - // Restore the prior chain state after leaving the branch. - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - if let Some(previous) = self.distance_stack.pop() { - self.current_distance = previous; - } - if let Some(previous) = self.non_null_stack.pop() { - self.non_null_columns = previous; - } - } - - pub(crate) fn on_limit_enter(&mut self, limit: &datafusion_expr::logical_plan::Limit) { - self.limit_stack.push(self.current_limit.take()); - self.current_limit = Self::extract_limit_info(limit); - } - - pub(crate) fn on_limit_exit(&mut self) { - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - } - - pub(crate) fn on_sort_enter(&mut self, sort: &datafusion_expr::logical_plan::Sort) { - // Distance is scoped to the nearest sort, while limit may be inherited from parents. - self.distance_stack.push(self.current_distance.take()); - self.limit_stack.push(self.current_limit.clone()); - let distance = Self::extract_distance_from_sort(sort); - self.current_distance = distance.clone(); - // Sort.fetch is a TopK limit, so we can infer k when no LIMIT is present. - if self.current_limit.is_none() { - self.current_limit = distance - .as_ref() - .and_then(|_| Self::extract_limit_from_sort(sort)); - } - } - - pub(crate) fn on_sort_exit(&mut self) { - if let Some(previous) = self.limit_stack.pop() { - self.current_limit = previous; - } - if let Some(previous) = self.distance_stack.pop() { - self.current_distance = previous; - } - } - - pub(crate) fn on_table_scan(&mut self, table_scan: &datafusion_expr::logical_plan::TableScan) { - self.record_vector_hint(table_scan); - } - - pub(crate) fn on_filter_enter(&mut self, predicate: &Expr) { - self.non_null_stack.push(self.non_null_columns.clone()); - for expr in split_conjunction(predicate) { - if let Expr::IsNotNull(inner) = expr - && let Expr::Column(col) = inner.as_ref() - { - self.non_null_columns.insert(ColumnKey { - table_reference: col.relation.clone(), - column_name: col.name.clone(), - }); - } - } - // TODO: detect non-null constraints from more complex predicates (casts/functions). - } - - pub(crate) fn on_filter_exit(&mut self) { - if let Some(previous) = self.non_null_stack.pop() { - self.non_null_columns = previous; - } - } - - pub(crate) fn take_vector_request_from_dummy( - &mut self, - provider: &DummyTableProvider, - table_name: &TableReference, - ) -> Option { - let hint = self.take_vector_hint(table_name)?; - self.build_vector_request_from_dummy(provider, table_name, &hint) - } - - fn build_vector_request_from_dummy( - &self, - provider: &DummyTableProvider, - table_name: &TableReference, - hint: &VectorHintEntry, - ) -> Option { - let info = &hint.distance; - let k = hint.limit.k()?; - - if let Some(ref hint_table) = info.table_reference - && table_name != hint_table - { - return None; - } - - let metadata = provider.region_metadata(); - let column = metadata.column_by_name(&info.column_name)?; - if column.column_schema.is_nullable() && !hint.non_null_constraint { - debug!( - "Skip vector hint: column '{}' is nullable without IS NOT NULL filter", - info.column_name - ); - return None; - } - - Some(VectorSearchRequest { - column_id: column.column_id, - query_vector: info.query_vector.clone(), - k, - metric: info.metric, - }) - } - - fn extract_distance_info(expr: &Expr) -> Option { - let Expr::ScalarFunction(func) = expr else { - return None; - }; - - let func_name = func.name().to_lowercase(); - let metric = match func_name.as_str() { - VEC_L2SQ_DISTANCE => VectorDistanceMetric::L2sq, - VEC_COS_DISTANCE => VectorDistanceMetric::Cosine, - VEC_DOT_PRODUCT => VectorDistanceMetric::InnerProduct, - _ => return None, - }; - - if func.args.len() != 2 { - return None; - } - - let (table_reference, column_name) = match &func.args[0] { - Expr::Column(col) => (col.relation.clone(), col.name.clone()), - _ => return None, - }; - - let query_vector = Self::extract_query_vector(&func.args[1])?; - - Some(VectorDistanceInfo { - table_reference, - column_name, - query_vector, - metric, - }) - } - - fn extract_distance_from_sort( - sort: &datafusion_expr::logical_plan::Sort, - ) -> Option { - if sort.expr.is_empty() { - debug!("Skip vector hint: Sort has no expressions"); - return None; - } - let sort_expr: &SortExpr = &sort.expr[0]; - let info = Self::extract_distance_info(&sort_expr.expr)?; - let expected_asc = info.metric != VectorDistanceMetric::InnerProduct; - if sort_expr.asc != expected_asc { - return None; - } - - if Self::tie_breakers_allowed(&sort.expr[1..], &info) { - Some(info) - } else { - if sort.expr.len() > 1 { - debug!( - "Skip vector hint: Sort has unsupported tie-breakers ({} expressions)", - sort.expr.len() - ); - } - None - } - } - - fn tie_breakers_allowed(sort_exprs: &[SortExpr], distance_info: &VectorDistanceInfo) -> bool { - if sort_exprs.is_empty() { - return true; - } - - sort_exprs.iter().all(|sort_expr| { - let Expr::Column(col) = &sort_expr.expr else { - return false; - }; - - match &distance_info.table_reference { - Some(table) => col.relation.as_ref() == Some(table), - None => col.relation.is_none(), - } - }) - } - - fn extract_limit_info(limit: &datafusion_expr::logical_plan::Limit) -> Option { - let fetch = match limit.get_fetch_type().ok()? { - FetchType::Literal(fetch) => fetch?, - FetchType::UnsupportedExpr => return None, - }; - let skip = match limit.get_skip_type().ok()? { - datafusion_expr::logical_plan::SkipType::Literal(skip) => skip, - datafusion_expr::logical_plan::SkipType::UnsupportedExpr => return None, - }; - Some(VectorLimitInfo { fetch, skip }) - } - - fn extract_limit_from_sort( - sort: &datafusion_expr::logical_plan::Sort, - ) -> Option { - let fetch = sort.fetch?; - Some(VectorLimitInfo { fetch, skip: 0 }) - } - - fn record_vector_hint(&mut self, table_scan: &datafusion_expr::logical_plan::TableScan) { - let Some(limit) = self.current_limit.as_ref() else { - return; - }; - let Some(distance) = self.current_distance.as_ref() else { - return; - }; - // Only emit hints when distance+limit are present and the table matches the sort target. - if let Some(ref hint_table) = distance.table_reference - && hint_table != &table_scan.table_name - { - return; - } - - let non_null_constraint = - self.is_column_non_null(&distance.table_reference, &distance.column_name); - self.vector_hints - .entry(table_scan.table_name.clone()) - .or_default() - .push_back(VectorHintEntry { - distance: distance.clone(), - limit: limit.clone(), - non_null_constraint, - }); - } - - fn take_vector_hint(&mut self, table_name: &TableReference) -> Option { - let hints = self.vector_hints.get_mut(table_name)?; - let hint = hints.pop_front(); - if hints.is_empty() { - self.vector_hints.remove(table_name); - } - hint - } - - fn is_column_non_null( - &self, - table_reference: &Option, - column_name: &str, - ) -> bool { - self.non_null_columns.contains(&ColumnKey { - table_reference: table_reference.clone(), - column_name: column_name.to_string(), - }) - } - - fn extract_query_vector(expr: &Expr) -> Option> { - match expr { - Expr::Literal(scalar, _) => match scalar { - ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => { - Self::parse_json_vector(s) - } - ScalarValue::Binary(Some(bytes)) | ScalarValue::LargeBinary(Some(bytes)) => { - Self::parse_binary_vector(bytes) - } - _ => None, - }, - _ => None, - } - } - - fn parse_json_vector(s: &str) -> Option> { - let trimmed = s.trim(); - if !trimmed.starts_with('[') || !trimmed.ends_with(']') { - return None; - } - parse_string_to_vector_type_value(trimmed, None) - .ok() - .and_then(|bytes| if bytes.is_empty() { None } else { Some(bytes) }) - .and_then(|bytes| Self::parse_binary_vector(&bytes)) - } - - fn parse_binary_vector(bytes: &[u8]) -> Option> { - if bytes.is_empty() || !bytes.len().is_multiple_of(4) { - return None; - } - Some( - bytes - .chunks(4) - .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])) - .collect(), - ) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use api::v1::SemanticType; - use common_function::function::Function; - use common_function::scalars::udf::create_udf; - use common_function::scalars::vector::distance::{VEC_DOT_PRODUCT, VEC_L2SQ_DISTANCE}; - use datafusion::datasource::DefaultTableSource; - use datafusion::logical_expr::ColumnarValue; - use datafusion_common::{Column, DataFusionError, Result, ScalarValue}; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::logical_plan::JoinType; - use datafusion_expr::{ - Expr, Extension, LogicalPlan, LogicalPlanBuilder, Signature, Subquery, Volatility, col, lit, - }; - use datafusion_optimizer::{OptimizerContext, OptimizerRule}; - use datatypes::schema::ColumnSchema; - use promql::extension_plan::InstantManipulate; - use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; - use store_api::storage::{ConcreteDataType, VectorDistanceMetric, VectorSearchRequest}; - - use super::VectorSearchState; - use crate::dummy_catalog::DummyTableProvider; - use crate::optimizer::scan_hint::ScanHintRule; - use crate::optimizer::test_util::MetaRegionEngine; - - struct TestVectorFunction { - name: &'static str, - signature: Signature, - } - - impl TestVectorFunction { - fn new(name: &'static str) -> Self { - Self { - name, - signature: Signature::any(2, Volatility::Immutable), - } - } - } - - impl std::fmt::Display for TestVectorFunction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl Function for TestVectorFunction { - fn name(&self) -> &str { - self.name - } - - fn return_type( - &self, - _input_types: &[datatypes::arrow::datatypes::DataType], - ) -> Result { - Ok(datatypes::arrow::datatypes::DataType::Float32) - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn invoke_with_args( - &self, - _args: datafusion_expr::ScalarFunctionArgs, - ) -> Result { - Err(DataFusionError::Execution( - "test udf should not be invoked".to_string(), - )) - } - } - - fn vec_distance_expr(function_name: &'static str) -> Expr { - vec_distance_expr_with_query(function_name, "[1.0, 2.0]") - } - - fn vec_distance_expr_with_query(function_name: &'static str, query_vector: &str) -> Expr { - let udf = create_udf(Arc::new(TestVectorFunction::new(function_name))); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - col("v"), - lit(ScalarValue::Utf8(Some(query_vector.to_string()))), - ], - )) - } - - fn scan_request_from_plan(plan: &LogicalPlan) -> store_api::storage::ScanRequest { - let mut request = None; - plan.apply_with_subqueries(|node| { - if let LogicalPlan::TableScan(scan) = node - && let Some(source) = scan.source.downcast_ref::() - && let Some(provider) = source.table_provider.downcast_ref::() - { - request = Some(provider.scan_request()); - } - Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) - }) - .unwrap(); - request.unwrap() - } - - fn vector_hint_from_plan(plan: &LogicalPlan) -> Option { - scan_request_from_plan(plan).vector_search - } - - fn build_dummy_provider(column_id: u32) -> Arc { - build_dummy_provider_with_nullable(column_id, false) - } - - fn build_dummy_provider_with_nullable( - column_id: u32, - nullable_vector: bool, - ) -> Arc { - let mut builder = RegionMetadataBuilder::new(0.into()); - builder - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new("k0", ConcreteDataType::string_datatype(), true), - semantic_type: SemanticType::Tag, - column_id: 1, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "ts", - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), - semantic_type: SemanticType::Timestamp, - column_id: 2, - }) - .push_column_metadata(ColumnMetadata { - column_schema: ColumnSchema::new( - "v", - ConcreteDataType::vector_datatype(2), - nullable_vector, - ), - semantic_type: SemanticType::Field, - column_id, - }) - .primary_key(vec![1]); - let metadata = Arc::new(builder.build().unwrap()); - let engine = Arc::new(MetaRegionEngine::with_metadata(metadata.clone())); - Arc::new(DummyTableProvider::new(0.into(), engine, metadata)) - } - - #[test] - fn test_parse_json_vector() { - assert_eq!( - VectorSearchState::parse_json_vector("[1.0, 2.0, 3.0]"), - Some(vec![1.0, 2.0, 3.0]) - ); - assert_eq!( - VectorSearchState::parse_json_vector("[1.5, -2.3, 0.0]"), - Some(vec![1.5, -2.3, 0.0]) - ); - assert_eq!(VectorSearchState::parse_json_vector("invalid"), None); - assert_eq!(VectorSearchState::parse_json_vector("[]"), None); - assert_eq!(VectorSearchState::parse_json_vector(""), None); - assert_eq!(VectorSearchState::parse_json_vector("["), None); - assert_eq!(VectorSearchState::parse_json_vector("[1.0, abc]"), None); - } - - #[test] - fn test_parse_binary_vector() { - let v1: f32 = 1.0; - let v2: f32 = 2.0; - let mut bytes = Vec::new(); - bytes.extend_from_slice(&v1.to_le_bytes()); - bytes.extend_from_slice(&v2.to_le_bytes()); - - let result = VectorSearchState::parse_binary_vector(&bytes); - assert_eq!(result, Some(vec![1.0, 2.0])); - - let result = VectorSearchState::parse_binary_vector(&[1, 2, 3]); - assert_eq!(result, None); - } - - #[test] - fn test_dummy_provider_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 10); - assert_eq!(hint.k, 5); - assert_eq!(hint.metric, VectorDistanceMetric::L2sq); - assert_eq!(hint.query_vector, vec![1.0, 2.0]); - } - - #[test] - fn test_vector_hint_is_applied_when_single_evaluation_limit_blocks_last_row() { - let provider = build_dummy_provider(10); - let source = Arc::new(DefaultTableSource::new(provider)); - let distance = vec_distance_expr(VEC_L2SQ_DISTANCE); - let scan = LogicalPlanBuilder::scan_with_filters("t", source, None, vec![]) - .unwrap() - .sort(vec![distance.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - let plan = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1000, - 1000, - 1000, - 1000, - 0, - "ts".to_string(), - vec![], - Some("v".to_string()), - scan, - )), - }); - - let rewritten = ScanHintRule - .rewrite(plan, &OptimizerContext::default()) - .unwrap() - .data; - let request = scan_request_from_plan(&rewritten); - assert_eq!(request.series_row_selector, None); - assert_eq!( - request.vector_search, - Some(VectorSearchRequest { - column_id: 10, - query_vector: vec![1.0, 2.0], - k: 5, - metric: VectorDistanceMetric::L2sq, - }) - ); - } - - #[test] - fn test_single_evaluation_limit_keeps_vector_hints_on_both_branches() { - let provider_a = build_dummy_provider(10); - let provider_b = build_dummy_provider(20); - let distance_a = vec_distance_expr_with_query(VEC_L2SQ_DISTANCE, "[1.0, 2.0]"); - let distance_b = vec_distance_expr_with_query(VEC_L2SQ_DISTANCE, "[3.0, 4.0]"); - - let branch_a = LogicalPlanBuilder::scan_with_filters( - "t", - Arc::new(DefaultTableSource::new(provider_a)), - None, - vec![], - ) - .unwrap() - .sort(vec![distance_a.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - let branch_a = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1000, - 1000, - 1000, - 1000, - 0, - "ts".to_string(), - vec![], - Some("v".to_string()), - branch_a, - )), - }); - - let branch_b = LogicalPlanBuilder::scan_with_filters( - "t", - Arc::new(DefaultTableSource::new(provider_b)), - None, - vec![], - ) - .unwrap() - .sort(vec![distance_b.sort(true, false)]) - .unwrap() - .limit(0, Some(9)) - .unwrap() - .build() - .unwrap(); - - let plan = LogicalPlanBuilder::from(branch_a) - .union(branch_b) - .unwrap() - .build() - .unwrap(); - let rewritten = ScanHintRule - .rewrite(plan, &OptimizerContext::default()) - .unwrap() - .data; - - let LogicalPlan::Union(union) = rewritten else { - panic!("expected union plan") - }; - let request_a = scan_request_from_plan(&union.inputs[0]); - let request_b = scan_request_from_plan(&union.inputs[1]); - - assert_eq!(request_a.series_row_selector, None); - assert_eq!( - request_a.vector_search, - Some(VectorSearchRequest { - column_id: 10, - query_vector: vec![1.0, 2.0], - k: 5, - metric: VectorDistanceMetric::L2sq, - }) - ); - assert_eq!(request_b.series_row_selector, None); - assert_eq!( - request_b.vector_search, - Some(VectorSearchRequest { - column_id: 20, - query_vector: vec![3.0, 4.0], - k: 9, - metric: VectorDistanceMetric::L2sq, - }) - ); - - // The catalog providers are not mutated; the requests above belong to - // the two rewritten scan use-sites. - assert_ne!(request_a.vector_search, request_b.vector_search); - } - - #[test] - fn test_limit_offset_for_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(5, Some(10)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 15); - } - - #[test] - fn test_inner_product_sort_direction() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_DOT_PRODUCT); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source.clone(), None, vec![]) - .unwrap() - .sort(vec![expr.clone().sort(true, false)]) - .unwrap() - .limit(0, Some(3)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - assert!(dummy_provider.get_vector_search_hint().is_none()); - - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(false, false)]) - .unwrap() - .limit(0, Some(3)) - .unwrap() - .build() - .unwrap(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.metric, VectorDistanceMetric::InnerProduct); - } - - #[test] - fn test_no_limit_clause() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(dummy_provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_nullable_vector_requires_is_not_null_filter() { - let dummy_provider = build_dummy_provider_with_nullable(10, true); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(dummy_provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_nullable_vector_with_is_not_null_filter() { - let dummy_provider = build_dummy_provider_with_nullable(10, true); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .filter(col("v").is_not_null()) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 10); - assert_eq!(hint.k, 5); - } - - #[test] - fn test_sort_fetch_for_vector_hint() { - let dummy_provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(dummy_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .sort_with_limit(vec![expr.sort(true, false)], Some(4)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 4); - } - - #[test] - fn test_limit_scoped_to_sort_branch() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - - let left = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let right = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .build() - .unwrap(); - - let join_keys: (Vec, Vec) = (vec![], vec![]); - let plan = LogicalPlanBuilder::from(left) - .join(right, JoinType::Inner, join_keys, None) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(t1_provider.get_vector_search_hint().is_none()); - assert!(t2_provider.get_vector_search_hint().is_none()); - } - - fn vec_distance_expr_qualified( - function_name: &'static str, - table_name: &str, - column_name: &str, - ) -> Expr { - use datafusion_common::Column; - - let udf = create_udf(Arc::new(TestVectorFunction::new(function_name))); - let qualified_col = Expr::Column(Column::new(Some(table_name.to_string()), column_name)); - Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::new(udf), - vec![ - qualified_col, - lit(ScalarValue::Utf8(Some("[1.0, 2.0]".to_string()))), - ], - )) - } - - #[test] - fn test_no_vector_hint_above_join() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - - let left = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let right = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let join_keys: (Vec, Vec) = (vec![], vec![]); - let join_plan = LogicalPlanBuilder::from(left) - .join(right, JoinType::Inner, join_keys, None) - .unwrap() - .build() - .unwrap(); - - let expr = vec_distance_expr_qualified(VEC_L2SQ_DISTANCE, "t1", "v"); - let plan = LogicalPlanBuilder::from(join_plan) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - assert!(t1_provider.get_vector_search_hint().is_none()); - assert!(t2_provider.get_vector_search_hint().is_none()); - } - - // Simple subqueries (without non-inlineable ops like Limit/Sort/Distinct/Aggregate/Window) - // allow hints to propagate through. See known limitations in VectorSearchState docs. - #[test] - fn test_simple_subquery_allows_hint_propagation() { - let provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(provider.clone())); - let scan_plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .build() - .unwrap(); - - let subquery = LogicalPlan::Subquery(Subquery { - subquery: Arc::new(scan_plan), - outer_ref_columns: vec![], - spans: Default::default(), - }); - - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::from(subquery) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let rewritten = ScanHintRule.rewrite(plan, &context).unwrap().data; - - // Hint propagates through simple subquery - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.k, 5); - } - - // Subqueries with non-inlineable ops (Limit/Sort/Distinct/Aggregate/Window) block hint propagation. - #[test] - fn test_subquery_with_limit_blocks_hint() { - let provider = build_dummy_provider(10); - let table_source = Arc::new(DefaultTableSource::new(provider.clone())); - let scan_plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) - .unwrap() - .limit(0, Some(100)) // non-inlineable op inside subquery - .unwrap() - .build() - .unwrap(); - - let subquery = LogicalPlan::Subquery(Subquery { - subquery: Arc::new(scan_plan), - outer_ref_columns: vec![], - spans: Default::default(), - }); - - let expr = vec_distance_expr(VEC_L2SQ_DISTANCE); - let plan = LogicalPlanBuilder::from(subquery) - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - let _ = ScanHintRule.rewrite(plan, &context).unwrap(); - - // Hint does NOT propagate through subquery with non-inlineable ops - assert!(provider.get_vector_search_hint().is_none()); - } - - #[test] - fn test_qualified_column_scopes_hint_to_correct_table() { - let t1_provider = build_dummy_provider(10); - let t2_provider = build_dummy_provider(20); - let t1_source = Arc::new(DefaultTableSource::new(t1_provider.clone())); - let t2_source = Arc::new(DefaultTableSource::new(t2_provider.clone())); - - let expr = vec_distance_expr_qualified(VEC_L2SQ_DISTANCE, "t2", "v"); - - let t1_plan = LogicalPlanBuilder::scan_with_filters("t1", t1_source, None, vec![]) - .unwrap() - .sort(vec![expr.clone().sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let t2_plan = LogicalPlanBuilder::scan_with_filters("t2", t2_source, None, vec![]) - .unwrap() - .sort(vec![expr.sort(true, false)]) - .unwrap() - .limit(0, Some(5)) - .unwrap() - .build() - .unwrap(); - - let context = OptimizerContext::default(); - - let _ = ScanHintRule.rewrite(t1_plan, &context).unwrap(); - assert!(t1_provider.get_vector_search_hint().is_none()); - - let rewritten = ScanHintRule.rewrite(t2_plan, &context).unwrap().data; - let hint = vector_hint_from_plan(&rewritten).unwrap(); - assert_eq!(hint.column_id, 20); - assert_eq!(hint.k, 5); - } -} diff --git a/src/query/src/sql/show_create_table.rs b/src/query/src/sql/show_create_table.rs index eb0392f15a1..c56e83d385a 100644 --- a/src/query/src/sql/show_create_table.rs +++ b/src/query/src/sql/show_create_table.rs @@ -23,11 +23,8 @@ use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, - COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, COMMENT_KEY, ColumnDefaultConstraint, ColumnSchema, - FulltextBackend, SchemaRef, + COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, + ColumnDefaultConstraint, ColumnSchema, FulltextBackend, SchemaRef, }; use datatypes::types::JsonFormat; use snafu::ResultExt; @@ -45,7 +42,7 @@ use table::requests::{ use crate::error::{ ConvertSqlTypeSnafu, ConvertSqlValueSnafu, GetFulltextOptionsSnafu, - GetSkippingIndexOptionsSnafu, GetVectorIndexOptionsSnafu, Result, SqlSnafu, + GetSkippingIndexOptionsSnafu, Result, SqlSnafu, }; /// Generates CREATE TABLE options from given table metadata and schema-level options. @@ -168,35 +165,6 @@ fn create_column(column_schema: &ColumnSchema, quote_style: char) -> Result StatusCode::Unsupported, PermissionDenied { .. } => StatusCode::PermissionDenied, - SetFulltextOption { .. } - | SetSkippingIndexOption { .. } - | SetVectorIndexOption { .. } => StatusCode::Unexpected, + SetFulltextOption { .. } | SetSkippingIndexOption { .. } => StatusCode::Unexpected, } } diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index bb73870b45f..d701415f3d2 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -48,7 +48,7 @@ use crate::parser::{FLOW, ParserContext}; use crate::parsers::tql_parser; use crate::parsers::utils::{ self, parse_with_options, validate_column_fulltext_create_option, - validate_column_skipping_index_create_option, validate_column_vector_index_create_option, + validate_column_skipping_index_create_option, }; use crate::statements::create::{ Column, ColumnExtensions, CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, @@ -66,7 +66,6 @@ pub const EXPIRE: &str = "EXPIRE"; pub const AFTER: &str = "AFTER"; pub const INVERTED: &str = "INVERTED"; pub const SKIPPING: &str = "SKIPPING"; -pub const VECTOR: &str = "VECTOR"; pub type RawIntervalExpr = String; @@ -1068,61 +1067,6 @@ impl<'a> ParserContext<'a> { is_index_declared |= true; } - // vector index - if let Token::Word(word) = parser.peek_token().token - && word.value.eq_ignore_ascii_case(VECTOR) - { - parser.next_token(); - // Consume `INDEX` keyword - ensure!( - parser.parse_keyword(Keyword::INDEX), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "expect INDEX after VECTOR keyword", - } - ); - - ensure!( - column_extensions.vector_index_options.is_none(), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "duplicated VECTOR INDEX option", - } - ); - - // Check that column is a vector type - let column_type = get_unalias_type(column_type); - let data_type = sql_data_type_to_concrete_data_type(&column_type)?; - ensure!( - matches!(data_type, ConcreteDataType::Vector(_)), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: "VECTOR INDEX only supports Vector type columns", - } - ); - - let options = parser - .parse_options(Keyword::WITH) - .context(error::SyntaxSnafu)? - .into_iter() - .map(parse_option_string) - .collect::>>()?; - - for (key, _) in options.iter() { - ensure!( - validate_column_vector_index_create_option(key), - InvalidColumnOptionSnafu { - name: column_name.to_string(), - msg: format!("invalid VECTOR INDEX option: {key}"), - } - ); - } - - let options = OptionMap::new(options); - column_extensions.vector_index_options = Some(options); - is_index_declared |= true; - } - Ok(is_index_declared) } @@ -3526,6 +3470,23 @@ CREATE TABLE log ( } } + #[test] + fn test_reject_vector_index() { + for index in ["VECTOR INDEX", "VECTOR INDEX WITH (metric = 'cosine')"] { + let sql = format!( + "CREATE TABLE vectors (ts TIMESTAMP TIME INDEX, embedding VECTOR(3) {index})" + ); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err() + ); + } + } + #[test] fn test_parse_interval_cast() { let s = "select '10s'::INTERVAL"; @@ -3534,174 +3495,4 @@ CREATE TABLE log ( .unwrap(); assert_eq!("SELECT '10 seconds'::INTERVAL", &stmts[0].to_string()); } - - #[test] - fn test_parse_create_table_vector_index_options() { - // Test basic vector index - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - if let Statement::CreateTable(c) = &result[0] { - c.columns.iter().for_each(|col| { - if col.name().value == "vec" { - assert!( - col.extensions - .vector_index_options - .as_ref() - .unwrap() - .is_empty() - ); - } - }); - } else { - panic!("should be create_table statement"); - } - - // Test vector index with options - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX WITH (metric='cosine', connectivity='32', expansion_add='256', expansion_search='128') -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - - if let Statement::CreateTable(c) = &result[0] { - c.columns.iter().for_each(|col| { - if col.name().value == "vec" { - let options = col.extensions.vector_index_options.as_ref().unwrap(); - assert_eq!(options.len(), 4); - assert_eq!(options.get("metric").unwrap(), "cosine"); - assert_eq!(options.get("connectivity").unwrap(), "32"); - assert_eq!(options.get("expansion_add").unwrap(), "256"); - assert_eq!(options.get("expansion_search").unwrap(), "128"); - } - }); - } else { - panic!("should be create_table statement"); - } - } - - #[test] - fn test_parse_create_table_vector_index_invalid_type() { - // Test vector index on non-vector type (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - col INT VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("VECTOR INDEX only supports Vector type columns") - ); - } - - #[test] - fn test_parse_create_table_vector_index_duplicate() { - // Test duplicate vector index (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX VECTOR INDEX, -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("duplicated VECTOR INDEX option") - ); - } - - #[test] - fn test_parse_create_table_vector_index_invalid_option() { - // Test invalid option key (should fail) - let sql = r" -CREATE TABLE vectors ( - ts TIMESTAMP TIME INDEX, - vec VECTOR(128) VECTOR INDEX WITH (metric='l2sq', invalid_option='foo') -)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("invalid VECTOR INDEX option") - ); - } - - #[test] - fn test_parse_column_extensions_vector_index() { - // Test vector index on vector type - { - let sql = "VECTOR INDEX WITH (metric = 'l2sq')"; - let dialect = GenericDialect {}; - let mut tokenizer = Tokenizer::new(&dialect, sql); - let tokens = tokenizer.tokenize().unwrap(); - let mut parser = Parser::new(&dialect).with_tokens(tokens); - let name = Ident::new("vec_col"); - let data_type = - DataType::Custom(vec![Ident::new("VECTOR")].into(), vec!["128".to_string()]); - // First, parse the vector type to set vector_options - let mut extensions = ColumnExtensions { - vector_options: Some(OptionMap::from([( - VECTOR_OPT_DIM.to_string(), - "128".to_string(), - )])), - ..Default::default() - }; - - let result = ParserContext::parse_column_extensions( - &mut parser, - &name, - &data_type, - &mut extensions, - ); - assert!(result.is_ok()); - assert!(extensions.vector_index_options.is_some()); - let vi_options = extensions.vector_index_options.unwrap(); - assert_eq!(vi_options.get("metric"), Some("l2sq")); - } - - // Test vector index on non-vector type (should fail) - { - let sql = "VECTOR INDEX"; - let dialect = GenericDialect {}; - let mut tokenizer = Tokenizer::new(&dialect, sql); - let tokens = tokenizer.tokenize().unwrap(); - let mut parser = Parser::new(&dialect).with_tokens(tokens); - let name = Ident::new("num_col"); - let data_type = DataType::Int(None); // Non-vector type - let mut extensions = ColumnExtensions::default(); - let result = ParserContext::parse_column_extensions( - &mut parser, - &name, - &data_type, - &mut extensions, - ); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("VECTOR INDEX only supports Vector type columns") - ); - } - } } diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index e91a2344643..afba437045c 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -32,9 +32,6 @@ use datatypes::schema::{ COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, }; use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; @@ -378,17 +375,6 @@ pub fn validate_column_skipping_index_create_option(key: &str) -> bool { .contains(&key) } -pub fn validate_column_vector_index_create_option(key: &str) -> bool { - [ - COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, - COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, - COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD, - COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, - ] - .contains(&key) -} - /// Convert an [`IntervalMonthDayNano`] to a [`Duration`]. #[cfg(feature = "enterprise")] pub fn convert_month_day_nano_to_duration( diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index e3f684bfcc2..2ea716678c7 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -56,7 +56,7 @@ use crate::ast::{ use crate::error::{ self, ConvertToGrpcDataTypeSnafu, ConvertValueSnafu, Result, SerializeColumnDefaultConstraintSnafu, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, - SetVectorIndexOptionSnafu, SqlCommonSnafu, + SqlCommonSnafu, }; use crate::statements::create::Column; pub use crate::statements::option_map::OptionMap; @@ -146,12 +146,6 @@ pub fn column_to_schema( .context(SetSkippingIndexOptionSnafu)?; } - if let Some(options) = column.extensions.build_vector_index_options()? { - column_schema = column_schema - .with_vector_index_options(&options) - .context(SetVectorIndexOptionSnafu)?; - } - column_schema.set_inverted_index(column.extensions.inverted_index_options.is_some()); let is_json2_column = if let SqlDataType::Custom(object_name, _) = column.data_type() { @@ -781,74 +775,4 @@ mod tests { assert_eq!(fulltext_options.analyzer, FulltextAnalyzer::English); assert!(fulltext_options.case_sensitive); } - - #[test] - fn test_column_to_schema_with_vector_index() { - use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; - - // Test with custom metric and parameters - let column = Column { - column_def: ColumnDef { - name: "embedding".into(), - data_type: SqlDataType::Custom( - vec![Ident::new(VECTOR_TYPE_NAME)].into(), - vec!["128".to_string()], - ), - options: vec![], - }, - extensions: ColumnExtensions { - vector_index_options: Some(OptionMap::from([ - ("metric".to_string(), "cosine".to_string()), - ("connectivity".to_string(), "32".to_string()), - ("expansion_add".to_string(), "200".to_string()), - ("expansion_search".to_string(), "100".to_string()), - ])), - ..Default::default() - }, - }; - - let column_schema = column_to_schema(&column, "ts", None).unwrap(); - assert_eq!("embedding", column_schema.name); - assert!(column_schema.is_vector_indexed()); - - let vector_options = column_schema.vector_index_options().unwrap().unwrap(); - assert_eq!(vector_options.engine, VectorIndexEngineType::Usearch); - assert_eq!(vector_options.metric, VectorDistanceMetric::Cosine); - assert_eq!(vector_options.connectivity, 32); - assert_eq!(vector_options.expansion_add, 200); - assert_eq!(vector_options.expansion_search, 100); - } - - #[test] - fn test_column_to_schema_with_vector_index_defaults() { - use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; - - // Test with default values (empty options map) - let column = Column { - column_def: ColumnDef { - name: "vec".into(), - data_type: SqlDataType::Custom( - vec![Ident::new(VECTOR_TYPE_NAME)].into(), - vec!["64".to_string()], - ), - options: vec![], - }, - extensions: ColumnExtensions { - vector_index_options: Some(OptionMap::default()), - ..Default::default() - }, - }; - - let column_schema = column_to_schema(&column, "ts", None).unwrap(); - assert_eq!("vec", column_schema.name); - assert!(column_schema.is_vector_indexed()); - - let vector_options = column_schema.vector_index_options().unwrap().unwrap(); - // Verify defaults - assert_eq!(vector_options.engine, VectorIndexEngineType::Usearch); - assert_eq!(vector_options.metric, VectorDistanceMetric::L2sq); - assert_eq!(vector_options.connectivity, 16); - assert_eq!(vector_options.expansion_add, 128); - assert_eq!(vector_options.expansion_search, 64); - } } diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 7c31790deaf..cb048a94b9e 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -18,10 +18,7 @@ use std::fmt::{Display, Formatter}; use common_catalog::consts::FILE_ENGINE; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; use datatypes::prelude::ConcreteDataType; -use datatypes::schema::{ - FulltextOptions, SkippingIndexOptions, VectorDistanceMetric, VectorIndexEngineType, - VectorIndexOptions, -}; +use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use itertools::Itertools; use serde::Serialize; use snafu::ResultExt; @@ -131,8 +128,6 @@ pub struct ColumnExtensions { /// /// Inverted index doesn't have options at present. There won't be any options in that map. pub inverted_index_options: Option, - /// Vector index options for HNSW-based vector similarity search. - pub vector_index_options: Option, /// JSON2-specific column options. pub json2_options: Option, } @@ -250,14 +245,6 @@ impl Display for Column { } } - if let Some(vector_index_options) = &self.extensions.vector_index_options { - if !vector_index_options.is_empty() { - let options = vector_index_options.kv_pairs(); - write!(f, " VECTOR INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " VECTOR INDEX")?; - } - } Ok(()) } } @@ -283,89 +270,6 @@ impl ColumnExtensions { )) } - pub fn build_vector_index_options(&self) -> Result> { - let Some(options) = self.vector_index_options.as_ref() else { - return Ok(None); - }; - - let options_map: HashMap = options.clone().into_map(); - let mut result = VectorIndexOptions::default(); - - if let Some(s) = options_map.get("engine") { - result.engine = s.parse::().map_err(|e| { - InvalidSqlSnafu { - msg: format!("invalid VECTOR INDEX engine: {e}"), - } - .build() - })?; - } - - if let Some(s) = options_map.get("metric") { - result.metric = s.parse::().map_err(|e| { - InvalidSqlSnafu { - msg: format!("invalid VECTOR INDEX metric: {e}"), - } - .build() - })?; - } - - if let Some(s) = options_map.get("connectivity") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX connectivity: {s}, expected positive integer" - ), - } - .build() - })?; - if !(2..=2048).contains(&value) { - return InvalidSqlSnafu { - msg: "VECTOR INDEX connectivity must be in the range [2, 2048].".to_string(), - } - .fail(); - } - result.connectivity = value; - } - - if let Some(s) = options_map.get("expansion_add") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX expansion_add: {s}, expected positive integer" - ), - } - .build() - })?; - if value == 0 { - return InvalidSqlSnafu { - msg: "VECTOR INDEX expansion_add must be greater than 0".to_string(), - } - .fail(); - } - result.expansion_add = value; - } - - if let Some(s) = options_map.get("expansion_search") { - let value = s.parse::().map_err(|_| { - InvalidSqlSnafu { - msg: format!( - "invalid VECTOR INDEX expansion_search: {s}, expected positive integer" - ), - } - .build() - })?; - if value == 0 { - return InvalidSqlSnafu { - msg: "VECTOR INDEX expansion_search must be greater than 0".to_string(), - } - .fail(); - } - result.expansion_search = value; - } - - Ok(Some(result)) - } - pub fn build_json_settings(&self) -> Result> { let Some(options) = &self.json2_options else { return Ok(None); @@ -1287,76 +1191,4 @@ AS SELECT number FROM numbers_input where number > 10"#, _ => unreachable!(), } } - - #[test] - fn test_vector_index_options_validation() { - use super::{ColumnExtensions, OptionMap}; - - // Test zero connectivity should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "connectivity".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("connectivity must be in the range [2, 2048]") - ); - - // Test zero expansion_add should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "expansion_add".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("expansion_add must be greater than 0") - ); - - // Test zero expansion_search should fail - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([( - "expansion_search".to_string(), - "0".to_string(), - )])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("expansion_search must be greater than 0") - ); - - // Test valid values should succeed - let extensions = ColumnExtensions { - vector_index_options: Some(OptionMap::from([ - ("connectivity".to_string(), "32".to_string()), - ("expansion_add".to_string(), "200".to_string()), - ("expansion_search".to_string(), "100".to_string()), - ])), - ..Default::default() - }; - let result = extensions.build_vector_index_options(); - assert!(result.is_ok()); - let options = result.unwrap().unwrap(); - assert_eq!(options.connectivity, 32); - assert_eq!(options.expansion_add, 200); - assert_eq!(options.expansion_search, 100); - } } diff --git a/src/store-api/src/metadata.rs b/src/store-api/src/metadata.rs index 21fc3295c44..6983bafb504 100644 --- a/src/store-api/src/metadata.rs +++ b/src/store-api/src/metadata.rs @@ -31,7 +31,7 @@ use common_macro::stack_trace_debug; use datatypes::arrow; use datatypes::arrow::datatypes::FieldRef; use datatypes::extension::json::json2_metadata_with_updated_settings; -use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef, VectorIndexOptions}; +use datatypes::schema::{ColumnSchema, FulltextOptions, Schema, SchemaRef}; use datatypes::types::TimestampType; use itertools::Itertools; use serde::de::Error; @@ -423,22 +423,6 @@ impl RegionMetadata { inverted_index } - /// Gets the column IDs that have vector indexes along with their options. - /// Returns a map from column ID to the vector index options. - pub fn vector_indexed_column_ids(&self) -> HashMap { - self.column_metadatas - .iter() - .filter_map(|column| { - column - .column_schema - .vector_index_options() - .ok() - .flatten() - .map(|options| (column.column_id, options)) - }) - .collect() - } - /// Checks whether the metadata is valid. fn validate(&self) -> Result<()> { // Id to name. diff --git a/src/store-api/src/storage.rs b/src/store-api/src/storage.rs index cae2439168a..e0761f8a006 100644 --- a/src/store-api/src/storage.rs +++ b/src/store-api/src/storage.rs @@ -29,8 +29,5 @@ pub use datatypes::schema::{ pub use self::descriptors::*; pub use self::file::{FileId, FileRef, FileRefsManifest, GcReport, IndexVersion, ParseIdError}; pub use self::projection::NestedPath; -pub use self::requests::{ - ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector, VectorDistanceMetric, - VectorIndexEngine, VectorIndexEngineType, VectorSearchMatches, VectorSearchRequest, -}; +pub use self::requests::{ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; pub use self::types::{SequenceNumber, SequenceRange}; diff --git a/src/store-api/src/storage/requests.rs b/src/store-api/src/storage/requests.rs index b35a1907d43..3c9d4951d4b 100644 --- a/src/store-api/src/storage/requests.rs +++ b/src/store-api/src/storage/requests.rs @@ -15,68 +15,13 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter}; -use common_error::ext::BoxedError; use common_recordbatch::OrderOption; use datafusion_expr::expr::Expr; -// Re-export vector types from datatypes to avoid duplication -pub use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType}; use datatypes::types::json_type::JsonNativeType; use itertools::Itertools; use strum::Display; -use crate::storage::{ColumnId, SequenceNumber}; - -/// A hint for KNN vector search. -#[derive(Debug, Clone, PartialEq)] -pub struct VectorSearchRequest { - /// Column ID of the vector column to search. - pub column_id: ColumnId, - /// The query vector to search for. - pub query_vector: Vec, - /// Number of nearest neighbors to return. - pub k: usize, - /// Distance metric to use (matches the index metric). - pub metric: VectorDistanceMetric, -} - -/// Search results from vector index. -#[derive(Debug, Clone, PartialEq)] -pub struct VectorSearchMatches { - /// Keys (row offsets in the index). - pub keys: Vec, - /// Distances from the query vector. - pub distances: Vec, -} - -/// Trait for vector index engines (HNSW implementations). -/// -/// This trait defines the interface for pluggable vector index engines. -/// Implementations (e.g., UsearchEngine) are provided by storage engines like mito2. -pub trait VectorIndexEngine: Send + Sync { - /// Adds a vector with the given key. - fn add(&mut self, key: u64, vector: &[f32]) -> Result<(), BoxedError>; - - /// Searches for k nearest neighbors. - fn search(&self, query: &[f32], k: usize) -> Result; - - /// Returns the serialized length. - fn serialized_length(&self) -> usize; - - /// Serializes the index to a buffer. - fn save_to_buffer(&self, buffer: &mut [u8]) -> Result<(), BoxedError>; - - /// Reserves capacity for vectors. - fn reserve(&mut self, capacity: usize) -> Result<(), BoxedError>; - - /// Returns current size (number of vectors). - fn size(&self) -> usize; - - /// Returns current capacity. - fn capacity(&self) -> usize; - - /// Returns memory usage in bytes. - fn memory_usage(&self) -> usize; -} +use crate::storage::SequenceNumber; /// A hint on how to select rows from a time-series. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display)] @@ -144,9 +89,6 @@ pub struct ScanRequest { pub exact_sequence_range: bool, /// Optional hint for the distribution of time-series data. pub distribution: Option, - /// Optional hint for KNN vector search. When set, the scan should use - /// vector index to find the k nearest neighbors. - pub vector_search: Option, /// Optional hint from query-driven JSON type concretization. pub json_type_hint: HashMap, /// Whether Mito should keep string primary-key columns dictionary encoded in its output. @@ -249,16 +191,6 @@ impl Display for ScanRequest { if let Some(distribution) = &self.distribution { write!(f, "{}distribution: {}", delimiter.as_str(), distribution)?; } - if let Some(vector_search) = &self.vector_search { - write!( - f, - "{}vector_search: column_id={}, k={}, metric={}", - delimiter.as_str(), - vector_search.column_id, - vector_search.k, - vector_search.metric - )?; - } if !self.json_type_hint.is_empty() { write!( f, diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 16c187d38b6..5bbd8c3c6d0 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -29,7 +29,7 @@ use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datatypes::types::json_type::JsonNativeType; use snafu::ResultExt; -use store_api::storage::{ScanRequest, VectorSearchRequest}; +use store_api::storage::ScanRequest; use crate::error::TablesRecordBatchSnafu; use crate::table::{TableRef, TableType}; @@ -61,18 +61,10 @@ impl DfTableProviderAdapter { self.scan_req.lock().unwrap().output_ordering = Some(order_opts.to_vec()); } - pub fn with_vector_search_hint(&self, hint: VectorSearchRequest) { - self.scan_req.lock().unwrap().vector_search = Some(hint); - } - pub fn with_json_type_hint(&self, hint: std::collections::HashMap) { self.scan_req.lock().unwrap().json_type_hint = hint; } - pub fn get_vector_search_hint(&self) -> Option { - self.scan_req.lock().unwrap().vector_search.clone() - } - #[cfg(feature = "testing")] pub fn get_scan_req(&self) -> ScanRequest { self.scan_req.lock().unwrap().clone() diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index a3d694f758b..d281f3cecc8 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -11,7 +11,6 @@ path = "tests/main.rs" [features] dashboard = ["servers/dashboard"] -vector_index = [] enterprise = ["frontend/enterprise", "meta-srv/enterprise", "common-meta/enterprise"] [lints] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index a4841afad93..fcafbc0ce52 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2394,19 +2394,6 @@ providers = []"#, ) }; - let vector_index_config = if cfg!(feature = "vector_index") { - r#" -[region_engine.mito.vector_index] -create_on_flush = "auto" -create_on_compaction = "auto" -apply_on_query = "auto" -mem_threshold_on_create = "auto" - -"# - } else { - "\n" - }; - let expected_toml_str = format!( r#" enable_telemetry = true @@ -2611,7 +2598,8 @@ create_on_flush = "auto" create_on_compaction = "auto" apply_on_query = "auto" mem_threshold_on_create = "auto" -{vector_index_config}[region_engine.mito.gc] + +[region_engine.mito.gc] enable = false lingering_time = "1h" unknown_file_lingering_time = "1day" diff --git a/tests/cases/standalone/common/function/vector/vector_index.result b/tests/cases/standalone/common/function/vector/vector_index.result deleted file mode 100644 index 44c6e447054..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index.result +++ /dev/null @@ -1,602 +0,0 @@ --- Test vector index creation and KNN search --- ============================================ --- Part 1: Basic L2 squared distance tests --- ============================================ --- Create a table with vector column and L2sq vector index -CREATE TABLE vectors_l2sq ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert test vectors -INSERT INTO vectors_l2sq VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.0, 0.0, 1.0]'), - (5, '2024-01-01 00:00:04', '[1.0, 1.0, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.0, 1.0, 1.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.0, 0.0, 1.0, 1.0]'), - (8, '2024-01-01 00:00:07', '[1.0, 0.0, 0.0, 1.0]'); - -Affected Rows: 8 - --- Query BEFORE flush (memtable search) --- Expected: vec_id=1 (distance=0), vec_id=5 (distance=1), vec_id=8 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 5 | 1.0 | -| 8 | 1.0 | -+--------+----------+ - --- Flush to create SST files with vector index -ADMIN FLUSH_TABLE('vectors_l2sq'); - -+-----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_l2sq') | -+-----------------------------------+ -| 0 | -+-----------------------------------+ - --- Query AFTER flush (SST index search) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 5 | 1.0 | -| 8 | 1.0 | -+--------+----------+ - --- Query with different target vector --- Expected: vec_id=6 (distance=0), vec_id=2 (distance=1), vec_id=3 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 6 | 0.0 | -| 2 | 1.0 | -| 3 | 1.0 | -+--------+----------+ - -DROP TABLE vectors_l2sq; - -Affected Rows: 0 - --- ============================================ --- Part 2: Cosine distance tests --- ============================================ -CREATE TABLE vectors_cosine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert vectors with different magnitudes but same/different directions -INSERT INTO vectors_cosine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with cosine distance --- vec_id=1 and vec_id=2 should have same cosine distance (0) since they point same direction -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+------------+ -| vec_id | distance | -+--------+------------+ -| 1 | 0.0 | -| 2 | 0.0 | -| 4 | 0.29289323 | -+--------+------------+ - -ADMIN FLUSH_TABLE('vectors_cosine'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_cosine') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with cosine distance -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+------------+ -| vec_id | distance | -+--------+------------+ -| 1 | 0.0 | -| 2 | 0.0 | -| 4 | 0.29289323 | -+--------+------------+ - -DROP TABLE vectors_cosine; - -Affected Rows: 0 - --- ============================================ --- Part 3: Dot product (inner product) tests --- ============================================ -CREATE TABLE vectors_dot ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'dot'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_dot VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with dot product --- Larger dot product means more similar, so we use negative for ordering --- vec_id=2 should be best (dot=2), vec_id=1 and vec_id=4 have dot=1 -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | dot_product | -+--------+-------------+ -| 2 | 2.0 | -| 1 | 1.0 | -| 4 | 1.0 | -+--------+-------------+ - -ADMIN FLUSH_TABLE('vectors_dot'); - -+----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_dot') | -+----------------------------------+ -| 0 | -+----------------------------------+ - --- SST index search with dot product -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | dot_product | -+--------+-------------+ -| 2 | 2.0 | -| 1 | 1.0 | -| 4 | 1.0 | -+--------+-------------+ - -DROP TABLE vectors_dot; - -Affected Rows: 0 - --- ============================================ --- Part 4: NULL vector handling tests --- ============================================ -CREATE TABLE vectors_null ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert vectors with some NULLs -INSERT INTO vectors_null VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', NULL), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', NULL), - (5, '2024-01-01 00:00:04', '[0.0, 0.0, 1.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search should skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 3 | 2.0 | -| 5 | 2.0 | -+--------+----------+ - -ADMIN FLUSH_TABLE('vectors_null'); - -+-----------------------------------+ -| ADMIN FLUSH_TABLE('vectors_null') | -+-----------------------------------+ -| 0 | -+-----------------------------------+ - --- SST index search should also skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+----------+ -| vec_id | distance | -+--------+----------+ -| 1 | 0.0 | -| 3 | 2.0 | -| 5 | 2.0 | -+--------+----------+ - -DROP TABLE vectors_null; - -Affected Rows: 0 - --- ============================================ --- Part 5: Mixed memtable + SST search tests --- ============================================ -CREATE TABLE vectors_mixed ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert first batch and flush to SST -INSERT INTO vectors_mixed VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'); - -Affected Rows: 3 - -ADMIN FLUSH_TABLE('vectors_mixed'); - -+------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_mixed') | -+------------------------------------+ -| 0 | -+------------------------------------+ - --- Insert second batch (stays in memtable) -INSERT INTO vectors_mixed VALUES - (4, '2024-01-01 00:00:03', '[0.5, 0.5, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.9, 0.1, 0.0, 0.0]'); - -Affected Rows: 2 - --- Query should search both SST (vec_id 1,2,3) and memtable (vec_id 4,5) --- Target: [1.0, 0.0, 0.0, 0.0] --- Expected: vec_id=1 (dist=0), vec_id=5 (dist=0.02), vec_id=4 (dist=0.5) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_mixed -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -+--------+-------------+ -| vec_id | distance | -+--------+-------------+ -| 1 | 0.0 | -| 5 | 0.020000005 | -| 4 | 0.5 | -+--------+-------------+ - -DROP TABLE vectors_mixed; - -Affected Rows: 0 - --- ============================================ --- Part 6: KNN with WHERE clause tests --- ============================================ -CREATE TABLE vectors_filter ( - vec_id INT, - "category" STRING, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_filter VALUES - (1, 'A', '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, 'B', '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, 'A', '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, 'B', '2024-01-01 00:00:03', '[0.1, 0.9, 0.0, 0.0]'), - (5, 'A', '2024-01-01 00:00:04', '[0.5, 0.5, 0.0, 0.0]'); - -Affected Rows: 5 - --- Memtable search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+----------+----------+ -| vec_id | category | distance | -+--------+----------+----------+ -| 1 | A | 0.0 | -| 5 | A | 0.5 | -+--------+----------+----------+ - -ADMIN FLUSH_TABLE('vectors_filter'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_filter') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+----------+----------+ -| vec_id | category | distance | -+--------+----------+----------+ -| 1 | A | 0.0 | -| 5 | A | 0.5 | -+--------+----------+----------+ - --- Filter with time range -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE ts >= '2024-01-01 00:00:02' AND ts <= '2024-01-01 00:00:04' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+-----------+ -| vec_id | distance | -+--------+-----------+ -| 5 | 0.5 | -| 4 | 1.6199999 | -+--------+-----------+ - -DROP TABLE vectors_filter; - -Affected Rows: 0 - --- ============================================ --- Part 7: Higher dimension vectors --- ============================================ -CREATE TABLE vectors_high_dim ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(128) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert high-dimensional vectors (simplified: first few elements differ) -INSERT INTO vectors_high_dim VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'); - -Affected Rows: 3 - --- Memtable search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -ADMIN FLUSH_TABLE('vectors_high_dim'); - -+---------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_high_dim') | -+---------------------------------------+ -| 0 | -+---------------------------------------+ - --- SST index search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -DROP TABLE vectors_high_dim; - -Affected Rows: 0 - --- ============================================ --- Part 8: Different k values (LIMIT) --- ============================================ -CREATE TABLE vectors_k ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_k VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.8, 0.2, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.7, 0.3, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.6, 0.4, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.5, 0.5, 0.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.4, 0.6, 0.0, 0.0]'), - (8, '2024-01-01 00:00:07', '[0.3, 0.7, 0.0, 0.0]'), - (9, '2024-01-01 00:00:08', '[0.2, 0.8, 0.0, 0.0]'), - (10, '2024-01-01 00:00:09', '[0.1, 0.9, 0.0, 0.0]'); - -Affected Rows: 10 - -ADMIN FLUSH_TABLE('vectors_k'); - -+--------------------------------+ -| ADMIN FLUSH_TABLE('vectors_k') | -+--------------------------------+ -| 0 | -+--------------------------------+ - --- k=1 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 1; - -+--------+ -| vec_id | -+--------+ -| 1 | -+--------+ - --- k=5 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 5; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 2 | -| 3 | -| 4 | -| 5 | -+--------+ - --- k=10 (all vectors) -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 10; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 2 | -| 3 | -| 4 | -| 5 | -| 6 | -| 7 | -| 8 | -| 9 | -| 10 | -+--------+ - -DROP TABLE vectors_k; - -Affected Rows: 0 - --- ============================================ --- Part 9: Engine parameter tests --- ============================================ --- Create table with explicit engine parameter (usearch is default) -CREATE TABLE vectors_engine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (engine = 'usearch', metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - --- Insert test vectors -INSERT INTO vectors_engine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0]'); - -Affected Rows: 3 - --- Memtable search -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -ADMIN FLUSH_TABLE('vectors_engine'); - -+-------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_engine') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SST index search with usearch engine -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -+--------+ -| vec_id | -+--------+ -| 1 | -| 3 | -+--------+ - -DROP TABLE vectors_engine; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index.sql b/tests/cases/standalone/common/function/vector/vector_index.sql deleted file mode 100644 index bbcf02d7f30..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index.sql +++ /dev/null @@ -1,341 +0,0 @@ --- Test vector index creation and KNN search - --- ============================================ --- Part 1: Basic L2 squared distance tests --- ============================================ - --- Create a table with vector column and L2sq vector index -CREATE TABLE vectors_l2sq ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert test vectors -INSERT INTO vectors_l2sq VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.0, 0.0, 1.0]'), - (5, '2024-01-01 00:00:04', '[1.0, 1.0, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.0, 1.0, 1.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.0, 0.0, 1.0, 1.0]'), - (8, '2024-01-01 00:00:07', '[1.0, 0.0, 0.0, 1.0]'); - --- Query BEFORE flush (memtable search) --- Expected: vec_id=1 (distance=0), vec_id=5 (distance=1), vec_id=8 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - --- Flush to create SST files with vector index -ADMIN FLUSH_TABLE('vectors_l2sq'); - --- Query AFTER flush (SST index search) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - --- Query with different target vector --- Expected: vec_id=6 (distance=0), vec_id=2 (distance=1), vec_id=3 (distance=1) -SELECT vec_id, vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]') as distance -FROM vectors_l2sq -ORDER BY vec_l2sq_distance(embedding, '[0.0, 1.0, 1.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_l2sq; - --- ============================================ --- Part 2: Cosine distance tests --- ============================================ - -CREATE TABLE vectors_cosine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - --- Insert vectors with different magnitudes but same/different directions -INSERT INTO vectors_cosine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - --- Memtable search with cosine distance --- vec_id=1 and vec_id=2 should have same cosine distance (0) since they point same direction -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_cosine'); - --- SST index search with cosine distance -SELECT vec_id, vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_cosine -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_cosine; - --- ============================================ --- Part 3: Dot product (inner product) tests --- ============================================ - -CREATE TABLE vectors_dot ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'dot'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_dot VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[2.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[1.0, 1.0, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[-1.0, 0.0, 0.0, 0.0]'); - --- Memtable search with dot product --- Larger dot product means more similar, so we use negative for ordering --- vec_id=2 should be best (dot=2), vec_id=1 and vec_id=4 have dot=1 -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_dot'); - --- SST index search with dot product -SELECT vec_id, vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') as dot_product -FROM vectors_dot -ORDER BY vec_dot_product(embedding, '[1.0, 0.0, 0.0, 0.0]') DESC, vec_id -LIMIT 3; - -DROP TABLE vectors_dot; - --- ============================================ --- Part 4: NULL vector handling tests --- ============================================ - -CREATE TABLE vectors_null ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert vectors with some NULLs -INSERT INTO vectors_null VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', NULL), - (3, '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', NULL), - (5, '2024-01-01 00:00:04', '[0.0, 0.0, 1.0, 0.0]'); - --- Memtable search should skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -ADMIN FLUSH_TABLE('vectors_null'); - --- SST index search should also skip NULL vectors -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_null -WHERE embedding IS NOT NULL -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_null; - --- ============================================ --- Part 5: Mixed memtable + SST search tests --- ============================================ - -CREATE TABLE vectors_mixed ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert first batch and flush to SST -INSERT INTO vectors_mixed VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 0.0, 1.0, 0.0]'); - -ADMIN FLUSH_TABLE('vectors_mixed'); - --- Insert second batch (stays in memtable) -INSERT INTO vectors_mixed VALUES - (4, '2024-01-01 00:00:03', '[0.5, 0.5, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.9, 0.1, 0.0, 0.0]'); - --- Query should search both SST (vec_id 1,2,3) and memtable (vec_id 4,5) --- Target: [1.0, 0.0, 0.0, 0.0] --- Expected: vec_id=1 (dist=0), vec_id=5 (dist=0.02), vec_id=4 (dist=0.5) -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_mixed -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 3; - -DROP TABLE vectors_mixed; - --- ============================================ --- Part 6: KNN with WHERE clause tests --- ============================================ - -CREATE TABLE vectors_filter ( - vec_id INT, - "category" STRING, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_filter VALUES - (1, 'A', '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, 'B', '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, 'A', '2024-01-01 00:00:02', '[0.0, 1.0, 0.0, 0.0]'), - (4, 'B', '2024-01-01 00:00:03', '[0.1, 0.9, 0.0, 0.0]'), - (5, 'A', '2024-01-01 00:00:04', '[0.5, 0.5, 0.0, 0.0]'); - --- Memtable search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_filter'); - --- SST index search with filter -SELECT vec_id, category, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE category = 'A' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - --- Filter with time range -SELECT vec_id, vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]') as distance -FROM vectors_filter -WHERE ts >= '2024-01-01 00:00:02' AND ts <= '2024-01-01 00:00:04' -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_filter; - --- ============================================ --- Part 7: Higher dimension vectors --- ============================================ - -CREATE TABLE vectors_high_dim ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(128) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert high-dimensional vectors (simplified: first few elements differ) -INSERT INTO vectors_high_dim VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'); - --- Memtable search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_high_dim'); - --- SST index search -SELECT vec_id FROM vectors_high_dim -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_high_dim; - --- ============================================ --- Part 8: Different k values (LIMIT) --- ============================================ - -CREATE TABLE vectors_k ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_k VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.1, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.8, 0.2, 0.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.7, 0.3, 0.0, 0.0]'), - (5, '2024-01-01 00:00:04', '[0.6, 0.4, 0.0, 0.0]'), - (6, '2024-01-01 00:00:05', '[0.5, 0.5, 0.0, 0.0]'), - (7, '2024-01-01 00:00:06', '[0.4, 0.6, 0.0, 0.0]'), - (8, '2024-01-01 00:00:07', '[0.3, 0.7, 0.0, 0.0]'), - (9, '2024-01-01 00:00:08', '[0.2, 0.8, 0.0, 0.0]'), - (10, '2024-01-01 00:00:09', '[0.1, 0.9, 0.0, 0.0]'); - -ADMIN FLUSH_TABLE('vectors_k'); - --- k=1 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 1; - --- k=5 -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 5; - --- k=10 (all vectors) -SELECT vec_id FROM vectors_k -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 10; - -DROP TABLE vectors_k; - --- ============================================ --- Part 9: Engine parameter tests --- ============================================ - --- Create table with explicit engine parameter (usearch is default) -CREATE TABLE vectors_engine ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(4) NOT NULL VECTOR INDEX WITH (engine = 'usearch', metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - --- Insert test vectors -INSERT INTO vectors_engine VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0, 0.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0, 0.0, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.5, 0.5, 0.0, 0.0]'); - --- Memtable search -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -ADMIN FLUSH_TABLE('vectors_engine'); - --- SST index search with usearch engine -SELECT vec_id FROM vectors_engine -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0, 0.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_engine; diff --git a/tests/cases/standalone/common/function/vector/vector_index_explain.result b/tests/cases/standalone/common/function/vector/vector_index_explain.result deleted file mode 100644 index 60f01ce157e..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_explain.result +++ /dev/null @@ -1,272 +0,0 @@ --- Vector index explain analyze coverage --- ============================================ --- Part 1: Single table KNN explain --- ============================================ -CREATE TABLE vectors_explain ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_explain'); - -+--------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain') | -+--------------------------------------+ -| 0 | -+--------------------------------------+ - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":0, "files":1, "file_ranges":1}, "projection": ["vec_id", "embedding"]REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain; - -Affected Rows: 0 - --- ============================================ --- Part 2: Join with vector order/limit --- ============================================ -CREATE TABLE vectors_explain_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_explain_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -INSERT INTO vectors_explain_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -Affected Rows: 2 - -ADMIN FLUSH_TABLE('vectors_explain_left'); - -+-------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain_left') | -+-------------------------------------------+ -| 0 | -+-------------------------------------------+ - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT l.vec_id -FROM vectors_explain_left l -JOIN vectors_explain_right r ON l.vec_id = r.vec_id -ORDER BY vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), l.vec_id -LIMIT 1; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=1 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=1), expr=[vec_l2sq_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(vec_id@0, vec_id@0)], projection=[vec_id@0, embedding@1] metrics=REDACTED_| -|_|_|_RepartitionExec: partitioning=Hash([vec_id@0],REDACTED -|_|_|_ProjectionExec: expr=[vec_id@0 as vec_id, embedding@2 as embedding] metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_RepartitionExec: partitioning=Hash([vec_id@0],REDACTED -|_|_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_FilterExec: vec_id@0 IS NOT NULL metrics=REDACTED_| -|_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -| 1_| 0_|_FilterExec: vec_id@0 IS NOT NULL metrics=REDACTED_| -|_|_|_CooperativeExec metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain_left; - -Affected Rows: 0 - -DROP TABLE vectors_explain_right; - -Affected Rows: 0 - --- ============================================ --- Part 3: Cosine and dot explain coverage --- ============================================ -CREATE TABLE vectors_explain_metric ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_explain_metric VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0]'), - (3, '2024-01-01 00:00:02', '[-1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, -1.0]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_explain_metric'); - -+---------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_explain_metric') | -+---------------------------------------------+ -| 0 | -+---------------------------------------------+ - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_cos_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_cos_distance(embedding@1, [1.0, 0.0]) ASC NULLS LAST, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_dot_product(embedding, '[1.0, 0.0]') DESC, vec_id -LIMIT 2; - -+-+-+-+ -| stage | node | plan_| -+-+-+-+ -| 0_| 0_|_CooperativeExec metrics=REDACTED_| -|_|_|_MergeScanExec: REDACTED -|_|_|_| -| 1_| 0_|_ProjectionExec: expr=[vec_id@0 as vec_id] metrics=REDACTED_| -|_|_|_SortPreservingMergeExec: [vec_dot_product(embedding@1, [1.0, 0.0]) DESC, vec_id@0 ASC NULLS LAST], fetch=2 metrics=REDACTED_| -|_|_|_SortExec: TopK(fetch=2), expr=[vec_dot_product(embedding@1, [1.0, 0.0]) DESC, vec_id@0 ASC NULLS LAST], preserve_partitioning=[true] metrics=REDACTED_| -|_|_|_SeqScan: REDACTED -|_|_|_| -|_|_| Total rows: REDACTED_| -+-+-+-+ - -DROP TABLE vectors_explain_metric; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index_explain.sql b/tests/cases/standalone/common/function/vector/vector_index_explain.sql deleted file mode 100644 index 7a0330b0ea2..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_explain.sql +++ /dev/null @@ -1,163 +0,0 @@ --- Vector index explain analyze coverage - - --- ============================================ --- Part 1: Single table KNN explain --- ============================================ -CREATE TABLE vectors_explain ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -ADMIN FLUSH_TABLE('vectors_explain'); - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE ,\s"dyn_filters":\s\[.* REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain -ORDER BY vec_l2sq_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - -DROP TABLE vectors_explain; - --- ============================================ --- Part 2: Join with vector order/limit --- ============================================ -CREATE TABLE vectors_explain_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_explain_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -INSERT INTO vectors_explain_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -ADMIN FLUSH_TABLE('vectors_explain_left'); - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT l.vec_id -FROM vectors_explain_left l -JOIN vectors_explain_right r ON l.vec_id = r.vec_id -ORDER BY vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), l.vec_id -LIMIT 1; - -DROP TABLE vectors_explain_left; -DROP TABLE vectors_explain_right; - --- ============================================ --- Part 3: Cosine and dot explain coverage --- ============================================ -CREATE TABLE vectors_explain_metric ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'cosine'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_explain_metric VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.0, 1.0]'), - (3, '2024-01-01 00:00:02', '[-1.0, 0.0]'), - (4, '2024-01-01 00:00:03', '[0.0, -1.0]'); - -ADMIN FLUSH_TABLE('vectors_explain_metric'); - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED --- SQLNESS REPLACE ,\s"vector_index_k":\s\d+ --- SQLNESS REPLACE "index_size":\d+ "index_size":REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_cos_distance(embedding, '[1.0, 0.0]'), vec_id -LIMIT 2; - --- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= --- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED --- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED --- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED --- SQLNESS REPLACE Hash\(\[vec_id@0\],.* Hash([vec_id@0],REDACTED --- SQLNESS REPLACE (-+) - --- SQLNESS REPLACE (\s\s+) _ --- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" --- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED --- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED --- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= --- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED --- SQLNESS REPLACE (peers.*) REDACTED --- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED --- SQLNESS REPLACE SeqScan:.* SeqScan: REDACTED -EXPLAIN ANALYZE VERBOSE -SELECT vec_id -FROM vectors_explain_metric -ORDER BY vec_dot_product(embedding, '[1.0, 0.0]') DESC, vec_id -LIMIT 2; - -DROP TABLE vectors_explain_metric; diff --git a/tests/cases/standalone/common/function/vector/vector_index_join.result b/tests/cases/standalone/common/function/vector/vector_index_join.result deleted file mode 100644 index 51fb3934058..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_join.result +++ /dev/null @@ -1,240 +0,0 @@ --- Vector index join/subquery coverage --- ============================================ --- Part 1: Join should not pre-limit left table --- ============================================ -CREATE TABLE vectors_join_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_join_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_join_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -INSERT INTO vectors_join_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -Affected Rows: 2 - -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_join_left l -JOIN vectors_join_right r ON l.vec_id = r.vec_id -ORDER BY dist, l.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_join_left; - -Affected Rows: 0 - -DROP TABLE vectors_join_right; - -Affected Rows: 0 - --- ============================================ --- Part 2: Subquery should be a barrier --- ============================================ -CREATE TABLE vectors_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_subquery WHERE vec_id >= 3 -) s -ORDER BY dist, s.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_subquery; - -Affected Rows: 0 - --- ============================================ --- Part 3: LEFT JOIN should not pre-limit --- ============================================ -CREATE TABLE vectors_left_join ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -CREATE TABLE vectors_left_join_filter ( - vec_id INT, - keep BOOLEAN, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_left_join VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - --- Only vec_id 3,4 have matching rows in filter table -INSERT INTO vectors_left_join_filter VALUES - (3, true, '2024-01-01 00:00:02'), - (4, true, '2024-01-01 00:00:03'); - -Affected Rows: 2 - --- LEFT JOIN then filter by IS NOT NULL --- Should return vec_id=4 (dist=1.81), not vec_id=1 or 2 -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_left_join l -LEFT JOIN vectors_left_join_filter r ON l.vec_id = r.vec_id -WHERE r.vec_id IS NOT NULL -ORDER BY dist, l.vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_left_join; - -Affected Rows: 0 - -DROP TABLE vectors_left_join_filter; - -Affected Rows: 0 - --- ============================================ --- Part 4: Inlineable subquery should allow hint --- ============================================ -CREATE TABLE vectors_inline_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_inline_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - -ADMIN FLUSH_TABLE('vectors_inline_subquery'); - -+----------------------------------------------+ -| ADMIN FLUSH_TABLE('vectors_inline_subquery') | -+----------------------------------------------+ -| 0 | -+----------------------------------------------+ - --- Subquery without LIMIT/DISTINCT/aggregation can be inlined --- Vector hint should be able to push down -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_inline_subquery WHERE vec_id >= 1 -) s -ORDER BY dist, s.vec_id -LIMIT 2; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 1 | 0.0 | -| 2 | 0.01 | -+--------+------+ - -DROP TABLE vectors_inline_subquery; - -Affected Rows: 0 - --- ============================================ --- Part 5: CTE should be a barrier --- ============================================ -CREATE TABLE vectors_cte ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -Affected Rows: 0 - -INSERT INTO vectors_cte VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -Affected Rows: 4 - --- CTE acts as optimization barrier --- Filter in CTE limits to vec_id >= 3, so result should be vec_id=4 (dist=1.81) -WITH filtered AS ( - SELECT * FROM vectors_cte WHERE vec_id >= 3 -) -SELECT vec_id, round(vec_l2sq_distance(embedding, '[1.0, 0.0]'), 2) AS dist -FROM filtered -ORDER BY dist, vec_id -LIMIT 1; - -+--------+------+ -| vec_id | dist | -+--------+------+ -| 4 | 1.81 | -+--------+------+ - -DROP TABLE vectors_cte; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/vector/vector_index_join.sql b/tests/cases/standalone/common/function/vector/vector_index_join.sql deleted file mode 100644 index 699942b3001..00000000000 --- a/tests/cases/standalone/common/function/vector/vector_index_join.sql +++ /dev/null @@ -1,159 +0,0 @@ --- Vector index join/subquery coverage - --- ============================================ --- Part 1: Join should not pre-limit left table --- ============================================ -CREATE TABLE vectors_join_left ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_join_right ( - vec_id INT, - note STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_join_left VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -INSERT INTO vectors_join_right VALUES - (3, 'keep', '2024-01-01 00:00:02'), - (4, 'keep', '2024-01-01 00:00:03'); - -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_join_left l -JOIN vectors_join_right r ON l.vec_id = r.vec_id -ORDER BY dist, l.vec_id -LIMIT 1; - -DROP TABLE vectors_join_left; -DROP TABLE vectors_join_right; - --- ============================================ --- Part 2: Subquery should be a barrier --- ============================================ -CREATE TABLE vectors_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_subquery WHERE vec_id >= 3 -) s -ORDER BY dist, s.vec_id -LIMIT 1; - -DROP TABLE vectors_subquery; - --- ============================================ --- Part 3: LEFT JOIN should not pre-limit --- ============================================ -CREATE TABLE vectors_left_join ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -CREATE TABLE vectors_left_join_filter ( - vec_id INT, - keep BOOLEAN, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_left_join VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - --- Only vec_id 3,4 have matching rows in filter table -INSERT INTO vectors_left_join_filter VALUES - (3, true, '2024-01-01 00:00:02'), - (4, true, '2024-01-01 00:00:03'); - --- LEFT JOIN then filter by IS NOT NULL --- Should return vec_id=4 (dist=1.81), not vec_id=1 or 2 -SELECT l.vec_id, round(vec_l2sq_distance(l.embedding, '[1.0, 0.0]'), 2) AS dist -FROM vectors_left_join l -LEFT JOIN vectors_left_join_filter r ON l.vec_id = r.vec_id -WHERE r.vec_id IS NOT NULL -ORDER BY dist, l.vec_id -LIMIT 1; - -DROP TABLE vectors_left_join; -DROP TABLE vectors_left_join_filter; - --- ============================================ --- Part 4: Inlineable subquery should allow hint --- ============================================ -CREATE TABLE vectors_inline_subquery ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_inline_subquery VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - -ADMIN FLUSH_TABLE('vectors_inline_subquery'); - --- Subquery without LIMIT/DISTINCT/aggregation can be inlined --- Vector hint should be able to push down -SELECT s.vec_id, round(vec_l2sq_distance(s.embedding, '[1.0, 0.0]'), 2) AS dist -FROM ( - SELECT * FROM vectors_inline_subquery WHERE vec_id >= 1 -) s -ORDER BY dist, s.vec_id -LIMIT 2; - -DROP TABLE vectors_inline_subquery; - --- ============================================ --- Part 5: CTE should be a barrier --- ============================================ -CREATE TABLE vectors_cte ( - vec_id INT, - ts TIMESTAMP TIME INDEX, - embedding VECTOR(2) NOT NULL VECTOR INDEX WITH (metric = 'l2sq'), - PRIMARY KEY (vec_id) -); - -INSERT INTO vectors_cte VALUES - (1, '2024-01-01 00:00:00', '[1.0, 0.0]'), - (2, '2024-01-01 00:00:01', '[0.9, 0.0]'), - (3, '2024-01-01 00:00:02', '[0.0, 1.0]'), - (4, '2024-01-01 00:00:03', '[0.0, 0.9]'); - --- CTE acts as optimization barrier --- Filter in CTE limits to vec_id >= 3, so result should be vec_id=4 (dist=1.81) -WITH filtered AS ( - SELECT * FROM vectors_cte WHERE vec_id >= 3 -) -SELECT vec_id, round(vec_l2sq_distance(embedding, '[1.0, 0.0]'), 2) AS dist -FROM filtered -ORDER BY dist, vec_id -LIMIT 1; - -DROP TABLE vectors_cte; diff --git a/tests/cases/standalone/common/types/vector/vector.result b/tests/cases/standalone/common/types/vector/vector.result index bc9a38c8167..1da7ad913ac 100644 --- a/tests/cases/standalone/common/types/vector/vector.result +++ b/tests/cases/standalone/common/types/vector/vector.result @@ -18,6 +18,34 @@ INSERT INTO t VALUES Affected Rows: 3 +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.004 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + +ADMIN FLUSH_TABLE('t'); + ++------------------------+ +| ADMIN FLUSH_TABLE('t') | ++------------------------+ +| 0 | ++------------------------+ + +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + ++-------------------------+ +| ts | ++-------------------------+ +| 1970-01-01T00:00:00.004 | +| 1970-01-01T00:00:00.002 | +| 1970-01-01T00:00:00.005 | ++-------------------------+ + SELECT ts, v, vec_to_string(v) FROM t; +-------------------------+--------------------------+--------------------+ diff --git a/tests/cases/standalone/common/types/vector/vector.sql b/tests/cases/standalone/common/types/vector/vector.sql index 1f81d5cd56f..46af4661947 100644 --- a/tests/cases/standalone/common/types/vector/vector.sql +++ b/tests/cases/standalone/common/types/vector/vector.sql @@ -12,6 +12,12 @@ INSERT INTO t VALUES (5, parse_vec('[4.0, 5.0, 6.0]')), (6, parse_vec('[7.0, 8.0, 9.0]')); +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + +ADMIN FLUSH_TABLE('t'); + +SELECT ts FROM t WHERE ts > 1 ORDER BY vec_l2sq_distance(v, '[1.0, 2.0, 3.0]'), ts LIMIT 3; + SELECT ts, v, vec_to_string(v) FROM t; SELECT round(vec_cos_distance(v, '[0.0, 0.0, 0.0]'), 2) FROM t; diff --git a/tests/runner/src/env/bare.rs b/tests/runner/src/env/bare.rs index dccd3673b16..43bbd987279 100644 --- a/tests/runner/src/env/bare.rs +++ b/tests/runner/src/env/bare.rs @@ -694,7 +694,7 @@ impl Env { "--bin", "greptime", "--features", - "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions", + "pg_kvbackend,mysql_kvbackend,ai_functions", ]) .output() .expect("Failed to start GreptimeDB"); From def0c2ec5ca0992083d955f6bcbc5ffaf914fe90 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 06:49:52 +0000 Subject: [PATCH 030/133] fix: only push down aggregates grouped by the partition columns themselves (#9337) * fix: only push down aggregates grouped by the partition columns themselves Signed-off-by: Dennis Zhuang * fix: keep grouping sets on the frontend for partitioned tables Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/dist_plan/analyzer/test.rs | 35 ++++++++++++++ src/query/src/dist_plan/commutativity.rs | 37 +++++++++----- .../common/aggregate/multi_regions.result | 48 +++++++++++++++++++ .../common/aggregate/multi_regions.sql | 16 +++++++ 4 files changed, 125 insertions(+), 11 deletions(-) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 5171688993c..694f2bcde59 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -40,6 +40,7 @@ use datafusion_expr::{ }; use datafusion_functions::datetime::date_bin; use datafusion_functions::datetime::expr_fn::now; +use datafusion_functions::unicode::expr_fn::substring; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; use futures::Stream; @@ -1001,6 +1002,40 @@ fn expand_proj_alias_aliased_part_col_aggr() { assert_eq!(expected, result.to_string()); } +/// `substr(pk1, 1, 1)` maps rows from different partitions to the same group, so the +/// aggregate has to be merged on the frontend even though it references `pk1`. +#[test] +fn expand_expr_over_part_col_aggr() { + init_default_ut_logging(); + let test_table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .aggregate( + vec![substring(col("pk1"), lit(1i64), lit(1i64)), col("pk2")], + vec![min(col("number"))], + ) + .unwrap() + .build() + .unwrap(); + + let config = ConfigOptions::default(); + let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + + let expected = [ + "Projection: substr(t.pk1,Int64(1),Int64(1)), t.pk2, min(t.number)", + " Aggregate: groupBy=[[substr(t.pk1,Int64(1),Int64(1)), t.pk2]], aggr=[[__min_merge(__min_state(t.number)) AS min(t.number)]]", + " MergeScan [is_placeholder=false, remote_input=[", + "Aggregate: groupBy=[[substr(t.pk1, Int64(1), Int64(1)), t.pk2]], aggr=[[__min_state(t.number)]]", + " TableScan: t", + "]]", + ] + .join("\n"); + assert_eq!(expected, result.to_string()); +} + /// notice that step aggr then part col aggr seems impossible as the partition columns for part col aggr /// can't pass through the step aggr without making step aggr also a part col aggr /// so here only test part col aggr -> step aggr case diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index a236118822b..d02b746a8bf 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -117,7 +117,14 @@ impl Categorizer { LogicalPlan::Filter(filter) => Self::check_expr(&filter.predicate), LogicalPlan::Window(_) => Commutativity::Unimplemented, LogicalPlan::Aggregate(aggr) => { - let is_all_steppable = is_all_aggr_exprs_steppable(&aggr.aggr_expr); + // The state/merge split maps each group expression to one output column, + // which doesn't hold for grouping sets. + let has_grouping_set = aggr + .group_expr + .iter() + .any(|expr| matches!(expr, Expr::GroupingSet(_))); + let is_all_steppable = + !has_grouping_set && is_all_aggr_exprs_steppable(&aggr.aggr_expr); let matches_partition = Self::check_partition(&aggr.group_expr, &partition_cols); if !matches_partition && is_all_steppable { debug!("Plan is steppable: {plan}"); @@ -303,26 +310,34 @@ impl Categorizer { /// Return true if the given expr and partition cols satisfied the rule. /// In this case the plan can be treated as fully commutative. /// - /// So only if all partition columns show up in `exprs`, return true. + /// So only if every partition column is itself one of `exprs`, return true. /// Otherwise return false. /// + /// An expression that only references a partition column, like `substr(host, 3, 1)` + /// or `k % 2`, doesn't count: it can put rows from different partitions into the same + /// group. fn check_partition(exprs: &[Expr], partition_cols: &AliasMapping) -> bool { - let mut ref_cols = HashSet::new(); - for expr in exprs { - expr.add_column_refs(&mut ref_cols); - } - let ref_cols = ref_cols - .into_iter() - .map(|c| c.name.clone()) + let group_cols = exprs + .iter() + .filter_map(|expr| { + let mut expr = expr; + while let Expr::Alias(alias) = expr { + expr = &alias.expr; + } + match expr { + Expr::Column(column) => Some(column.name.clone()), + _ => None, + } + }) .collect::>(); for all_alias in partition_cols.values() { let all_alias = all_alias .iter() .map(|c| c.name.clone()) .collect::>(); - // check if ref columns intersect with all alias of partition columns + // check if group columns intersect with all alias of partition columns // is empty, if it's empty, not all partition columns show up in `exprs` - if ref_cols.intersection(&all_alias).count() == 0 { + if group_cols.intersection(&all_alias).count() == 0 { return false; } } diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index fa0bf929319..79ae85fbc04 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -113,6 +113,54 @@ select sum(val) from t group by idc; |_|_| Total rows: 0_| +-+-+-+ +insert into t values + (1000, 1, '1000-x', 'a'), + (1000, 2, '1000-y', 'a'), + (1000, 3, '2000-x', 'a'), + (1000, 4, '2000-y', 'a'); + +Affected Rows: 4 + +-- group keys derived from the partition column span regions +select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; + ++---+----------+------------+ +| g | count(*) | sum(t.val) | ++---+----------+------------+ +| x | 2 | 4.0 | +| y | 2 | 6.0 | ++---+----------+------------+ + +-- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them +-- SQLNESS PROTOCOL POSTGRES +select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; + ++--------+-----+------------+ +| host | idc | sum(t.val) | ++--------+-----+------------+ +| 1000-x | a | 1.0 | +| 1000-x | | 1.0 | +| 1000-y | a | 2.0 | +| 1000-y | | 2.0 | +| 2000-x | a | 3.0 | +| 2000-x | | 3.0 | +| 2000-y | a | 4.0 | +| 2000-y | | 4.0 | ++--------+-----+------------+ + +-- SQLNESS PROTOCOL POSTGRES +select host, sum(val) from t group by rollup(host) order by host; + ++--------+------------+ +| host | sum(t.val) | ++--------+------------+ +| 1000-x | 1.0 | +| 1000-y | 2.0 | +| 2000-x | 3.0 | +| 2000-y | 4.0 | +| | 10.0 | ++--------+------------+ + drop table t; Affected Rows: 0 diff --git a/tests/cases/standalone/common/aggregate/multi_regions.sql b/tests/cases/standalone/common/aggregate/multi_regions.sql index 3396bdddf19..4f99e1ed465 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.sql +++ b/tests/cases/standalone/common/aggregate/multi_regions.sql @@ -45,4 +45,20 @@ select sum(val) from t; explain analyze select sum(val) from t group by idc; +insert into t values + (1000, 1, '1000-x', 'a'), + (1000, 2, '1000-y', 'a'), + (1000, 3, '2000-x', 'a'), + (1000, 4, '2000-y', 'a'); + +-- group keys derived from the partition column span regions +select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g; + +-- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them +-- SQLNESS PROTOCOL POSTGRES +select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc; + +-- SQLNESS PROTOCOL POSTGRES +select host, sum(val) from t group by rollup(host) order by host; + drop table t; From affc0a1b1ddf57a79183b65335882071656d416e Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 07:02:14 +0000 Subject: [PATCH 031/133] fix: align ordered aggregate state type with the accumulator output (#9340) * fix: align ordered aggregate state type with the accumulator output Signed-off-by: Dennis Zhuang * fix: reject mismatched aggregate states and keep hard-ordered aggregates unsplit Signed-off-by: Dennis Zhuang * fix: keep WITHIN GROUP aggregates splittable and relabel state fields by position Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/aggrs/aggr_wrapper.rs | 163 ++++++++++-------- .../src/aggrs/aggr_wrapper/fix_order.rs | 6 +- .../function/src/aggrs/aggr_wrapper/tests.rs | 123 +++++++++++++ .../common/aggregate/array_agg.result | 91 ++++++++++ .../standalone/common/aggregate/array_agg.sql | 38 ++++ 5 files changed, 348 insertions(+), 73 deletions(-) diff --git a/src/common/function/src/aggrs/aggr_wrapper.rs b/src/common/function/src/aggrs/aggr_wrapper.rs index 70651eee41d..710c383a834 100644 --- a/src/common/function/src/aggrs/aggr_wrapper.rs +++ b/src/common/function/src/aggrs/aggr_wrapper.rs @@ -25,7 +25,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, StructArray}; +use arrow::array::{ArrayData, ArrayRef, BooleanArray, StructArray, make_array}; use arrow_schema::{FieldRef, Fields}; use common_telemetry::debug; use datafusion::functions_aggregate::all_default_aggregate_functions; @@ -84,6 +84,19 @@ pub fn is_all_aggr_exprs_steppable(aggr_exprs: &[Expr]) -> bool { return false; } + // DataFusion only sorts the input of an aggregate with a hard ordering requirement + // when the requirement is already satisfied or the aggregate has a reverse + // expression (apache/datafusion#25676). The state wrapper has none, so e.g. + // `nth_value(.. ORDER BY ..)` would read unsorted input on datanodes. Ordered-set + // aggregates like `approx_percentile_cont(..) WITHIN GROUP (ORDER BY ..)` are + // exempt: their ORDER BY names the value, and they don't need sorted input. + if !aggr_func.params.order_by.is_empty() + && aggr_func.func.order_sensitivity().hard_requires() + && !aggr_func.func.supports_within_group_clause() + { + return false; + } + // whether the corresponding state function exists in the registry FUNCTION_REGISTRY.is_aggr_func_exist(&aggr_state_func_name(aggr_func.func.name())) } else { @@ -526,43 +539,7 @@ impl StateGroupsAccum { } fn wrap_state_arrays(&self, arrays: Vec) -> datafusion_common::Result { - let array_type = arrays - .iter() - .map(|array| array.data_type().clone()) - .collect::>(); - let expected_type = self - .state_fields - .iter() - .map(|field| field.data_type().clone()) - .collect::>(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - arrays.len(), - self.state_fields, - array_type, - ); - let guess_schema = arrays - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let array = StructArray::try_new(guess_schema, arrays, None)?; - return Ok(Arc::new(array)); - } - - Ok(Arc::new(StructArray::try_new( - self.state_fields.clone(), - arrays, - None, - )?)) + Ok(Arc::new(state_struct_array(&self.state_fields, arrays)?)) } } @@ -610,6 +587,81 @@ impl GroupsAccumulator for StateGroupsAccum { } } +/// Wraps the state arrays of an accumulator into a struct of the declared state fields. +/// +/// The declared state type is derived from logical expressions, while the accumulator names +/// nested fields after physical expressions. For example, `array_agg(v ORDER BY ts)` declares +/// its orderings as `List(Struct("ts": ..))` but produces `List(Struct("ts@0": ..))`. Arrays that +/// differ only in nested field names are cast to the declared type; any other difference is an +/// error. +fn state_struct_array( + state_fields: &Fields, + arrays: Vec, +) -> datafusion_common::Result { + if arrays.len() != state_fields.len() { + return Err(datafusion_common::DataFusionError::Internal(format!( + "Expected {} state arrays for fields {:?}, got {}", + state_fields.len(), + state_fields, + arrays.len() + ))); + } + let arrays = arrays + .into_iter() + .zip(state_fields.iter()) + .map(|(array, field)| { + let expected = field.data_type(); + if array.data_type() == expected { + Ok(array) + } else if array.data_type().equals_datatype(expected) { + Ok(make_array(relabel_nested_fields( + array.to_data(), + expected, + )?)) + } else { + Err(datafusion_common::DataFusionError::Internal(format!( + "State field `{}` expects type {expected}, but the accumulator produced {}", + field.name(), + array.data_type() + ))) + } + }) + .collect::>>()?; + Ok(StructArray::try_new(state_fields.clone(), arrays, None)?) +} + +/// Rebuilds `data` with the type `target`, which must match it position by position apart +/// from nested field names and metadata. Unlike a cast, children are never matched by name. +fn relabel_nested_fields( + data: ArrayData, + target: &DataType, +) -> datafusion_common::Result { + let child_types = match target { + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => vec![field.data_type()], + DataType::Struct(fields) => fields.iter().map(|f| f.data_type()).collect(), + _ if data.child_data().is_empty() => vec![], + _ => { + return Err(datafusion_common::DataFusionError::NotImplemented(format!( + "Relabeling nested fields of {target}" + ))); + } + }; + let children = data + .child_data() + .iter() + .zip(child_types) + .map(|(child, child_type)| relabel_nested_fields(child.clone(), child_type)) + .collect::>>()?; + Ok(data + .into_builder() + .data_type(target.clone()) + .child_data(children) + .build()?) +} + impl StateAccum { pub fn new( inner: Box, @@ -636,40 +688,7 @@ impl Accumulator for StateAccum { .iter() .map(|s| s.to_array()) .collect::, _>>()?; - let array_type = array - .iter() - .map(|a| a.data_type().clone()) - .collect::>(); - let expected_type: Vec<_> = self - .state_fields - .iter() - .map(|f| f.data_type().clone()) - .collect(); - if array_type != expected_type { - debug!( - "State mismatch, expected: {}, got: {} for expected fields: {:?} and given array types: {:?}", - self.state_fields.len(), - array.len(), - self.state_fields, - array_type, - ); - let guess_schema = array - .iter() - .enumerate() - .map(|(index, array)| { - Field::new( - format!("col_{index}[mismatch_state]").as_str(), - array.data_type().clone(), - true, - ) - }) - .collect::(); - let arr = StructArray::try_new(guess_schema, array, None)?; - - return Ok(ScalarValue::Struct(Arc::new(arr))); - } - - let struct_array = StructArray::try_new(self.state_fields.clone(), array, None)?; + let struct_array = state_struct_array(&self.state_fields, array)?; Ok(ScalarValue::Struct(Arc::new(struct_array))) } diff --git a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs index 480b0b5a490..715b051407d 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/fix_order.rs @@ -153,13 +153,17 @@ fn rewrite_expr( if is_fix { // then always fix the ordering field&distinct flag and more let order_by = aggregate_function.params.order_by.clone(); + // DataFusion always makes ordering fields nullable when building the physical + // aggregate (`ordering_fields` in `datafusion-functions-aggregate-common`), and some + // accumulators (e.g. `array_agg`) nest these fields in their state type. Keep the + // declared state type in line with what the accumulator produces. let ordering_fields: Vec<_> = order_by .iter() .map(|sort_expr| { sort_expr .expr .to_field(&aggregate_input.schema()) - .map(|(_, f)| f) + .map(|(_, f)| Arc::new(f.as_ref().clone().with_nullable(true))) }) .collect::>>()?; let distinct = aggregate_function.params.distinct; diff --git a/src/common/function/src/aggrs/aggr_wrapper/tests.rs b/src/common/function/src/aggrs/aggr_wrapper/tests.rs index 4274a03cd89..b68a735f182 100644 --- a/src/common/function/src/aggrs/aggr_wrapper/tests.rs +++ b/src/common/function/src/aggrs/aggr_wrapper/tests.rs @@ -27,8 +27,11 @@ use datafusion::catalog::{Session, TableProvider}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::DefaultTableSource; use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; +use datafusion::functions_aggregate::array_agg::array_agg_udaf; use datafusion::functions_aggregate::average::avg_udaf; use datafusion::functions_aggregate::count::count_udaf; +use datafusion::functions_aggregate::nth_value::nth_value_udaf; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::optimizer::AnalyzerRule; use datafusion::optimizer::analyzer::type_coercion::TypeCoercion; @@ -1289,6 +1292,40 @@ async fn test_udaf_correct_eval_result() { order_by: vec![], null_treatment: None, }, + // The ordering state of `array_agg` nests the ORDER BY fields in `List(Struct(..))`. + TestCase { + func: array_agg_udaf(), + input_schema: Arc::new(arrow_schema::Schema::new(vec![ + Field::new("number", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + false, + ), + ])), + args: vec![Expr::Column(Column::new_unqualified("number"))], + input: vec![ + Arc::new(Float64Array::from(vec![Some(3.), Some(1.), Some(2.)])), + Arc::new(TimestampMillisecondArray::from(vec![3000, 1000, 2000])), + ], + expected_output: Some(ScalarValue::List(ScalarValue::new_list_nullable( + &[ + ScalarValue::Float64(Some(1.)), + ScalarValue::Float64(Some(2.)), + ScalarValue::Float64(Some(3.)), + ], + &DataType::Float64, + ))), + expected_fn: None, + distinct: false, + filter: None, + order_by: vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )], + null_treatment: None, + }, // TODO(discord9): udd_merge/hll_merge/geo_path/quantile_aggr tests ]; let test_table_ref = TableReference::bare("TestTable"); @@ -1396,3 +1433,89 @@ async fn execute_phy_plan( } Ok(batches) } + +#[test] +fn test_state_struct_array_rejects_mismatched_state() { + let fields = Fields::from(vec![Field::new("sum", DataType::Int64, true)]); + let arrays: Vec = vec![Arc::new(Float64Array::from(vec![1.0]))]; + let err = state_struct_array(&fields, arrays).unwrap_err(); + assert!( + err.to_string() + .contains("State field `sum` expects type Int64, but the accumulator produced Float64"), + "{err}" + ); +} + +#[test] +fn test_state_struct_array_keeps_child_order() { + let int_field = |name: &str| Field::new(name, DataType::Int64, true); + // Same child types, names crossed: children must stay in place, not be matched by name. + let produced = StructArray::from(vec![ + ( + Arc::new(int_field("a")), + Arc::new(Int64Array::from(vec![10])) as ArrayRef, + ), + ( + Arc::new(int_field("b")), + Arc::new(Int64Array::from(vec![20])) as ArrayRef, + ), + ]); + let declared = Fields::from(vec![Field::new( + "state", + DataType::Struct(Fields::from(vec![int_field("b"), int_field("a")])), + true, + )]); + + let state = state_struct_array(&declared, vec![Arc::new(produced)]).unwrap(); + let state = state.column(0).as_struct(); + assert_eq!( + state + .column_by_name("b") + .unwrap() + .as_primitive::() + .value(0), + 10 + ); + assert_eq!( + state + .column_by_name("a") + .unwrap() + .as_primitive::() + .value(0), + 20 + ); +} + +#[test] +fn test_hard_ordered_aggr_not_steppable() { + let order_by = vec![SortExpr::new( + Expr::Column(Column::new_unqualified("ts")), + true, + true, + )]; + let aggr = |func: Arc, args: Vec| { + Expr::AggregateFunction(AggregateFunction::new_udf( + func, + args, + false, + None, + order_by.clone(), + None, + )) + }; + let number = Expr::Column(Column::new_unqualified("number")); + + assert!(!is_all_aggr_exprs_steppable(&[aggr( + nth_value_udaf(), + vec![number.clone(), lit(2i64)], + )])); + assert!(is_all_aggr_exprs_steppable(&[aggr( + array_agg_udaf(), + vec![number.clone()] + )])); + // WITHIN GROUP (ORDER BY number) + assert!(is_all_aggr_exprs_steppable(&[aggr( + approx_percentile_cont_udaf(), + vec![number, lit(0.5f64)] + )])); +} diff --git a/tests/cases/standalone/common/aggregate/array_agg.result b/tests/cases/standalone/common/aggregate/array_agg.result index 5a73c289899..236a4582f14 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.result +++ b/tests/cases/standalone/common/aggregate/array_agg.result @@ -144,3 +144,94 @@ DROP TABLE doubles; Affected Rows: 0 +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +Affected Rows: 7 + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + ++-----------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | ++-----------------------------------------------------------------------------------------+ +| [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] | ++-----------------------------------------------------------------------------------------+ + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + ++-------+-------------------------------------------------------------------------------------------+ +| g | array_agg(array_agg_partitioned.lat) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++-------+-------------------------------------------------------------------------------------------+ +| false | [3.0, 2.0, 1.0] | +| true | [7.0, 6.0, 5.0, 4.0] | ++-------+-------------------------------------------------------------------------------------------+ + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| array_agg(array_agg_partitioned.k) ORDER BY [array_agg_partitioned.k % Int64(10) ASC NULLS LAST, array_agg_partitioned.ts DESC NULLS FIRST] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ +| [21, 11, 1, 22, 12, 2, 3] | ++---------------------------------------------------------------------------------------------------------------------------------------------+ + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| nth_value(array_agg_partitioned.lat,Int64(2)) ORDER BY [array_agg_partitioned.ts ASC NULLS LAST] | nth_value(array_agg_partitioned.lat,Int64(3)) ORDER BY [array_agg_partitioned.ts DESC NULLS FIRST] | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ +| 2.0 | 5.0 | ++--------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------+ + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Aggregate: groupBy=[[]], aggr=[[__sum_merge(__sum_state(array_agg_partitioned.lat)) AS sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) AS approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Aggregate: groupBy=[[]], aggr=[[__sum_state(array_agg_partitioned.lat), __approx_percentile_cont_state(array_agg_partitioned.lat, Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]]]_| +|_|_TableScan: array_agg_partitioned_| +|_| ]]_| +| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[__sum_merge(__sum_state(array_agg_partitioned.lat)) as sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) as approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]]_| +|_|_CoalescePartitionsExec_| +|_|_AggregateExec: mode=Partial, gby=[], aggr=[__sum_merge(__sum_state(array_agg_partitioned.lat)) as sum(array_agg_partitioned.lat), __approx_percentile_cont_merge(__approx_percentile_cont_state(array_agg_partitioned.lat,Float64(0.75)) ORDER BY [array_agg_partitioned.lat DESC NULLS FIRST]) as approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST]] | +|_|_RepartitionExec: partitioning=REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| sum(array_agg_partitioned.lat) | approx_percentile_cont(Float64(0.75)) WITHIN GROUP [array_agg_partitioned.lat DESC NULLS FIRST] | ++--------------------------------+-------------------------------------------------------------------------------------------------+ +| 28.0 | 2.25 | ++--------------------------------+-------------------------------------------------------------------------------------------------+ + +DROP TABLE array_agg_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/array_agg.sql b/tests/cases/standalone/common/aggregate/array_agg.sql index dedabf1a18f..d0d7caa574f 100644 --- a/tests/cases/standalone/common/aggregate/array_agg.sql +++ b/tests/cases/standalone/common/aggregate/array_agg.sql @@ -54,3 +54,41 @@ DROP TABLE integers; DROP TABLE strings; DROP TABLE doubles; + +-- On partitioned tables the aggregate is split into partial state and merge +CREATE TABLE array_agg_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + lat DOUBLE, + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO array_agg_partitioned VALUES + (1000, 1, 1), + (2000, 11, 2), + (3000, 21, 3), + (4000, 2, 4), + (5000, 12, 5), + (6000, 22, 6), + (7000, 3, 7); + +SELECT array_agg(lat ORDER BY ts) FROM array_agg_partitioned; + +SELECT lat > 3 AS g, array_agg(lat ORDER BY ts DESC) FROM array_agg_partitioned GROUP BY g ORDER BY g; + +SELECT array_agg(k ORDER BY k % 10, ts DESC) FROM array_agg_partitioned; + +-- nth_value needs sorted input, so it isn't split into partial state and merge +SELECT nth_value(lat, 2 ORDER BY ts), nth_value(lat, 3 ORDER BY ts DESC) FROM array_agg_partitioned; + +-- WITHIN GROUP aggregates don't need sorted input and are still split +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +EXPLAIN SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +SELECT sum(lat), approx_percentile_cont(0.75) WITHIN GROUP (ORDER BY lat DESC) FROM array_agg_partitioned; + +DROP TABLE array_agg_partitioned; From d59725b04ac9862830b82733e0746d6027e035c3 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 07:03:50 +0000 Subject: [PATCH 032/133] test(sqlness): run environments concurrently and drop redundant restarts and sleeps (#9333) * test(sqlness): run environments concurrently with external backends The runner forced both environment and instance parallelism to 1 whenever etcd/PG/MySQL or an external Kafka was set up, so the standalone and distributed environments ran one after the other. Only distributed uses the kv backend, and the two environments use different Kafka topic prefixes, so they can run at the same time. Keep one instance per environment, since instances would share the metadata table and Kafka topics. Runs with a runner-managed Kafka, an external server, a test filter, or `-j 1` stay fully serial. Signed-off-by: Dennis Zhuang * test(sqlness): drop redundant restarts and sleeps - flow_rebuild: drop the R2 block, which repeats R1 from the same state, and the restart in R5, whose result R1 already asserts. - flow_advance_ttl, flow_view, ttl_instant: drop sleeps that no assertion depends on. - flow_basic, flow_null, flow_call_df_func: drop the bytes_log section (covered by flow_insert), a duplicate state_size query, an unchecked insert, and flushes that run with no new data. - alter_table_options, skip_wal: share restarts between independent tables. - session_skip_wal, copy_skip_wal: check every case after a single restart. Tables copied or inserted with skip_wal = true first get a WAL row that is truncated, so the check that truncated WAL entries are not replayed stays. - region_statistics: wait once for the statistics of all three tables. - build_index_table: fold its index_size checks into build_index_table_restart, which runs the same fixture and waits. Signed-off-by: Dennis Zhuang * test(sqlness): restore session_skip_wal and copy_skip_wal Checking every case after a single restart dropped two state transitions the original cases cover: recovering a region whose memtable only held skipped rows, and writing to the recovered region before restarting again. Restore the original cases. Signed-off-by: Dennis Zhuang * test(sqlness): restore the one-restart check in flow_auto_sink_table #5112 checked the auto-created sink before a restart and the flow and sink after it. #5987 moved the restart in front of the first SHOW and added a second one. Flow recovery only creates the sink when it is missing, so the second restart recovers from the same persisted state as the first. Restore the original before/after check with one restart. Signed-off-by: Dennis Zhuang * test(sqlness): keep one instance per environment with external store addresses Instances of the distributed environment share the etcd behind --store-addrs, so run one instance per environment when it is set, same as --setup-etcd. Signed-off-by: Dennis Zhuang * test(sqlness): restore the false-branch batch in flow_basic and assert it The kept batch (20, 22) is all above the threshold. The second batch (10, 23) is the only input that hits the false branch of the CASE and makes the flow compute again. Restore it and check the result, which the original case never did. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .../common/alter/alter_table_options.result | 97 +++++---- .../common/alter/alter_table_options.sql | 25 ++- .../common/flow/flow_advance_ttl.result | 2 - .../common/flow/flow_advance_ttl.sql | 2 - .../common/flow/flow_auto_sink_table.result | 10 - .../common/flow/flow_auto_sink_table.sql | 3 - .../standalone/common/flow/flow_basic.result | 192 +----------------- .../standalone/common/flow/flow_basic.sql | 98 +-------- .../common/flow/flow_call_df_func.result | 72 ------- .../common/flow/flow_call_df_func.sql | 24 --- .../standalone/common/flow/flow_null.result | 12 -- .../standalone/common/flow/flow_null.sql | 6 - .../common/flow/flow_rebuild.result | 78 ------- .../standalone/common/flow/flow_rebuild.sql | 42 ---- .../standalone/common/flow/flow_view.result | 2 - .../standalone/common/flow/flow_view.sql | 2 - .../function/admin/build_index_table.result | 72 ------- .../function/admin/build_index_table.sql | 29 --- .../admin/build_index_table_restart.result | 42 ++++ .../admin/build_index_table_restart.sql | 20 ++ .../region_statistics.result | 57 +++--- .../information_schema/region_statistics.sql | 41 ++-- tests/cases/standalone/common/skip_wal.result | 89 ++++---- tests/cases/standalone/common/skip_wal.sql | 39 ++-- .../standalone/common/ttl/ttl_instant.result | 1 - .../standalone/common/ttl/ttl_instant.sql | 1 - tests/runner/src/cmd/bare.rs | 129 ++++++++++-- 27 files changed, 354 insertions(+), 833 deletions(-) delete mode 100644 tests/cases/standalone/common/function/admin/build_index_table.result delete mode 100644 tests/cases/standalone/common/function/admin/build_index_table.sql diff --git a/tests/cases/standalone/common/alter/alter_table_options.result b/tests/cases/standalone/common/alter/alter_table_options.result index 85e64db429f..09afb247055 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.result +++ b/tests/cases/standalone/common/alter/alter_table_options.result @@ -278,6 +278,36 @@ SHOW CREATE TABLE ato; | | ) | +-------+--------------------------------------------------------------+ +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +Affected Rows: 0 + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +Affected Rows: 0 + +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + -- SQLNESS ARG restart=true SHOW CREATE TABLE ato; @@ -303,6 +333,24 @@ SHOW CREATE TABLE ato; | | ) | +-------+--------------------------------------------------------------+ +SHOW CREATE TABLE write_buffer_size_options; + ++---------------------------+----------------------------------------------------------+ +| Table | Create Table | ++---------------------------+----------------------------------------------------------+ +| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | +| | "i" INT NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("i") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | write_buffer_size = '128.0MiB' | +| | ) | ++---------------------------+----------------------------------------------------------+ + ALTER TABLE ato UNSET 'compaction.twcs.active_window.trigger_file_num'; Affected Rows: 0 @@ -364,55 +412,6 @@ DROP TABLE ato; Affected Rows: 0 -CREATE TABLE write_buffer_size_options( - i INTEGER, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(i) -); - -Affected Rows: 0 - -ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; - -Affected Rows: 0 - -SHOW CREATE TABLE write_buffer_size_options; - -+---------------------------+----------------------------------------------------------+ -| Table | Create Table | -+---------------------------+----------------------------------------------------------+ -| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | -| | "i" INT NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("i") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | write_buffer_size = '128.0MiB' | -| | ) | -+---------------------------+----------------------------------------------------------+ - --- SQLNESS ARG restart=true -SHOW CREATE TABLE write_buffer_size_options; - -+---------------------------+----------------------------------------------------------+ -| Table | Create Table | -+---------------------------+----------------------------------------------------------+ -| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( | -| | "i" INT NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("i") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | write_buffer_size = '128.0MiB' | -| | ) | -+---------------------------+----------------------------------------------------------+ - ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; Affected Rows: 0 diff --git a/tests/cases/standalone/common/alter/alter_table_options.sql b/tests/cases/standalone/common/alter/alter_table_options.sql index dd32ee0d4c3..de6470037b0 100644 --- a/tests/cases/standalone/common/alter/alter_table_options.sql +++ b/tests/cases/standalone/common/alter/alter_table_options.sql @@ -62,9 +62,21 @@ ALTER TABLE ato SET 'compaction.twcs.inactive_window.l1_merge_trigger'='12'; SHOW CREATE TABLE ato; +CREATE TABLE write_buffer_size_options( + i INTEGER, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY(i) +); + +ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; + +SHOW CREATE TABLE write_buffer_size_options; + -- SQLNESS ARG restart=true SHOW CREATE TABLE ato; +SHOW CREATE TABLE write_buffer_size_options; + ALTER TABLE ato UNSET 'compaction.twcs.active_window.trigger_file_num'; ALTER TABLE ato UNSET 'compaction.twcs.active_window.l1_merge_trigger'; @@ -80,19 +92,6 @@ SHOW CREATE TABLE ato; DROP TABLE ato; -CREATE TABLE write_buffer_size_options( - i INTEGER, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(i) -); - -ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB'; - -SHOW CREATE TABLE write_buffer_size_options; - --- SQLNESS ARG restart=true -SHOW CREATE TABLE write_buffer_size_options; - ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size'; SHOW CREATE TABLE write_buffer_size_options; diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.result b/tests/cases/standalone/common/flow/flow_advance_ttl.result index 279a6552d21..eae65328a6b 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.result +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.result @@ -116,7 +116,6 @@ SELECT number FROM distinct_basic; +--------+ +--------+ --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); +-------------------------------------+ @@ -356,7 +355,6 @@ SELECT number FROM distinct_basic; | 22 | +--------+ --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); +-------------------------------------+ diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.sql b/tests/cases/standalone/common/flow/flow_advance_ttl.sql index 653fa16637f..2aacc0f9458 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.sql +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.sql @@ -51,7 +51,6 @@ FROM SELECT number FROM distinct_basic; --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); -- Recover the persisted streaming DISTINCT flow, then replan its first write @@ -158,7 +157,6 @@ FROM SELECT number FROM distinct_basic; --- SQLNESS SLEEP 6s ADMIN FLUSH_TABLE('distinct_basic'); INSERT INTO diff --git a/tests/cases/standalone/common/flow/flow_auto_sink_table.result b/tests/cases/standalone/common/flow/flow_auto_sink_table.result index 4ef3297125b..e2f28e0c441 100644 --- a/tests/cases/standalone/common/flow/flow_auto_sink_table.result +++ b/tests/cases/standalone/common/flow/flow_auto_sink_table.result @@ -114,16 +114,6 @@ GROUP BY Affected Rows: 0 --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s SHOW CREATE TABLE out_num_cnt_basic; +-------------------+---------------------------------------------------+ diff --git a/tests/cases/standalone/common/flow/flow_auto_sink_table.sql b/tests/cases/standalone/common/flow/flow_auto_sink_table.sql index e67aa961910..87be556f66c 100644 --- a/tests/cases/standalone/common/flow/flow_auto_sink_table.sql +++ b/tests/cases/standalone/common/flow/flow_auto_sink_table.sql @@ -47,10 +47,7 @@ FROM numbers_input_basic GROUP BY ts; --- SQLNESS ARG restart=true -SELECT 1; --- SQLNESS SLEEP 3s SHOW CREATE TABLE out_num_cnt_basic; -- SQLNESS ARG restart=true diff --git a/tests/cases/standalone/common/flow/flow_basic.result b/tests/cases/standalone/common/flow/flow_basic.result index 5144c459ad6..f8f98e45dcd 100644 --- a/tests/cases/standalone/common/flow/flow_basic.result +++ b/tests/cases/standalone/common/flow/flow_basic.result @@ -447,190 +447,6 @@ DROP TABLE out_distinct_basic; Affected Rows: 0 -CREATE TABLE bytes_log ( - byte INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - -- event time - TIME INDEX(ts) -); - -Affected Rows: 0 - --- TODO(discord9): remove this after auto infer table's time index is impl -CREATE TABLE approx_rate ( - rate DOUBLE, - time_window TIMESTAMP, - update_at TIMESTAMP, - TIME INDEX(time_window) -); - -Affected Rows: 0 - -CREATE FLOW find_approx_rate SINK TO approx_rate AS -SELECT - (max(byte) - min(byte)) / 30.0 as rate, - date_bin(INTERVAL '30 second', ts) as time_window -from - bytes_log -GROUP BY - time_window; - -Affected Rows: 0 - -SHOW CREATE TABLE approx_rate; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| approx_rate | CREATE TABLE IF NOT EXISTS "approx_rate" ( | -| | "rate" DOUBLE NULL, | -| | "time_window" TIMESTAMP(3) NOT NULL, | -| | "update_at" TIMESTAMP(3) NULL, | -| | TIME INDEX ("time_window") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+-------------+--------------------------------------------+ - -INSERT INTO - bytes_log -VALUES - (NULL, '2023-01-01 00:00:01'), - (300, '2023-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+------+---------------------+ -| rate | time_window | -+------+---------------------+ -| 0.0 | 2023-01-01T00:00:00 | -+------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (NULL, '2022-01-01 00:00:01'), - (NULL, '2022-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+------+---------------------+ -| rate | time_window | -+------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -+------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (101, '2025-01-01 00:00:01'), - (300, '2025-01-01 00:00:29'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+-------------------+---------------------+ -| rate | time_window | -+-------------------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -| 6.633333333333334 | 2025-01-01T00:00:00 | -+-------------------+---------------------+ - -INSERT INTO - bytes_log -VALUES - (450, '2025-01-01 00:00:32'), - (500, '2025-01-01 00:00:37'); - -Affected Rows: 2 - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -+--------------------------------------+ -| ADMIN FLUSH_FLOW('find_approx_rate') | -+--------------------------------------+ -| FLOW_FLUSHED | -+--------------------------------------+ - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -+--------------------+---------------------+ -| rate | time_window | -+--------------------+---------------------+ -| | 2022-01-01T00:00:00 | -| 0.0 | 2023-01-01T00:00:00 | -| 6.633333333333334 | 2025-01-01T00:00:00 | -| 1.6666666666666667 | 2025-01-01T00:00:30 | -+--------------------+---------------------+ - -DROP TABLE bytes_log; - -Affected Rows: 0 - -DROP FLOW find_approx_rate; - -Affected Rows: 0 - -DROP TABLE approx_rate; - -Affected Rows: 0 - -- input table CREATE TABLE ngx_access_log ( client STRING, @@ -1774,6 +1590,14 @@ ADMIN FLUSH_FLOW('test_numbers_basic'); | FLOW_FLUSHED | +----------------------------------------+ +SELECT avg_after_filter_num FROM out_num_cnt_basic; + ++----------------------+ +| avg_after_filter_num | ++----------------------+ +| 0 | ++----------------------+ + DROP FLOW test_numbers_basic; Affected Rows: 0 diff --git a/tests/cases/standalone/common/flow/flow_basic.sql b/tests/cases/standalone/common/flow/flow_basic.sql index ef95ad88a9f..62e09db6d27 100644 --- a/tests/cases/standalone/common/flow/flow_basic.sql +++ b/tests/cases/standalone/common/flow/flow_basic.sql @@ -180,102 +180,6 @@ DROP TABLE distinct_basic; DROP TABLE out_distinct_basic; -CREATE TABLE bytes_log ( - byte INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - -- event time - TIME INDEX(ts) -); - --- TODO(discord9): remove this after auto infer table's time index is impl -CREATE TABLE approx_rate ( - rate DOUBLE, - time_window TIMESTAMP, - update_at TIMESTAMP, - TIME INDEX(time_window) -); - -CREATE FLOW find_approx_rate SINK TO approx_rate AS -SELECT - (max(byte) - min(byte)) / 30.0 as rate, - date_bin(INTERVAL '30 second', ts) as time_window -from - bytes_log -GROUP BY - time_window; - -SHOW CREATE TABLE approx_rate; - -INSERT INTO - bytes_log -VALUES - (NULL, '2023-01-01 00:00:01'), - (300, '2023-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (NULL, '2022-01-01 00:00:01'), - (NULL, '2022-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (101, '2025-01-01 00:00:01'), - (300, '2025-01-01 00:00:29'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -INSERT INTO - bytes_log -VALUES - (450, '2025-01-01 00:00:32'), - (500, '2025-01-01 00:00:37'); - --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('find_approx_rate'); - -SELECT - rate, - time_window -FROM - approx_rate -ORDER BY time_window ASC; - -DROP TABLE bytes_log; - -DROP FLOW find_approx_rate; - -DROP TABLE approx_rate; - -- input table CREATE TABLE ngx_access_log ( client STRING, @@ -827,6 +731,8 @@ VALUES -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_numbers_basic'); +SELECT avg_after_filter_num FROM out_num_cnt_basic; + DROP FLOW test_numbers_basic; DROP TABLE numbers_input_basic; DROP TABLE out_num_cnt_basic; diff --git a/tests/cases/standalone/common/flow/flow_call_df_func.result b/tests/cases/standalone/common/flow/flow_call_df_func.result index 960957a8e47..3bddade9825 100644 --- a/tests/cases/standalone/common/flow/flow_call_df_func.result +++ b/tests/cases/standalone/common/flow/flow_call_df_func.result @@ -15,15 +15,6 @@ SELECT sum(abs(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00' Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -50,15 +41,6 @@ SELECT "sum(abs(numbers_input_df_func.number))", time_window FROM out_num_cnt_df | 42 | 2021-07-01T00:00:00 | +----------------------------------------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -114,15 +96,6 @@ SELECT abs(sum(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00' Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -148,15 +121,6 @@ SELECT "abs(sum(numbers_input_df_func.number))", time_window FROM out_num_cnt_df | 2 | 2021-07-01T00:00:00 | +----------------------------------------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -211,15 +175,6 @@ SELECT max(number) - min(number) as maxmin, date_bin(INTERVAL '1 second', ts, '2 Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -244,15 +199,6 @@ SELECT maxmin, time_window FROM out_num_cnt_df_func; | 2 | 2021-07-01T00:00:00 | +--------+---------------------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -307,15 +253,6 @@ SELECT date_trunc('second', ts) as time_window, sum(number) as sum_num FROM numb Affected Rows: 0 --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -340,15 +277,6 @@ SELECT time_window, sum_num FROM out_num_cnt; | 2021-07-01T00:00:00 | 42 | +---------------------+---------+ --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - -+------------------------------------------+ -| ADMIN FLUSH_FLOW('test_numbers_df_func') | -+------------------------------------------+ -| FLOW_FLUSHED | -+------------------------------------------+ - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), diff --git a/tests/cases/standalone/common/flow/flow_call_df_func.sql b/tests/cases/standalone/common/flow/flow_call_df_func.sql index 21a09005960..d32f2b7b921 100644 --- a/tests/cases/standalone/common/flow/flow_call_df_func.sql +++ b/tests/cases/standalone/common/flow/flow_call_df_func.sql @@ -11,9 +11,6 @@ SINK TO out_num_cnt_df_func AS SELECT sum(abs(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00') as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -26,9 +23,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); -- note that this quote-unquote column is a column-name, **not** a aggregation expr, generated by datafusion SELECT "sum(abs(numbers_input_df_func.number))", time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -57,9 +51,6 @@ SINK TO out_num_cnt_df_func AS SELECT abs(sum(number)), date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00') as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (-20, "2021-07-01 00:00:00.200"), @@ -71,9 +62,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT "abs(sum(numbers_input_df_func.number))", time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -101,9 +89,6 @@ SINK TO out_num_cnt_df_func AS SELECT max(number) - min(number) as maxmin, date_bin(INTERVAL '1 second', ts, '2021-07-01 00:00:00'::Timestamp) as time_window FROM numbers_input_df_func GROUP BY time_window; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -114,9 +99,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT maxmin, time_window FROM out_num_cnt_df_func; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), @@ -145,9 +127,6 @@ SINK TO out_num_cnt AS SELECT date_trunc('second', ts) as time_window, sum(number) as sum_num FROM numbers_input_df_func GROUP BY date_trunc('second', ts); --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (20, "2021-07-01 00:00:00.200"), @@ -158,9 +137,6 @@ ADMIN FLUSH_FLOW('test_numbers_df_func'); SELECT time_window, sum_num FROM out_num_cnt; --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_numbers_df_func'); - INSERT INTO numbers_input_df_func VALUES (23,"2021-07-01 00:00:01.000"), diff --git a/tests/cases/standalone/common/flow/flow_null.result b/tests/cases/standalone/common/flow/flow_null.result index 4ee94cfc099..2386d676cd5 100644 --- a/tests/cases/standalone/common/flow/flow_null.result +++ b/tests/cases/standalone/common/flow/flow_null.result @@ -66,18 +66,6 @@ FROM | 600 | 2024-10-18T19:01:30 | +---------+---------------------+ --- Test if FLOWS table works, but don't care about the result since it vary from runs -SELECT - count(CASE WHEN state_size > 0 THEN 1 ELSE 0 END) as active_flows, -FROM - INFORMATION_SCHEMA.FLOWS; - -+--------------+ -| active_flows | -+--------------+ -| 1 | -+--------------+ - DROP FLOW requests_long_term; Affected Rows: 0 diff --git a/tests/cases/standalone/common/flow/flow_null.sql b/tests/cases/standalone/common/flow/flow_null.sql index 680d404b781..4cc71628092 100644 --- a/tests/cases/standalone/common/flow/flow_null.sql +++ b/tests/cases/standalone/common/flow/flow_null.sql @@ -44,12 +44,6 @@ SELECT FROM sum_val_in_reqs; --- Test if FLOWS table works, but don't care about the result since it vary from runs -SELECT - count(CASE WHEN state_size > 0 THEN 1 ELSE 0 END) as active_flows, -FROM - INFORMATION_SCHEMA.FLOWS; - DROP FLOW requests_long_term; DROP TABLE sum_val_in_reqs; diff --git a/tests/cases/standalone/common/flow/flow_rebuild.result b/tests/cases/standalone/common/flow/flow_rebuild.result index 79c31fc8b11..a0b71b02c4f 100644 --- a/tests/cases/standalone/common/flow/flow_rebuild.result +++ b/tests/cases/standalone/common/flow/flow_rebuild.result @@ -349,73 +349,7 @@ SELECT wildcard FROM out_basic; | 2 | +----------+ -DROP TABLE input_basic; - -Affected Rows: 0 - -DROP TABLE out_basic; - -Affected Rows: 0 - -DROP FLOW test_wildcard_basic; - -Affected Rows: 0 - -- combination of different order of rebuild input table/flow -CREATE TABLE input_basic ( - "number" INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY(number), - TIME INDEX(ts) -); - -Affected Rows: 0 - -CREATE FLOW test_wildcard_basic sink TO out_basic EVAL INTERVAL '1m' AS -SELECT - COUNT(*) as wildcard -FROM - input_basic; - -Affected Rows: 0 - --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s -INSERT INTO - input_basic -VALUES - (23, "2021-07-01 00:00:01.000"), - (24, "2021-07-01 00:00:01.500"); - -Affected Rows: 2 - --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_wildcard_basic'); - -+-----------------------------------------+ -| ADMIN FLUSH_FLOW('test_wildcard_basic') | -+-----------------------------------------+ -| FLOW_FLUSHED | -+-----------------------------------------+ - -SELECT wildcard FROM out_basic; - -+----------+ -| wildcard | -+----------+ -| 2 | -+----------+ - DROP TABLE input_basic; Affected Rows: 0 @@ -557,16 +491,6 @@ FROM Affected Rows: 0 --- SQLNESS ARG restart=true -SELECT 1; - -+----------+ -| Int64(1) | -+----------+ -| 1 | -+----------+ - --- SQLNESS SLEEP 3s INSERT INTO input_basic VALUES @@ -575,8 +499,6 @@ VALUES Affected Rows: 2 --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_wildcard_basic'); diff --git a/tests/cases/standalone/common/flow/flow_rebuild.sql b/tests/cases/standalone/common/flow/flow_rebuild.sql index d7d4f43749f..bdafed019d3 100644 --- a/tests/cases/standalone/common/flow/flow_rebuild.sql +++ b/tests/cases/standalone/common/flow/flow_rebuild.sql @@ -189,44 +189,8 @@ ADMIN FLUSH_FLOW('test_wildcard_basic'); SELECT wildcard FROM out_basic; -DROP TABLE input_basic; - -DROP TABLE out_basic; - -DROP FLOW test_wildcard_basic; - -- combination of different order of rebuild input table/flow -CREATE TABLE input_basic ( - "number" INT, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY(number), - TIME INDEX(ts) -); - -CREATE FLOW test_wildcard_basic sink TO out_basic EVAL INTERVAL '1m' AS -SELECT - COUNT(*) as wildcard -FROM - input_basic; - --- SQLNESS ARG restart=true -SELECT 1; - --- SQLNESS SLEEP 3s -INSERT INTO - input_basic -VALUES - (23, "2021-07-01 00:00:01.000"), - (24, "2021-07-01 00:00:01.500"); - --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | -ADMIN FLUSH_FLOW('test_wildcard_basic'); - -SELECT wildcard FROM out_basic; - DROP TABLE input_basic; CREATE TABLE input_basic ( @@ -302,18 +266,12 @@ SELECT FROM input_basic; --- SQLNESS ARG restart=true -SELECT 1; - --- SQLNESS SLEEP 3s INSERT INTO input_basic VALUES (23, "2021-07-01 00:00:01.000"), (24, "2021-07-01 00:00:01.500"); --- give flownode a second to rebuild flow --- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('test_wildcard_basic'); diff --git a/tests/cases/standalone/common/flow/flow_view.result b/tests/cases/standalone/common/flow/flow_view.result index f823328d507..ab2577f93ae 100644 --- a/tests/cases/standalone/common/flow/flow_view.result +++ b/tests/cases/standalone/common/flow/flow_view.result @@ -10,12 +10,10 @@ INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512 Affected Rows: 4 --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); Affected Rows: 4 --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); Affected Rows: 4 diff --git a/tests/cases/standalone/common/flow/flow_view.sql b/tests/cases/standalone/common/flow/flow_view.sql index 420b9e84f4f..b10ef090908 100644 --- a/tests/cases/standalone/common/flow/flow_view.sql +++ b/tests/cases/standalone/common/flow/flow_view.sql @@ -4,10 +4,8 @@ CREATE TABLE user_agent_statistics ( user_agent STRING, total_count INT64, updat INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); --- SQLNESS SLEEP 1s INSERT INTO ngx_access_log VALUES ('192.168.1.1', 'GET', '/index.html', 200, 512, 'Mozilla/5.0', 1024, '2023-10-01T10:00:00Z'), ('192.168.1.2', 'POST', '/submit', 201, 256, 'curl/7.68.0', 512, '2023-10-01T10:01:00Z'), ('192.168.1.1', 'GET', '/about.html', 200, 128, 'Mozilla/5.0', 256, '2023-10-01T10:02:00Z'), ('192.168.1.3', 'GET', '/contact', 404, 64, 'curl/7.68.0', 128, '2023-10-01T10:03:00Z'); CREATE FLOW user_agent_flow SINK TO user_agent_statistics EVAL INTERVAL '1m' AS SELECT user_agent, COUNT(user_agent) AS total_count FROM ngx_access_log GROUP BY user_agent; diff --git a/tests/cases/standalone/common/function/admin/build_index_table.result b/tests/cases/standalone/common/function/admin/build_index_table.result deleted file mode 100644 index 4951df9e0da..00000000000 --- a/tests/cases/standalone/common/function/admin/build_index_table.result +++ /dev/null @@ -1,72 +0,0 @@ -CREATE TABLE test ( - ts TIMESTAMP TIME INDEX, - msg TEXT, -); - -Affected Rows: 0 - -INSERT INTO test VALUES -(1,"The quick brown fox jumps over the lazy dog"), -(2,"The quick brown fox jumps over the lazy cat"), -(3,"The quick brown fox jumps over the lazy mouse"), -(4,"The quick brown fox jumps over the lazy rabbit"), -(5,"The quick brown fox jumps over the lazy turtle"); - -Affected Rows: 5 - -SELECT * FROM test; - -+-------------------------+------------------------------------------------+ -| ts | msg | -+-------------------------+------------------------------------------------+ -| 1970-01-01T00:00:00.001 | The quick brown fox jumps over the lazy dog | -| 1970-01-01T00:00:00.002 | The quick brown fox jumps over the lazy cat | -| 1970-01-01T00:00:00.003 | The quick brown fox jumps over the lazy mouse | -| 1970-01-01T00:00:00.004 | The quick brown fox jumps over the lazy rabbit | -| 1970-01-01T00:00:00.005 | The quick brown fox jumps over the lazy turtle | -+-------------------------+------------------------------------------------+ - -ADMIN FLUSH_TABLE('test'); - -+---------------------------+ -| ADMIN FLUSH_TABLE('test') | -+---------------------------+ -| 0 | -+---------------------------+ - --- SQLNESS SLEEP 1s --- No fulltext index yet -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -+------------+ -| index_size | -+------------+ -| 0 | -+------------+ - -ALTER TABLE test MODIFY COLUMN msg SET FULLTEXT INDEX; - -Affected Rows: 0 - -ADMIN BUILD_INDEX('test'); - -+---------------------------+ -| ADMIN BUILD_INDEX('test') | -+---------------------------+ -| 0 | -+---------------------------+ - --- SQLNESS SLEEP 1s --- Fulltext index built -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -+------------+ -| index_size | -+------------+ -| 318 | -+------------+ - -DROP TABLE test; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/function/admin/build_index_table.sql b/tests/cases/standalone/common/function/admin/build_index_table.sql deleted file mode 100644 index da3b84f23da..00000000000 --- a/tests/cases/standalone/common/function/admin/build_index_table.sql +++ /dev/null @@ -1,29 +0,0 @@ -CREATE TABLE test ( - ts TIMESTAMP TIME INDEX, - msg TEXT, -); - -INSERT INTO test VALUES -(1,"The quick brown fox jumps over the lazy dog"), -(2,"The quick brown fox jumps over the lazy cat"), -(3,"The quick brown fox jumps over the lazy mouse"), -(4,"The quick brown fox jumps over the lazy rabbit"), -(5,"The quick brown fox jumps over the lazy turtle"); - -SELECT * FROM test; - -ADMIN FLUSH_TABLE('test'); - --- SQLNESS SLEEP 1s --- No fulltext index yet -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -ALTER TABLE test MODIFY COLUMN msg SET FULLTEXT INDEX; - -ADMIN BUILD_INDEX('test'); - --- SQLNESS SLEEP 1s --- Fulltext index built -SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS; - -DROP TABLE test; diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.result b/tests/cases/standalone/common/function/admin/build_index_table_restart.result index 78451e3121c..f6e02582189 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.result +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.result @@ -14,6 +14,18 @@ INSERT INTO build_index_restart_test VALUES Affected Rows: 5 +SELECT * FROM build_index_restart_test; + ++-------------------------+------------------------------------------------+ +| ts | msg | ++-------------------------+------------------------------------------------+ +| 1970-01-01T00:00:00.001 | The quick brown fox jumps over the lazy dog | +| 1970-01-01T00:00:00.002 | The quick brown fox jumps over the lazy cat | +| 1970-01-01T00:00:00.003 | The quick brown fox jumps over the lazy mouse | +| 1970-01-01T00:00:00.004 | The quick brown fox jumps over the lazy rabbit | +| 1970-01-01T00:00:00.005 | The quick brown fox jumps over the lazy turtle | ++-------------------------+------------------------------------------------+ + ADMIN FLUSH_TABLE('build_index_restart_test'); +-----------------------------------------------+ @@ -32,6 +44,21 @@ SHOW INDEX FROM build_index_restart_test; | build_index_restart_test | 1 | TIME INDEX | 1 | ts | A | | | | | TIME | | | YES | | +--------------------------+------------+------------+--------------+-------------+-----------+-------------+----------+--------+------+------------+---------+---------------+---------+------------+ +-- Index size is 0 before the fulltext index is built +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + ++------------+ +| index_size | ++------------+ +| 0 | ++------------+ + -- No physical fulltext index metadata yet. -- ssts_index_meta.index_type uses physical backend names (fulltext_bloom/fulltext_tantivy), -- unlike SHOW INDEX's greptime-fulltext-index-* display value. @@ -108,6 +135,21 @@ AND index_type LIKE 'fulltext%'; | 1 | +---------------------------+ +-- Fulltext index size is reported after the build +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + ++------------+ +| index_size | ++------------+ +| 318 | ++------------+ + -- SQLNESS ARG restart=true -- After restart, verify fulltext query still works SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; diff --git a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql index b5503ed18e2..745e98eee2c 100644 --- a/tests/cases/standalone/common/function/admin/build_index_table_restart.sql +++ b/tests/cases/standalone/common/function/admin/build_index_table_restart.sql @@ -10,12 +10,23 @@ INSERT INTO build_index_restart_test VALUES (4,"The quick brown fox jumps over the lazy rabbit"), (5,"The quick brown fox jumps over the lazy turtle"); +SELECT * FROM build_index_restart_test; + ADMIN FLUSH_TABLE('build_index_restart_test'); -- SQLNESS SLEEP 1s -- No fulltext index yet SHOW INDEX FROM build_index_restart_test; +-- Index size is 0 before the fulltext index is built +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + -- No physical fulltext index metadata yet. -- ssts_index_meta.index_type uses physical backend names (fulltext_bloom/fulltext_tantivy), -- unlike SHOW INDEX's greptime-fulltext-index-* display value. @@ -55,6 +66,15 @@ WHERE table_id = ( ) AND index_type LIKE 'fulltext%'; +-- Fulltext index size is reported after the build +SELECT index_size FROM INFORMATION_SCHEMA.REGION_STATISTICS +WHERE table_id = ( + SELECT table_id + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'build_index_restart_test' +); + -- SQLNESS ARG restart=true -- After restart, verify fulltext query still works SELECT msg FROM build_index_restart_test WHERE MATCHES(msg, 'fox') ORDER BY ts; diff --git a/tests/cases/standalone/common/information_schema/region_statistics.result b/tests/cases/standalone/common/information_schema/region_statistics.result index e6f9c9e47fc..5589b5d00e2 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.result +++ b/tests/cases/standalone/common/information_schema/region_statistics.result @@ -21,6 +21,34 @@ INSERT INTO test VALUES Affected Rows: 3 +-- The projected columns are millisecond-typed. A nanosecond time index proves +-- min_timestamp floors and max_timestamp ceils, so the reported window stays a +-- superset of the data and cannot hide a region from a time-bounded lookup. +CREATE TABLE precise ( + a int primary key, + ts timestamp(9) time index, +); + +Affected Rows: 0 + +-- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. +INSERT INTO precise VALUES + (1, 1500000), + (2, 2500000), + (3, 3999999); + +Affected Rows: 3 + +-- A region with no rows has nothing to bound, so both ends stay NULL rather +-- than collapsing onto the epoch. +CREATE TABLE empty_region ( + a int primary key, + ts timestamp time index, +); + +Affected Rows: 0 + +-- Wait once for the region statistics of all three tables to be reported. -- SQLNESS SLEEP 3s -- SQLNESS REPLACE (\s+\d+\s+) -- For regions using different WAL implementations, the manifest size may vary. @@ -49,25 +77,6 @@ DROP TABLE test; Affected Rows: 0 --- The projected columns are millisecond-typed. A nanosecond time index proves --- min_timestamp floors and max_timestamp ceils, so the reported window stays a --- superset of the data and cannot hide a region from a time-bounded lookup. -CREATE TABLE precise ( - a int primary key, - ts timestamp(9) time index, -); - -Affected Rows: 0 - --- 1500000ns = 1.5ms floors to 1ms; 3999999ns = 3.999999ms ceils to 4ms. -INSERT INTO precise VALUES - (1, 1500000), - (2, 2500000), - (3, 3999999); - -Affected Rows: 3 - --- SQLNESS SLEEP 3s SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); @@ -82,16 +91,6 @@ DROP TABLE precise; Affected Rows: 0 --- A region with no rows has nothing to bound, so both ends stay NULL rather --- than collapsing onto the epoch. -CREATE TABLE empty_region ( - a int primary key, - ts timestamp time index, -); - -Affected Rows: 0 - --- SQLNESS SLEEP 3s SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); diff --git a/tests/cases/standalone/common/information_schema/region_statistics.sql b/tests/cases/standalone/common/information_schema/region_statistics.sql index 60264fbfb57..6ce12b96417 100644 --- a/tests/cases/standalone/common/information_schema/region_statistics.sql +++ b/tests/cases/standalone/common/information_schema/region_statistics.sql @@ -16,20 +16,6 @@ INSERT INTO test VALUES (11, 'b', 11), (21, 'c', 21); --- SQLNESS SLEEP 3s --- SQLNESS REPLACE (\s+\d+\s+) --- For regions using different WAL implementations, the manifest size may vary. --- The remote WAL implementation additionally stores a flushed entry ID when creating the manifest. -SELECT SUM(region_rows), SUM(written_bytes_since_open), SUM(query_cpu_time_millis), - SUM(query_scanned_bytes), SUM(memtable_size), SUM(sst_size), SUM(index_size) - FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id - IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test' and table_schema = 'public'); - --- SQLNESS REPLACE (\s+\d+\s+) -SELECT data_length, index_length, avg_row_length, table_rows FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test'; - -DROP TABLE test; - -- The projected columns are millisecond-typed. A nanosecond time index proves -- min_timestamp floors and max_timestamp ceils, so the reported window stays a -- superset of the data and cannot hide a region from a time-bounded lookup. @@ -44,13 +30,6 @@ INSERT INTO precise VALUES (2, 2500000), (3, 3999999); --- SQLNESS SLEEP 3s -SELECT min_timestamp, max_timestamp - FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id - IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); - -DROP TABLE precise; - -- A region with no rows has nothing to bound, so both ends stay NULL rather -- than collapsing onto the epoch. CREATE TABLE empty_region ( @@ -58,7 +37,27 @@ CREATE TABLE empty_region ( ts timestamp time index, ); +-- Wait once for the region statistics of all three tables to be reported. -- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (\s+\d+\s+) +-- For regions using different WAL implementations, the manifest size may vary. +-- The remote WAL implementation additionally stores a flushed entry ID when creating the manifest. +SELECT SUM(region_rows), SUM(written_bytes_since_open), SUM(query_cpu_time_millis), + SUM(query_scanned_bytes), SUM(memtable_size), SUM(sst_size), SUM(index_size) + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test' and table_schema = 'public'); + +-- SQLNESS REPLACE (\s+\d+\s+) +SELECT data_length, index_length, avg_row_length, table_rows FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'test'; + +DROP TABLE test; + +SELECT min_timestamp, max_timestamp + FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id + IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'precise' and table_schema = 'public'); + +DROP TABLE precise; + SELECT min_timestamp, max_timestamp FROM INFORMATION_SCHEMA.REGION_STATISTICS WHERE table_id IN (SELECT TABLE_ID FROM INFORMATION_SCHEMA.TABLES WHERE table_name = 'empty_region' and table_schema = 'public'); diff --git a/tests/cases/standalone/common/skip_wal.result b/tests/cases/standalone/common/skip_wal.result index 25c18fed20e..792660b2f9a 100644 --- a/tests/cases/standalone/common/skip_wal.result +++ b/tests/cases/standalone/common/skip_wal.result @@ -19,50 +19,6 @@ VALUES Affected Rows: 3 --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -+------+-----+----------+-------------+-----------+----+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+------+-----+----------+-------------+-----------+----+ -+------+-----+----------+-------------+-----------+----+ - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -ADMIN flush_table('system_metrics'); - -+-------------------------------------+ -| ADMIN flush_table('system_metrics') | -+-------------------------------------+ -| 0 | -+-------------------------------------+ - --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - --- A table created without a real WAL provider cannot enable WAL later. -ALTER TABLE system_metrics SET 'skip_wal' = 'false'; - -Error: 1001(Unsupported), Unsupported operation setting skip_wal = 'false' without an existing WAL provider - -DROP TABLE system_metrics; - -Affected Rows: 0 - CREATE TABLE alter_skip_wal ( host STRING, val DOUBLE, @@ -152,8 +108,15 @@ SELECT * FROM alter_skip_wal ORDER BY ts; | host3 | 3.0 | 1970-01-01T00:00:03 | +-------+-----+---------------------+ --- Re-enabling WAL does not retroactively write the skipped row to WAL. -- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + ++------+-----+----------+-------------+-----------+----+ +| host | idc | cpu_util | memory_util | disk_util | ts | ++------+-----+----------+-------------+-----------+----+ ++------+-----+----------+-------------+-----------+----+ + +-- Re-enabling WAL does not retroactively write the skipped row to WAL. SHOW CREATE TABLE alter_skip_wal; +----------------+-----------------------------------------------+ @@ -186,3 +149,39 @@ DROP TABLE alter_skip_wal; Affected Rows: 0 +INSERT INTO system_metrics +VALUES + ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), + ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), + ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); + +Affected Rows: 3 + +ADMIN flush_table('system_metrics'); + ++-------------------------------------+ +| ADMIN flush_table('system_metrics') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +-- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + ++-------+-------+----------+-------------+-----------+-------------------------+ +| host | idc | cpu_util | memory_util | disk_util | ts | ++-------+-------+----------+-------------+-----------+-------------------------+ +| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | +| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | +| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | ++-------+-------+----------+-------------+-----------+-------------------------+ + +-- A table created without a real WAL provider cannot enable WAL later. +ALTER TABLE system_metrics SET 'skip_wal' = 'false'; + +Error: 1001(Unsupported), Unsupported operation setting skip_wal = 'false' without an existing WAL provider + +DROP TABLE system_metrics; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/skip_wal.sql b/tests/cases/standalone/common/skip_wal.sql index 2e9cac1f272..e0bfa4b41d8 100644 --- a/tests/cases/standalone/common/skip_wal.sql +++ b/tests/cases/standalone/common/skip_wal.sql @@ -15,25 +15,6 @@ VALUES ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -ADMIN flush_table('system_metrics'); - --- SQLNESS ARG restart=true -SELECT * FROM system_metrics; - --- A table created without a real WAL provider cannot enable WAL later. -ALTER TABLE system_metrics SET 'skip_wal' = 'false'; - -DROP TABLE system_metrics; - CREATE TABLE alter_skip_wal ( host STRING, val DOUBLE, @@ -65,10 +46,28 @@ INSERT INTO alter_skip_wal VALUES ('host3', 3, 3000); SELECT * FROM alter_skip_wal ORDER BY ts; --- Re-enabling WAL does not retroactively write the skipped row to WAL. -- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + +-- Re-enabling WAL does not retroactively write the skipped row to WAL. SHOW CREATE TABLE alter_skip_wal; SELECT * FROM alter_skip_wal ORDER BY ts; DROP TABLE alter_skip_wal; + +INSERT INTO system_metrics +VALUES + ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), + ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), + ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); + +ADMIN flush_table('system_metrics'); + +-- SQLNESS ARG restart=true +SELECT * FROM system_metrics; + +-- A table created without a real WAL provider cannot enable WAL later. +ALTER TABLE system_metrics SET 'skip_wal' = 'false'; + +DROP TABLE system_metrics; diff --git a/tests/cases/standalone/common/ttl/ttl_instant.result b/tests/cases/standalone/common/ttl/ttl_instant.result index 110fd92d92b..d6a3921d245 100644 --- a/tests/cases/standalone/common/ttl/ttl_instant.result +++ b/tests/cases/standalone/common/ttl/ttl_instant.result @@ -45,7 +45,6 @@ ORDER BY +-----+ +-----+ --- SQLNESS SLEEP 2s ADMIN flush_table('test_ttl'); +-------------------------------+ diff --git a/tests/cases/standalone/common/ttl/ttl_instant.sql b/tests/cases/standalone/common/ttl/ttl_instant.sql index 57bc4b625ea..be4831ba2b4 100644 --- a/tests/cases/standalone/common/ttl/ttl_instant.sql +++ b/tests/cases/standalone/common/ttl/ttl_instant.sql @@ -20,7 +20,6 @@ from ORDER BY val; --- SQLNESS SLEEP 2s ADMIN flush_table('test_ttl'); ADMIN compact_table('test_ttl'); diff --git a/tests/runner/src/cmd/bare.rs b/tests/runner/src/cmd/bare.rs index 24ba756cee6..dd54d1d2d9b 100644 --- a/tests/runner/src/cmd/bare.rs +++ b/tests/runner/src/cmd/bare.rs @@ -112,7 +112,53 @@ pub struct BareCommand { enable_gc: bool, } +#[derive(Debug, PartialEq, Eq)] +struct Parallelism { + /// How many environments (standalone, distributed) run at the same time. + env: usize, + /// How many database instances run at the same time within one environment. + instance: usize, +} + impl BareCommand { + fn parallelism(&self) -> Parallelism { + let self_managed_kafka = + matches!(self.wal, Wal::Kafka) && self.kafka_wal_broker_endpoints.is_none(); + // External servers are shared by both environments, a filtered run is for + // debugging a single case, and each environment starts and tears down the + // self-managed Kafka cluster on its own. + if self.server_addr.server_addr.is_some() + || self.config.test_filter != ".*" + || self_managed_kafka + { + return Parallelism { + env: 1, + instance: 1, + }; + } + + // Instances of one environment would share the same metadata table and + // Kafka topics. Only distributed uses the kv backend, and the two + // environments use different Kafka topic prefixes, so the environments + // themselves can still run concurrently. + if self.setup_etcd + || !self.store_addrs.is_empty() + || self.setup_pg.is_some() + || self.setup_mysql.is_some() + || self.kafka_wal_broker_endpoints.is_some() + { + return Parallelism { + env: self.jobs, + instance: 1, + }; + } + + Parallelism { + env: self.jobs, + instance: self.jobs, + } + } + pub async fn run(mut self) { let temp_dir = tempfile::Builder::new() .prefix("sqlness") @@ -134,22 +180,11 @@ impl BareCommand { if self.jobs == 0 { self.jobs = num_cpus::get() / 2; } - - // normalize parallelism to 1 if any of the following conditions are met: - // Note: parallelism in pg and mysql is possible, but need configuration. - if self.server_addr.server_addr.is_some() - || self.setup_etcd - || self.setup_pg.is_some() - || self.setup_mysql.is_some() - || matches!(self.wal, Wal::Kafka) - || self.kafka_wal_broker_endpoints.is_some() - || self.config.test_filter != ".*" - { - self.jobs = 1; - println!( - "Normalizing parallelism to 1 due to server addresses, etcd/pg/mysql/kafka setup, or test filter usage" - ); - } + let parallelism = self.parallelism(); + println!( + "Environment parallelism: {}, instance parallelism per environment: {}", + parallelism.env, parallelism.instance + ); let config = ConfigBuilder::default() .case_dir(util::get_case_dir(self.config.case_dir)) @@ -158,8 +193,8 @@ impl BareCommand { .follow_links(true) .env_config_file(self.config.env_config_file) .interceptor_registry(interceptor_registry) - .parallelism(self.jobs) - .env_parallelism(self.jobs) + .parallelism(parallelism.instance) + .env_parallelism(parallelism.env) .build() .unwrap(); @@ -216,3 +251,61 @@ impl BareCommand { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn parallelism(jobs: &str, args: &[&str]) -> Parallelism { + let cmd = BareCommand::try_parse_from(["bare", "-j", jobs].iter().chain(args)).unwrap(); + cmd.parallelism() + } + + #[test] + fn test_parallelism() { + let serial = Parallelism { + env: 1, + instance: 1, + }; + let env_only = Parallelism { + env: 4, + instance: 1, + }; + + assert_eq!( + parallelism("4", &[]), + Parallelism { + env: 4, + instance: 4 + } + ); + assert_eq!(parallelism("4", &["--setup-pg"]), env_only); + assert_eq!(parallelism("4", &["--setup-mysql"]), env_only); + assert_eq!(parallelism("4", &["--setup-etcd"]), env_only); + assert_eq!( + parallelism("4", &["--store-addrs", "127.0.0.1:2379"]), + env_only + ); + assert_eq!( + parallelism("4", &["-w", "kafka", "-k", "127.0.0.1:9092"]), + env_only + ); + assert_eq!(parallelism("4", &["-w", "kafka"]), serial); + assert_eq!(parallelism("4", &["-t", "basic"]), serial); + assert_eq!( + parallelism( + "4", + &[ + "-s", + "127.0.0.1:4001", + "-p", + "127.0.0.1:4003", + "-m", + "127.0.0.1:4002" + ] + ), + serial + ); + assert_eq!(parallelism("1", &["--setup-pg"]), serial); + } +} From abedeb2beadf982e8c2c0979221ba72a75a6e71f Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 24 Sep 2026 08:18:41 +0000 Subject: [PATCH 033/133] fix(query): keep count_values generated label in enclosing expressions (#9223) * fix(query): keep count_values generated label in enclosing expressions `count_values("v", m)` projects the generated label as a real output column, but did not register it in `ctx.tag_columns`. Enclosing expressions (abs/round/+1/topk/label_replace/vector join) rebuild their projection from `ctx.tag_columns` and silently drop the label. Register the generated label in `ctx.tag_columns` after the projection, and give it the same qualifier as other tag columns so qualified references resolve. PromQL overwrites an input label with the same name, so drop any existing tag with that name first to avoid duplicate column ambiguity. Fixes https://github.com/GreptimeTeam/greptimedb/issues/9181 Report: .e-agent/greptimedb_promql_compatibility_report_2026-09-16.md P0-1 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): use total_cmp in count_values test helper Silence clippy::needless_borrow on partial_cmp(&right.1); f64 sorting uses total_cmp, matching the other planner test helpers. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): keep count_values generated label as sink table primary key FindGroupByFinalName::f_up only renamed a group key when the projection aliased the key column directly. count_values now projects its generated label as a unary UDF over the sampled column (prom_float_to_string(value) AS label) while the aggregate still groups by the raw column, so the name match failed and the sink table lost the label from its PRIMARY KEY, demoting it to a DOUBLE value column. Allow a projection above the aggregate to rename a group key by deriving its output from a single group-key column (unary scalar function or cast over that column only). Multi-column expressions, case, aggregates, windows, subqueries and literals are still rejected so a computed column cannot be mistaken for the group key. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): group count_values by the formatted sample value The numeric count_values branch grouped by the raw sample column and only formatted the value into the generated label in the post-aggregate projection. Two distinct raw values that collapse to one label text (e.g. BIGINT 9007199254740992 and 9007199254740993, both 9007199254740992 in Float64) were split into two groups, each emitting the same label set at one timestamp, violating Prometheus' unique-label-set-per-timestamp invariant. Return the formatted value expression (prom_float_to_string, with a CAST to Float64 for non-Float64 inputs) as previous_field_expressions so the aggregate groups by the same expression that produces the label, mirroring the existing mixed float/native-histogram precedent. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): drop needless borrow in count_values test helper Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(flow): don't replace group key with derived unary expr FindGroupByFinalName treated any unary expression of a single group column (scalar fn, Cast, TryCast) as a rename of that group key and swapped it in as the sink primary key. A derived expression such as lower(host) is many-to-one, so distinct groups (HOST_A vs host_a) would collapse to the same primary key and be silently deduplicated. Narrow the matching to direct name-matched aliases of the actual group expression only, and drop the derived-unary path (is_unary_expr_of_column) and the allow_derived flag. count_values already groups by the formatted expression name, so it is unaffected. Add a regression asserting host survives and host_lc does not replace it under both optimizer settings. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/flow/src/batching_mode/table_creator.rs | 123 +++- src/flow/src/batching_mode/utils.rs | 26 +- src/flow/src/batching_mode/utils/test.rs | 86 +++ src/query/src/promql/planner.rs | 534 +++++++++++++++++- src/servers/tests/http/mod.rs | 1 + .../tests/http/prom_count_values_test.rs | 192 +++++++ .../distributed/flow-tql/flow_tql.result | 16 +- .../common/promql/count_values.result | 139 +++++ .../standalone/common/promql/count_values.sql | 46 ++ .../cases/standalone/flow-tql/flow_tql.result | 16 +- 10 files changed, 1128 insertions(+), 51 deletions(-) create mode 100644 src/servers/tests/http/prom_count_values_test.rs diff --git a/src/flow/src/batching_mode/table_creator.rs b/src/flow/src/batching_mode/table_creator.rs index 73a17f94967..778ee939a8e 100644 --- a/src/flow/src/batching_mode/table_creator.rs +++ b/src/flow/src/batching_mode/table_creator.rs @@ -254,15 +254,22 @@ mod test { use std::sync::Arc; use api::v1::column_def::try_as_column_schema; + use catalog::RegisterTableRequest; + use catalog::memory::new_memory_catalog_manager; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use datafusion::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit, }; use datafusion_common::DFSchema; use datafusion_expr::logical_plan::EmptyRelation; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::ColumnSchema; + use datatypes::schema::{ColumnSchema, Schema}; use pretty_assertions::assert_eq; + use query::options::QueryOptions; + use query::{QueryEngineFactory, QueryEngineRef}; use session::context::QueryContext; + use table::metadata::{TableInfoBuilder, TableMetaBuilder}; + use table::test_util::EmptyTable; use super::*; use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, AUTO_CREATED_UPDATE_AT_TS_COL}; @@ -308,6 +315,120 @@ mod test { assert!(columns[2].is_time_index()); } + /// Creates a query engine holding a Prometheus shaped table `http_requests`: tags + /// (`host`, `idc`), a single f64 value column (`val`) and a time index (`ts`), so that + /// TQL queries can be planned against it. + fn create_tql_test_query_engine() -> QueryEngineRef { + let catalog_list = new_memory_catalog_manager().unwrap(); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("idc", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ]))) + // `host` and `idc` are tags, `val` is the value column + .primary_key_indices(vec![1, 2]) + .value_indices(vec![3]) + .engine("mito".to_string()) + .next_column_id(1026) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("http_requests".to_string()) + .meta(table_meta) + .build() + .unwrap(); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "http_requests".to_string(), + table_id: 1026, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + + QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + } + + /// A TQL `count_values` flow must keep the generated label as a primary key of the auto + /// created sink table. `count_values("status_code", http_requests)` groups by the sample + /// value column and projects the label as a unary scalar expression of that column + /// (`prom_float_to_string(val) AS status_code`), so the label only derives from a group by + /// column instead of referencing it directly. + /// + /// The plan is built with the same path a flow task uses (including the DataFusion + /// optimizers, which may rewrite the shape of the alias), and the assertion covers both. + #[tokio::test] + async fn test_tql_count_values_generated_label_is_primary_key() { + let query_engine = create_tql_test_query_engine(); + let ctx = QueryContext::arc(); + + for optimize in [false, true] { + let plan = sql_to_df_plan( + ctx.clone(), + query_engine.clone(), + r#"TQL EVAL (0, 15, '5s') count_values("status_code", http_requests)"#, + optimize, + ) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + let expr = create_table_with_expr( + &plan, + &[ + "greptime".to_string(), + "public".to_string(), + "sink".to_string(), + ], + &QueryType::Tql, + ) + .unwrap(); + let columns = expr + .column_defs + .iter() + .map(|column| try_as_column_schema(column).unwrap()) + .collect::>(); + + assert_eq!( + vec!["status_code".to_string()], + expr.primary_keys, + "optimize={optimize}, plan:\n{plan_display}" + ); + assert_eq!( + "ts", expr.time_index, + "optimize={optimize}, plan:\n{plan_display}" + ); + // the aggregation output is a value column, the generated label is a tag column + assert_eq!( + "count(http_requests.val)", columns[0].name, + "optimize={optimize}, plan:\n{plan_display}" + ); + assert_eq!(ConcreteDataType::float64_datatype(), columns[0].data_type); + assert_eq!("ts", columns[1].name); + assert!(columns[1].is_time_index()); + assert_eq!("status_code", columns[2].name); + assert_eq!(ConcreteDataType::string_datatype(), columns[2].data_type); + } + } + #[tokio::test] async fn test_gen_create_table_sql() { let query_engine = create_test_query_engine(); diff --git a/src/flow/src/batching_mode/utils.rs b/src/flow/src/batching_mode/utils.rs index 0220f4ef48d..267fa170cdd 100644 --- a/src/flow/src/batching_mode/utils.rs +++ b/src/flow/src/batching_mode/utils.rs @@ -1164,6 +1164,16 @@ pub fn df_plan_to_sql(plan: &LogicalPlan) -> Result { Ok(sql.to_string()) } +/// Returns whether `expr` directly renames the group by expression `group_expr`. +/// +/// Renames are identified only by expression name. Derived unary expressions are not group +/// keys, even when they reference a single group key column. A PromQL `count_values("label", +/// metric)` plan lands here too: it groups by the formatted sample value +/// (`prom_float_to_string(value)`) and projects that group key column as the generated label. +fn is_alias_of_group_expr(group_expr: &Expr, expr: &Expr) -> DfResult { + Ok(group_expr.name_for_alias()? == expr.name_for_alias()?) +} + /// Helper to find the innermost group by expr in schema, return None if no group by expr #[derive(Debug, Clone, Default)] pub struct FindGroupByFinalName { @@ -1220,24 +1230,28 @@ impl TreeNodeVisitor<'_> for FindGroupByFinalName { Ok(TreeNodeRecursion::Continue) } - /// deal with projection when going up with group exprs + /// Applies only direct, name-matched aliases when propagating group expressions upward. fn f_up(&mut self, node: &Self::Node) -> datafusion_common::Result { if let LogicalPlan::Projection(projection) = node { + // Only direct aliases of group expressions rename group keys. Derived unary + // expressions remain ordinary projected columns. for expr in &projection.expr { let Some(group_exprs) = &mut self.group_exprs else { return Ok(TreeNodeRecursion::Continue); }; if let datafusion_expr::Expr::Alias(alias) = expr { // if a alias exist, replace with the new alias - let mut new_group_exprs = group_exprs.clone(); + let mut renamed_group_expr = None; for group_expr in group_exprs.iter() { - if group_expr.name_for_alias()? == alias.expr.name_for_alias()? { - new_group_exprs.remove(group_expr); - new_group_exprs.insert(expr.clone()); + if is_alias_of_group_expr(group_expr, alias.expr.as_ref())? { + renamed_group_expr = Some(group_expr.clone()); break; } } - *group_exprs = new_group_exprs; + if let Some(group_expr) = renamed_group_expr { + group_exprs.remove(&group_expr); + group_exprs.insert(expr.clone()); + } } } } diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 6315ae3d679..1f390d1594c 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -1033,6 +1033,34 @@ async fn test_find_group_by_exprs() { } } +#[tokio::test] +async fn test_find_group_by_exprs_does_not_replace_group_key_with_derived_expr() { + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + let sql = "SELECT host, lower(host) AS host_lc, ts, SUM(val) AS total \ + FROM (SELECT CAST(number AS STRING) AS host, ts, number AS val FROM numbers_with_ts) \ + GROUP BY host, ts"; + + for optimize in [false, true] { + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, optimize) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + let mut group_finder = FindGroupByFinalName::default(); + plan.visit(&mut group_finder).unwrap(); + let group_names = group_finder.get_group_expr_names().unwrap_or_default(); + + assert!( + group_names.contains("host"), + "optimize={optimize}, group keys {group_names:?} must keep host:\n{plan_display}" + ); + assert!( + !group_names.contains("host_lc"), + "optimize={optimize}, derived host_lc must not replace host:\n{plan_display}" + ); + } +} + #[tokio::test] async fn test_analyze_incremental_aggregate_plan() { let query_engine = create_test_query_engine(); @@ -1243,6 +1271,64 @@ async fn test_rewrite_incremental_aggregate_allows_alias_wrapped_scan() { assert_eq!(rewritten_fields, analysis.output_field_names); } +#[tokio::test] +async fn test_count_values_flow_plan_keeps_generated_label_as_group_key() { + // PromQL `count_values("label", metric)` groups by the *formatted* sample value + // (`prom_float_to_string(value)`), so such a flow carries a scalar UDF inside its GROUP + // BY instead of the raw value column: + // + // Projection: count(val), ts, prom_float_to_string(val) AS v + // Aggregate: groupBy=[[ts, prom_float_to_string(CAST(val AS Float64))]], aggr=[[count(val)]] + // + // Two things must keep working for that shape: the TQL flow transport (a Substrait + // encoded insert plan) and the group-key resolution that derives the sink schema, which + // is what keeps the generated label a primary key of the auto-created sink table. + let query_engine = create_test_query_engine(); + let ctx = QueryContext::arc(); + for optimize in [false, true] { + let plan = sql_to_df_plan( + ctx.clone(), + query_engine.clone(), + r#"TQL EVAL (0, 15, '5s') count_values("v", numbers_with_ts)"#, + optimize, + ) + .await + .unwrap(); + let plan_display = plan.display_indent_schema().to_string(); + + DFLogicalSubstraitConvertor {} + .encode(&plan, DefaultSerializer) + .unwrap_or_else(|err| panic!("optimize={optimize}: {err}\n{plan_display}")); + + let mut group_finder = FindGroupByFinalName::default(); + plan.visit(&mut group_finder).unwrap(); + let group_names = group_finder.get_group_expr_names().unwrap_or_default(); + assert!( + group_names.contains("v"), + "optimize={optimize}, group keys {group_names:?} must keep the generated label:\n{plan_display}" + ); + assert!( + group_names.contains("ts"), + "optimize={optimize}, group keys {group_names:?}:\n{plan_display}" + ); + + // The flow plan keeps the generated label as a group key name, and the plan shape + // (PromQL extension nodes above the aggregate) stays outside the incremental + // aggregate rewrite: `prepare_plan_for_incremental` only rewrites `QueryType::Sql` + // flows, so a `count_values` flow keeps running as a full snapshot. + let analysis = analyze_incremental_aggregate_plan(&plan).unwrap().unwrap(); + assert!( + analysis.group_key_names.contains(&"v".to_string()), + "optimize={optimize}, analysis {analysis:?}:\n{plan_display}" + ); + assert!( + !analysis.unsupported_exprs.is_empty(), + "optimize={optimize}, a count_values flow plan must not become an incremental \ + aggregate rewrite candidate: {analysis:?}\n{plan_display}" + ); + } +} + #[tokio::test] async fn test_analyze_incremental_aggregate_plan_rejects_having_filter() { let sql = diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 2142ee82371..73332c237c8 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -713,13 +713,76 @@ impl PromPlanner { let label = Self::get_param_value_as_str(*op, param)?; // `count_values` must be grouped by fields, // and project the fields to the new label. - let count_value_exprs = prev_field_exprs.iter().map(|expr| { - match expr { - DfExpr::Column(column) => DfExpr::Column(column.clone()), - _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())), - } - .alias(label) - }); + // + // The generated label is a real label (column) of the output, so it must be + // registered in `ctx.tag_columns` below. Otherwise enclosing expressions + // rebuild their projections from `ctx.tag_columns` and silently drop it. + // + // PromQL sets the generated label *before* the grouping key is built, so an + // input label with the same name is overwritten by the sample value and must + // not remain a grouping key either: samples are grouped by the generated + // label only. Dropping it from the projected tag columns is required as well: + // projecting both would emit two columns with the same name (rejected as an + // ambiguous reference). + self.ctx.tag_columns.retain(|tag| tag != label); + group_exprs.retain( + |expr| !matches!(expr, DfExpr::Column(column) if column.name == label), + ); + // The tag columns projected below are unqualified `Column` references, so they + // end up qualified with whatever qualifier they carry in the input plan. Give + // the generated label the same qualifier, otherwise qualified references to it + // (e.g. from an enclosing binary expression or a vector join) fail to resolve. + let label_qualifier = self + .ctx + .time_index_column + .as_deref() + .and_then(|time_index| { + builder + .schema() + .qualified_field_with_unqualified_name(time_index) + .ok() + }) + .and_then(|(qualifier, _)| qualifier.cloned()); + // The generated label carries the *sample value*, so it must be materialized + // as a string using Prometheus' format + // (`strconv.FormatFloat(value, 'f', -1, 64)`: shortest decimal form without + // an exponent, `1.0` becomes "1"). Labels are inferred from string columns + // when a result is converted into the Prometheus HTTP API JSON format, so a + // numeric label column would be mistaken for the sample value of the series. + let count_value_exprs = prev_field_exprs + .iter() + .map(|expr| { + let value = match expr { + DfExpr::Column(column) => { + let value = DfExpr::Column(column.clone()); + // `prom_float_to_string` formats exactly like Prometheus, + // while arrow's `Float64 -> Utf8` cast would render `1.0`. + let value = if matches!( + builder.schema().field_with_unqualified_name(&column.name), + Ok(field) if field.data_type() == &ArrowDataType::Float64 + ) { + value + } else { + DfExpr::Cast(Cast::new( + Box::new(value), + ArrowDataType::Float64, + )) + }; + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![value], + }) + } + // The value is already formatted (e.g. a native histogram is + // converted to its string form by the aggregation), and the + // aggregate output names it by its schema name. + _ => DfExpr::Column(Column::from_name( + expr.schema_name().to_string(), + )), + }; + DfExpr::Alias(Alias::new(value, label_qualifier.clone(), label)) + }) + .collect::>(); let aggregate_group_exprs = group_exprs .iter() .cloned() @@ -733,11 +796,14 @@ impl PromPlanner { .chain(Some(self.create_time_index_column_expr()?)) .chain(count_value_exprs); - builder + let builder = builder .aggregate(aggregate_group_exprs, aggr_exprs) .context(DataFusionPlanningSnafu)? .project(project_fields) - .context(DataFusionPlanningSnafu)? + .context(DataFusionPlanningSnafu)?; + // The label only exists in the output schema from here on. + self.ctx.tag_columns.push(label.to_string()); + builder } else { builder .aggregate(group_exprs.clone(), aggr_exprs) @@ -4445,9 +4511,10 @@ impl PromPlanner { /// /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where: /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields - /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty - /// only when the operation is `count_values`, as this operation requires preserving the original - /// values for grouping. + /// - `previous_field_expressions`: Field expressions naming the pre-aggregation values. This is + /// non-empty only when the operation is `count_values`, which groups by the sample value and + /// projects it as the generated label, so these expressions are passed through the same + /// formatting as that label (`prom_float_to_string`). /// fn create_aggregate_exprs( &mut self, @@ -4491,14 +4558,13 @@ impl PromPlanner { .collect::>>()?; // if the aggregator is `count_values`, it must be grouped by current fields. + // + // The grouping key is the *formatted* sample value, i.e. the same expression that + // produces the generated label below: PromQL groups by the value, and the label is + // that value in Prometheus' textual form (`strconv.FormatFloat(value, 'f', -1, 64)`), + // so grouping by the raw value would split samples that render to one label into + // several groups, each emitting the same label set for one timestamp. let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { - let prev_field_exprs: Vec<_> = self - .ctx - .field_columns - .iter() - .map(|col| DfExpr::Column(Column::from_name(col))) - .collect(); - ensure!( self.ctx.field_columns.len() == 1, UnsupportedExprSnafu { @@ -4506,7 +4572,28 @@ impl PromPlanner { } ); - prev_field_exprs + self.ctx + .field_columns + .iter() + .map(|col| { + let value = DfExpr::Column(Column::from_name(col)); + // Normalize non `Float64` inputs the same way the label projection does, + // so both sides agree on the formatted value: `prom_float_to_string` + // formats exactly like Prometheus, while arrow's `Float64 -> Utf8` cast + // would render `1.0`. + let value = if Self::field_column_type(input_plan.schema(), col) + == Some(&ArrowDataType::Float64) + { + value + } else { + DfExpr::Cast(Cast::new(Box::new(value), ArrowDataType::Float64)) + }; + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![value], + }) + }) + .collect() } else { vec![] }; @@ -7633,7 +7720,7 @@ mod test { use common_query::test_util::DummyDecoder; use common_recordbatch::RecordBatch as GreptimeRecordBatch; use datafusion::arrow::array::{ - Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, + Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; use datafusion::arrow::record_batch::RecordBatch; @@ -13081,9 +13168,9 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) .await .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ @@ -13129,10 +13216,10 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: .await .unwrap(); let expected = r#" -Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] +Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] @@ -14884,4 +14971,395 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); assert_eq!(sample_count, 2); } + + /// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: + /// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. + async fn build_count_values_table_provider() -> DfTableSourceProvider { + build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await + } + + /// Table provider with a single metric `cv_metric` holding one series per given value at + /// `ts=1000` (`k` names the series, the sample value is the given value). + async fn build_count_values_table_provider_with_values( + values: &[f64], + ) -> DfTableSourceProvider { + build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( + values.to_vec(), + ))) + .await + } + + /// Like [`build_count_values_table_provider_with_values`], but with a caller provided value + /// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). + async fn build_count_values_table_provider_with_value_array( + values: ArrayRef, + ) -> DfTableSourceProvider { + let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value().to_string(), value_data_type, true), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )), + Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), + values.clone(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + "cv_metric", + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + + /// Collects `(label, value)` pairs of a `count_values` result, where `label` is the + /// PromQL label generated by `count_values` and `value` is the aggregated sample value. + /// + /// The generated label holds the original sample value in PromQL's textual form, so it + /// is asserted as a string: comparing it as a number would not catch formatting bugs + /// (`1.0` instead of `1`, scientific notation, ...). + fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + // The aggregated value is the only numeric column that is not the generated label. + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| { + field.name() != label + && matches!( + field.data_type(), + ArrowDataType::Float64 + | ArrowDataType::Int64 + | ArrowDataType::UInt64 + ) + }) + .expect("no aggregated value column"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + let values = datafusion::arrow::compute::cast( + batch.column(value_index), + &ArrowDataType::Float64, + ) + .unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + labels + .iter() + .zip(values.iter()) + .map(|(label, value)| (label.unwrap(), value.unwrap())) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); + rows + } + + /// Asserts that a `count_values` result holds one sample per label set and evaluation + /// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat + /// a label set (that would mean the samples were still grouped by the overwritten input + /// label). + fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { + let mut seen = HashMap::>::new(); + for batch in batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + for (timestamp, label) in timestamps.iter().zip(labels.iter()) { + let timestamp = timestamp.unwrap(); + let label = label.unwrap(); + assert!( + seen.entry(timestamp).or_default().insert(label.to_string()), + "duplicated label set `{label}` at timestamp {timestamp}" + ); + } + } + } + + #[tokio::test] + async fn test_count_values_generated_label_survives_enclosing_expr() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + for (case, label) in [ + (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), + ( + r#"abs(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"round(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), + ( + r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, + "v", + ), + ( + r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, + "v", + ), + // The generated label overwrites an input label with the same name. + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, + "ip", + ), + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, + "ip", + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + &["ip"], + ) + .await, + &build_eval_stmt(case), + &build_query_engine_state(), + ) + .await + .unwrap(); + + let label_columns = plan + .schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(); + assert_eq!( + label_columns, + 1, + "{case}: the `{label}` label must survive: {}", + plan.display_indent() + ); + } + } + + #[tokio::test] + async fn test_count_values_generated_label_in_enclosing_expr_execute() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + let state = build_query_engine_state(); + // (query, generated label, expected `(label value, aggregated value)` pairs) + for (query, label, expected) in [ + ( + r#"count_values("v", cv_metric)"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"abs(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"round(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"count_values("v", cv_metric) + 1"#, + "v", + vec![("1", 3.0), ("2", 2.0)], + ), + ( + r#"sum by (v) (count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"topk(10, count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + // The generated label overwrites the input label with the same name, and the + // samples are grouped by the generated label only: `{k="1"}` holds the two + // samples of value `1.0` instead of one row per (overwritten label, value). + ( + r#"count_values("k", cv_metric) by (k)"#, + "k", + vec![("1", 2.0), ("2", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap_or_else(|err| panic!("{query}: {err}")); + + assert_eq!( + plan.schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(), + 1, + "{query}: {}", + plan.display_indent() + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(count_values_rows(&batches, label), expected, "{query}"); + assert_unique_label_set_per_timestamp(&batches, label); + } + } + + #[tokio::test] + async fn test_count_values_label_is_prometheus_formatted_value() { + // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: + // the shortest decimal form of the sample value without an exponent. The label is a + // label, so it must be a string column holding exactly that text: arrow's + // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. + let state = build_query_engine_state(); + // Samples are grouped by that formatted text, exactly like Prometheus groups by the + // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that + // round to the same text share one group. `0.0` formats as "0". + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_values(&[ + -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, + ]) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_unique_label_set_per_timestamp(&batches, "v"); + let mut labels = count_values_rows(&batches, "v") + .into_iter() + .map(|(label, _)| label) + .collect::>(); + labels.sort(); + assert_eq!( + labels, + vec![ + "-0", + "0", + "0.0000001", + "0.5", + "1", + "1000000000000000000000", + "2.5", + "200", + ] + ); + } + + #[tokio::test] + async fn test_count_values_groups_by_formatted_value_for_bigint_input() { + // The grouping key of `count_values` is the formatted sample value, not the raw input + // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and + // `2^53 + 1`) cast and format to the same label, so they must share one group and one + // count, exactly like Prometheus, which groups by the generated label text. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ + 9_007_199_254_740_992_i64, + 9_007_199_254_740_993_i64, + ]))) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + // One timestamp must never carry the same label set twice. + assert_unique_label_set_per_timestamp(&batches, "v"); + assert_eq!( + count_values_rows(&batches, "v"), + vec![("9007199254740992", 2.0)] + ); + } } diff --git a/src/servers/tests/http/mod.rs b/src/servers/tests/http/mod.rs index cca2dfd7873..ffd76b3753b 100644 --- a/src/servers/tests/http/mod.rs +++ b/src/servers/tests/http/mod.rs @@ -16,4 +16,5 @@ mod authorize; mod http_handler_test; mod influxdb_test; mod opentsdb_test; +mod prom_count_values_test; mod prom_store_test; diff --git a/src/servers/tests/http/prom_count_values_test.rs b/src/servers/tests/http/prom_count_values_test.rs new file mode 100644 index 00000000000..ab17ef8791f --- /dev/null +++ b/src/servers/tests/http/prom_count_values_test.rs @@ -0,0 +1,192 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Regression tests for the Prometheus HTTP API conversion of `count_values` results. +//! +//! `count_values` turns a sample value into a label, so the generated column must be a +//! *string* column holding Prometheus' textual form of the value: the conversion infers +//! labels from string columns and the sample value from the first numeric column. A numeric +//! label column is reported as the sample value of the series instead of a label. + +use std::sync::Arc; + +use catalog::RegisterTableRequest; +use catalog::memory::MemoryCatalogManager; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_query::prelude::greptime_value; +use common_recordbatch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector}; +use promql_parser::parser::value::ValueType; +use query::options::QueryOptions; +use query::parser::{PromQuery, QueryLanguageParser}; +use query::query_engine::QueryEngineFactory; +use servers::http::prometheus::{ + PromData, PromQueryResult, PrometheusJsonResponse, PrometheusResponse, +}; +use session::context::QueryContext; +use table::Table; +use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; +use table::test_util::MemTable; + +/// Catalog with a single metric `cv_metric` holding one series per value at `ts=1000ms`. +fn metric_catalog_manager(values: &[f64]) -> Arc { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(StringVector::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )) as _, + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000; + values.len() + ])) as _, + Arc::new(Float64Vector::from_vec(values.to_vec())) as _, + ], + ) + .unwrap(); + let backing = MemTable::new_with_catalog( + "cv_metric", + batch, + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + catalog_list +} + +#[tokio::test] +async fn count_values_generated_label_is_reported_as_a_prometheus_label() { + // Two samples carry the value 5.0, one carries 0.5 and one carries 200.0. The query + // evaluates at `ts=1s`, the only timestamp holding samples. + let catalog_list = metric_catalog_manager(&[5.0, 5.0, 0.5, 200.0]); + let query_engine = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + + let query_ctx = QueryContext::arc(); + let prom_query = PromQuery { + query: r#"count_values("v", cv_metric) + 1"#.to_string(), + start: "1".to_string(), + end: "1".to_string(), + step: "5s".to_string(), + lookback: "5m".to_string(), + alias: None, + }; + let statement = QueryLanguageParser::parse_promql(&prom_query, &query_ctx).unwrap(); + let plan = query_engine + .planner() + .plan(&statement, query_ctx.clone()) + .await + .unwrap(); + let output = query_engine.execute(plan, query_ctx).await.unwrap(); + + let response = + PrometheusJsonResponse::from_query_result(Ok(output), None, ValueType::Vector, None).await; + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response.data + else { + panic!("expected a vector response"); + }; + + // `v` is the generated label, so it must be reported as a label holding Prometheus' + // textual form of the sample value (`5`, not `5.0`), and the series carries the + // aggregation result (`count_values("v", cv_metric) + 1`) as its sample value. + let mut actual = series + .into_iter() + .map(|series| { + assert_eq!( + series.metric.keys().collect::>(), + vec!["v"], + "`v` must be the only label of the series" + ); + let label = series.metric["v"].clone(); + let (timestamp, value) = series + .value + .expect("every series of a vector result carries a sample"); + (label, timestamp, value) + }) + .collect::>(); + actual.sort_by(|left, right| left.0.cmp(&right.0)); + + assert_eq!( + actual, + vec![ + ("0.5".to_string(), 1.0, "2.0".to_string()), + ("200".to_string(), 1.0, "2.0".to_string()), + ("5".to_string(), 1.0, "3.0".to_string()), + ] + ); +} diff --git a/tests/cases/distributed/flow-tql/flow_tql.result b/tests/cases/distributed/flow-tql/flow_tql.result index 925fff3f242..30158a176f8 100644 --- a/tests/cases/distributed/flow-tql/flow_tql.result +++ b/tests/cases/distributed/flow-tql/flow_tql.result @@ -228,14 +228,14 @@ SELECT * FROM cnt_reqs ORDER BY ts, status_code; +--------------------------+---------------------+-------------+ | count(http_requests.val) | ts | status_code | +--------------------------+---------------------+-------------+ -| 3.0 | 1970-01-01T00:00:00 | 200.0 | -| 1.0 | 1970-01-01T00:00:00 | 401.0 | -| 1.0 | 1970-01-01T00:00:05 | 401.0 | -| 2.0 | 1970-01-01T00:00:05 | 404.0 | -| 1.0 | 1970-01-01T00:00:05 | 500.0 | -| 2.0 | 1970-01-01T00:00:10 | 200.0 | -| 2.0 | 1970-01-01T00:00:10 | 201.0 | -| 4.0 | 1970-01-01T00:00:15 | 500.0 | +| 3.0 | 1970-01-01T00:00:00 | 200 | +| 1.0 | 1970-01-01T00:00:00 | 401 | +| 1.0 | 1970-01-01T00:00:05 | 401 | +| 2.0 | 1970-01-01T00:00:05 | 404 | +| 1.0 | 1970-01-01T00:00:05 | 500 | +| 2.0 | 1970-01-01T00:00:10 | 200 | +| 2.0 | 1970-01-01T00:00:10 | 201 | +| 4.0 | 1970-01-01T00:00:15 | 500 | +--------------------------+---------------------+-------------+ DROP FLOW calc_reqs; diff --git a/tests/cases/standalone/common/promql/count_values.result b/tests/cases/standalone/common/promql/count_values.result index e2f65eae342..886c3a143b7 100644 --- a/tests/cases/standalone/common/promql/count_values.result +++ b/tests/cases/standalone/common/promql/count_values.result @@ -61,6 +61,145 @@ TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc); | 2 | idc2 | 1970-01-01T00:00:15 | 500 | +--------------------------+------+---------------------+-------------+ +-- https://github.com/GreptimeTeam/greptimedb/issues/9181 +-- The label generated by `count_values` must survive an enclosing expression, +-- e.g. a per-sample function, an arithmetic operation or an enclosing aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') abs(count_values("v", http_requests)); + ++---------------------+-------------------------------+-----+ +| ts | abs(count(http_requests.val)) | v | ++---------------------+-------------------------------+-----+ +| 1970-01-01T00:00:00 | 1 | 401 | +| 1970-01-01T00:00:00 | 3 | 200 | +| 1970-01-01T00:00:05 | 1 | 401 | +| 1970-01-01T00:00:05 | 1 | 500 | +| 1970-01-01T00:00:05 | 2 | 404 | +| 1970-01-01T00:00:10 | 2 | 200 | +| 1970-01-01T00:00:10 | 2 | 201 | +| 1970-01-01T00:00:15 | 4 | 500 | ++---------------------+-------------------------------+-----+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') round(count_values("v", http_requests)); + ++---------------------+-------------------------------------------------+-----+ +| ts | prom_round(count(http_requests.val),Float64(0)) | v | ++---------------------+-------------------------------------------------+-----+ +| 1970-01-01T00:00:00 | 1.0 | 401 | +| 1970-01-01T00:00:00 | 3.0 | 200 | +| 1970-01-01T00:00:05 | 1.0 | 401 | +| 1970-01-01T00:00:05 | 1.0 | 500 | +| 1970-01-01T00:00:05 | 2.0 | 404 | +| 1970-01-01T00:00:10 | 2.0 | 200 | +| 1970-01-01T00:00:10 | 2.0 | 201 | +| 1970-01-01T00:00:15 | 4.0 | 500 | ++---------------------+-------------------------------------------------+-----+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("v", http_requests) + 1; + ++-----+---------------------+---------------------------------------+ +| v | ts | count(http_requests.val) + Float64(1) | ++-----+---------------------+---------------------------------------+ +| 200 | 1970-01-01T00:00:00 | 4.0 | +| 200 | 1970-01-01T00:00:10 | 3.0 | +| 201 | 1970-01-01T00:00:10 | 3.0 | +| 401 | 1970-01-01T00:00:00 | 2.0 | +| 401 | 1970-01-01T00:00:05 | 2.0 | +| 404 | 1970-01-01T00:00:05 | 3.0 | +| 500 | 1970-01-01T00:00:05 | 2.0 | +| 500 | 1970-01-01T00:00:15 | 5.0 | ++-----+---------------------+---------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') sum by (v) (count_values("v", http_requests)); + ++-----+---------------------+-------------------------------+ +| v | ts | sum(count(http_requests.val)) | ++-----+---------------------+-------------------------------+ +| 200 | 1970-01-01T00:00:00 | 3 | +| 200 | 1970-01-01T00:00:10 | 2 | +| 201 | 1970-01-01T00:00:10 | 2 | +| 401 | 1970-01-01T00:00:00 | 1 | +| 401 | 1970-01-01T00:00:05 | 1 | +| 404 | 1970-01-01T00:00:05 | 2 | +| 500 | 1970-01-01T00:00:05 | 1 | +| 500 | 1970-01-01T00:00:15 | 4 | ++-----+---------------------+-------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') topk(10, count_values("v", http_requests)); + ++--------------------------+-----+---------------------+ +| count(http_requests.val) | v | ts | ++--------------------------+-----+---------------------+ +| 1 | 401 | 1970-01-01T00:00:00 | +| 1 | 401 | 1970-01-01T00:00:05 | +| 1 | 500 | 1970-01-01T00:00:05 | +| 2 | 200 | 1970-01-01T00:00:10 | +| 2 | 201 | 1970-01-01T00:00:10 | +| 2 | 404 | 1970-01-01T00:00:05 | +| 3 | 200 | 1970-01-01T00:00:00 | +| 4 | 500 | 1970-01-01T00:00:15 | ++--------------------------+-----+---------------------+ + +-- The generated label overwrites an input label with the same name (PromQL semantics), +-- and must not turn the generated column into an ambiguous reference. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("idc", http_requests) by (idc); + ++--------------------------+---------------------+-----+ +| count(http_requests.val) | ts | idc | ++--------------------------+---------------------+-----+ +| 1 | 1970-01-01T00:00:00 | 401 | +| 1 | 1970-01-01T00:00:05 | 401 | +| 1 | 1970-01-01T00:00:05 | 500 | +| 2 | 1970-01-01T00:00:05 | 404 | +| 2 | 1970-01-01T00:00:10 | 200 | +| 2 | 1970-01-01T00:00:10 | 201 | +| 3 | 1970-01-01T00:00:00 | 200 | +| 4 | 1970-01-01T00:00:15 | 500 | ++--------------------------+---------------------+-----+ + +-- The generated label is the sample value in PromQL's textual form +-- (`strconv.FormatFloat(value, 'f', -1, 64)`: the shortest decimal form without an +-- exponent), so it is a string: `200.0` becomes the label `200` and `1e21` becomes +-- `1000000000000000000000`. +CREATE TABLE http_requests_double ( + ts timestamp(3) time index, + host STRING, + idc STRING, + val DOUBLE, + PRIMARY KEY(host, idc), +); + +Affected Rows: 0 + +INSERT INTO TABLE http_requests_double VALUES + (0, 'host1', "idc1", 200.0), + (0, 'host2', "idc1", 0.5), + (5000, 'host1', "idc1", 200.0), + (5000, 'host2', "idc1", 1e21); + +Affected Rows: 4 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') count_values("val_label", http_requests_double); + ++---------------------------------+---------------------+------------------------+ +| count(http_requests_double.val) | ts | val_label | ++---------------------------------+---------------------+------------------------+ +| 1 | 1970-01-01T00:00:00 | 0.5 | +| 1 | 1970-01-01T00:00:00 | 200 | +| 1 | 1970-01-01T00:00:05 | 1000000000000000000000 | +| 1 | 1970-01-01T00:00:05 | 200 | ++---------------------------------+---------------------+------------------------+ + +DROP TABLE http_requests_double; + +Affected Rows: 0 + DROP TABLE http_requests; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/count_values.sql b/tests/cases/standalone/common/promql/count_values.sql index 09c2b136b83..f23cdec5fbe 100644 --- a/tests/cases/standalone/common/promql/count_values.sql +++ b/tests/cases/standalone/common/promql/count_values.sql @@ -28,4 +28,50 @@ TQL EVAL (0, 15, '5s') count_values("status_code", http_requests); TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc); +-- https://github.com/GreptimeTeam/greptimedb/issues/9181 +-- The label generated by `count_values` must survive an enclosing expression, +-- e.g. a per-sample function, an arithmetic operation or an enclosing aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') abs(count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') round(count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("v", http_requests) + 1; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') sum by (v) (count_values("v", http_requests)); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') topk(10, count_values("v", http_requests)); + +-- The generated label overwrites an input label with the same name (PromQL semantics), +-- and must not turn the generated column into an ambiguous reference. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') count_values("idc", http_requests) by (idc); + +-- The generated label is the sample value in PromQL's textual form +-- (`strconv.FormatFloat(value, 'f', -1, 64)`: the shortest decimal form without an +-- exponent), so it is a string: `200.0` becomes the label `200` and `1e21` becomes +-- `1000000000000000000000`. +CREATE TABLE http_requests_double ( + ts timestamp(3) time index, + host STRING, + idc STRING, + val DOUBLE, + PRIMARY KEY(host, idc), +); + +INSERT INTO TABLE http_requests_double VALUES + (0, 'host1', "idc1", 200.0), + (0, 'host2', "idc1", 0.5), + (5000, 'host1', "idc1", 200.0), + (5000, 'host2', "idc1", 1e21); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 5, '5s') count_values("val_label", http_requests_double); + +DROP TABLE http_requests_double; + DROP TABLE http_requests; diff --git a/tests/cases/standalone/flow-tql/flow_tql.result b/tests/cases/standalone/flow-tql/flow_tql.result index de93fb7c518..13239bf4e1f 100644 --- a/tests/cases/standalone/flow-tql/flow_tql.result +++ b/tests/cases/standalone/flow-tql/flow_tql.result @@ -228,14 +228,14 @@ SELECT * FROM cnt_reqs ORDER BY ts, status_code; +--------------------------+---------------------+-------------+ | count(http_requests.val) | ts | status_code | +--------------------------+---------------------+-------------+ -| 3.0 | 1970-01-01T00:00:00 | 200.0 | -| 1.0 | 1970-01-01T00:00:00 | 401.0 | -| 1.0 | 1970-01-01T00:00:05 | 401.0 | -| 2.0 | 1970-01-01T00:00:05 | 404.0 | -| 1.0 | 1970-01-01T00:00:05 | 500.0 | -| 2.0 | 1970-01-01T00:00:10 | 200.0 | -| 2.0 | 1970-01-01T00:00:10 | 201.0 | -| 4.0 | 1970-01-01T00:00:15 | 500.0 | +| 3.0 | 1970-01-01T00:00:00 | 200 | +| 1.0 | 1970-01-01T00:00:00 | 401 | +| 1.0 | 1970-01-01T00:00:05 | 401 | +| 2.0 | 1970-01-01T00:00:05 | 404 | +| 1.0 | 1970-01-01T00:00:05 | 500 | +| 2.0 | 1970-01-01T00:00:10 | 200 | +| 2.0 | 1970-01-01T00:00:10 | 201 | +| 4.0 | 1970-01-01T00:00:15 | 500 | +--------------------------+---------------------+-------------+ DROP FLOW calc_reqs; From 20619273eede955f62f252cc2e01e3038c27b0bd Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 24 Sep 2026 08:37:43 +0000 Subject: [PATCH 034/133] chore(toolchain): switch to stable Rust 1.96.1 and remove all nightly feature gates (#9303) * chore(toolchain): switch to stable Rust 1.96.1 and remove all nightly feature gates Move the workspace from the pinned nightly-2026-03-21 to stable 1.96.1 and drop all 23 '#![feature]' gates across 13 crates, rewriting the still-unstable API usages with stable equivalents: - try_blocks: closures / an async block (table, query, common-function, servers) - duration_constructors: Duration::from_secs(n * 86400) / (n * 60) - iterator_try_collect: collect::, _>>() - box_patterns: as_deref() + matches! chains (sql) - error_iter: error_chain_root() source-chain walker (common-error); sources() includes the error itself, so the walker never panics - int_roundings: div_floor -> div_euclid (equal for positive divisors) - iter_partition_in_place: stable sort_by_key partition helper (index) - hash_set_entry: HashSet::insert bool / contains+insert - trait_alias: lifetime-parameterized dyn FnOnce type aliases (puffin) - string_from_utf8_lossy_owned: from_utf8_lossy(&v).into_owned() - never_type: Infallible (common-recordbatch) - debug_closure_helpers: closure-backed DebugFmt newtype (mito2) - binary_heap_pop_if: peek().is_some_and() + pop() - exclusive_wrapper: drop Exclusive; C: Send + Unpin already in bounds - stmt_expr_attributes: stale gate, no usages Also fix release-dev-builder-images.yaml, which parsed rust-toolchain.toml with a date-only regex and would produce empty image versions with a stable channel; it now extracts the full channel token. Dev-builder images verified against stable 1.96.1 (image build, default-toolchain behavior, binstall/nextest, riscv64 and android targets, in-image cargo check). Validated on 1.96.1: cargo check --workspace --all-targets, clippy --workspace --all-targets --all-features -D warnings, cargo fmt --check, and nextest on all 13 affected crates (4586 passed). Part of #9289. Depends on #9298 (fuzz nightly quarantine) merging first. Signed-off-by: Ning Sun * chore: update flake checksum * chore: use wild for linker in flake --------- Signed-off-by: Ning Sun --- .../workflows/release-dev-builder-images.yaml | 5 +- flake.nix | 4 +- rust-toolchain.toml | 2 +- src/common/error/src/ext.rs | 18 +++- src/common/error/src/lib.rs | 2 - src/common/event-recorder/src/lib.rs | 2 - src/common/event-recorder/src/recorder.rs | 4 +- src/common/function/src/lib.rs | 2 - src/common/function/src/scalars/matches.rs | 5 +- src/common/recordbatch/src/adapter.rs | 2 +- src/common/recordbatch/src/lib.rs | 2 - src/common/telemetry/src/lib.rs | 2 - src/common/telemetry/src/logging.rs | 2 +- src/common/wal/src/config.rs | 2 +- src/common/wal/src/config/kafka/common.rs | 2 +- src/index/src/bloom_filter/creator.rs | 8 +- src/index/src/inverted_index/search.rs | 10 ++ .../search/fst_apply/keys_apply.rs | 12 +-- .../search/index_apply/predicates_apply.rs | 7 +- src/index/src/lib.rs | 3 - src/meta-srv/src/gc/options.rs | 10 +- .../src/handler/persist_stats_handler.rs | 4 +- src/meta-srv/src/lib.rs | 3 - src/meta-srv/src/metasrv.rs | 2 +- src/meta-srv/src/metasrv/builder.rs | 2 +- src/meta-srv/src/procedure/repartition.rs | 2 +- .../src/procedure/wal_prune/manager.rs | 17 ++-- src/mito2/src/compaction/run.rs | 6 +- src/mito2/src/lib.rs | 4 - src/mito2/src/memtable/time_partition.rs | 2 +- src/mito2/src/sst/file.rs | 50 ++++++---- .../src/etl/transform/transformer/greptime.rs | 6 +- src/pipeline/src/lib.rs | 2 - src/pipeline/src/tablesuffix.rs | 2 +- src/puffin/src/lib.rs | 2 - src/puffin/src/puffin_manager/stager.rs | 8 +- .../puffin_manager/stager/bounded_stager.rs | 14 +-- src/query/src/lib.rs | 6 -- src/query/src/log_query/planner.rs | 16 ++-- src/query/src/optimizer/windowed_sort.rs | 8 +- src/query/src/planner.rs | 8 +- src/query/src/range_select/plan.rs | 4 +- src/query/src/window_sort.rs | 7 +- .../src/http/result/prometheus_resp.rs | 7 +- src/servers/src/lib.rs | 3 - src/servers/src/postgres/auth_handler.rs | 7 +- src/sql/src/lib.rs | 2 - src/sql/src/statements/insert.rs | 95 ++++++++++--------- src/table/src/lib.rs | 2 - src/table/src/table/scan.rs | 8 +- 50 files changed, 209 insertions(+), 196 deletions(-) diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index bebd5ed61e7..41325f61ab9 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -59,7 +59,10 @@ jobs: commitShortSHA=`echo ${{ github.sha }} | cut -c1-8` buildTime=`date +%Y%m%d%H%M%S` BUILD_VERSION="$commitShortSHA-$buildTime" - RUST_TOOLCHAIN_VERSION=$(cat rust-toolchain.toml | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + # Extract the full channel string (e.g. "1.96.1" or + # "nightly-2026-03-21") so the image version always carries the + # toolchain it was built with. + RUST_TOOLCHAIN_VERSION=$(grep -E '^channel' rust-toolchain.toml | cut -d'"' -f2) IMAGE_VERSION="${RUST_TOOLCHAIN_VERSION}-${BUILD_VERSION}" echo "VERSION=${IMAGE_VERSION}" >> $GITHUB_ENV echo "version=$IMAGE_VERSION" >> $GITHUB_OUTPUT diff --git a/flake.nix b/flake.nix index 2d95d8dcfd6..6adb69678b2 100644 --- a/flake.nix +++ b/flake.nix @@ -22,7 +22,7 @@ lib = nixpkgs.lib; rustToolchain = fenix.packages.${system}.fromToolchainName { name = (lib.importTOML ./rust-toolchain.toml).toolchain.channel; - sha256 = "sha256-rboGKQLH4eDuiY01SINOqmXUFUNr9F4awoFZGzib17o="; + sha256 = "sha256-h+t2xTBz5yt2YIO+1VMIIGlCU7gyp2LYOFvaV1nwOXU="; }; in { @@ -34,7 +34,7 @@ gcc protobuf gnumake - mold + wild (rustToolchain.withComponents [ "cargo" "clippy" diff --git a/rust-toolchain.toml b/rust-toolchain.toml index d16edecca80..2d45363a5be 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "nightly-2026-03-21" +channel = "1.96.1" diff --git a/src/common/error/src/ext.rs b/src/common/error/src/ext.rs index 59895551156..0c267dc15b2 100644 --- a/src/common/error/src/ext.rs +++ b/src/common/error/src/ext.rs @@ -21,6 +21,20 @@ use std::sync::Arc; use serde::{Deserialize, Deserializer, Serializer}; use snafu::{FromString, Snafu}; +/// Returns the root cause of an error's source chain, i.e. the last error +/// reachable via [`std::error::Error::source`]. For an error without any +/// source the error itself is returned. +/// +/// Mirrors `err.sources().last().unwrap()` (unstable `error_iter`, which +/// yields the error itself followed by its sources). +fn error_chain_root(err: &dyn std::error::Error) -> &dyn std::error::Error { + let mut root = err; + while let Some(source) = root.source() { + root = source; + } + root +} + use crate::status_code::StatusCode; /// Describes whether an error instance is safe and useful to retry. @@ -148,7 +162,7 @@ pub trait ErrorExt: StackError { _ => { let error = self.last(); if let Some(external_error) = error.source() { - let external_root = external_error.sources().last().unwrap(); + let external_root = error_chain_root(external_error); if error.transparent() { format!("{external_root}") @@ -169,7 +183,7 @@ pub trait ErrorExt: StackError { { let error = self.last(); if let Some(external_error) = error.source() { - let external_root = external_error.sources().last().unwrap(); + let external_root = error_chain_root(external_error); Some(external_root) } else { None diff --git a/src/common/error/src/lib.rs b/src/common/error/src/lib.rs index eeedc637133..29ae5c3b1cc 100644 --- a/src/common/error/src/lib.rs +++ b/src/common/error/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(error_iter)] - pub mod ext; pub mod mock; pub mod status_code; diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 7ee71722229..0a5d051a01e 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(duration_constructors)] - pub mod context; pub mod error; pub mod event_table; diff --git a/src/common/event-recorder/src/recorder.rs b/src/common/event-recorder/src/recorder.rs index 3ea5c7e55a8..040b2fbcbb3 100644 --- a/src/common/event-recorder/src/recorder.rs +++ b/src/common/event-recorder/src/recorder.rs @@ -108,9 +108,9 @@ fn event_type_filter_is_all(event_types: &EventTypeFilterRef) -> bool { /// The time interval for flushing batched events to the event handler. pub const DEFAULT_FLUSH_INTERVAL_SECONDS: Duration = Duration::from_secs(5); /// The default TTL(90 days) for the events table. -const DEFAULT_EVENTS_TABLE_TTL: Duration = Duration::from_days(90); +const DEFAULT_EVENTS_TABLE_TTL: Duration = Duration::from_secs(90 * 86400); /// The default compaction time window for the events table. -pub const DEFAULT_COMPACTION_TIME_WINDOW: Duration = Duration::from_days(1); +pub const DEFAULT_COMPACTION_TIME_WINDOW: Duration = Duration::from_secs(86400); // The capacity of the tokio channel for transmitting events to background processor. const DEFAULT_CHANNEL_SIZE: usize = 2048; // The size of the buffer for batching events before flushing to event handler. diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index 336f8d489ca..7309dba981d 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] - mod admin; mod flush_flow; mod macros; diff --git a/src/common/function/src/scalars/matches.rs b/src/common/function/src/scalars/matches.rs index 14889957ee9..109ad1daeb2 100644 --- a/src/common/function/src/scalars/matches.rs +++ b/src/common/function/src/scalars/matches.rs @@ -1119,11 +1119,12 @@ mod test { ]; for (query, expected) in cases { - let result: Result<()> = try { + let result: Result<()> = (|| { let parser = ParserContext { stack: vec![] }; let ast = parser.parse_pattern(query)?; let _ast = ast.transform_ast()?; - }; + Ok(()) + })(); assert!(result.is_err(), "{query}"); let actual_error = result.unwrap_err().to_string(); diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index a7f6e4a10ac..f8d645bd3e1 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -570,7 +570,7 @@ impl MetricCollector { } impl ExecutionPlanVisitor for MetricCollector { - type Error = !; + type Error = std::convert::Infallible; fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> std::result::Result { // skip if no metric available diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 4fb37cd5b11..0ccb8a4c28e 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(never_type)] - pub mod adapter; pub mod cursor; pub mod error; diff --git a/src/common/telemetry/src/lib.rs b/src/common/telemetry/src/lib.rs index 26bf5d53b3c..6daf2662abe 100644 --- a/src/common/telemetry/src/lib.rs +++ b/src/common/telemetry/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(duration_constructors)] - pub mod logging; mod macros; pub mod metric; diff --git a/src/common/telemetry/src/logging.rs b/src/common/telemetry/src/logging.rs index a9b50930786..dd99e901f75 100644 --- a/src/common/telemetry/src/logging.rs +++ b/src/common/telemetry/src/logging.rs @@ -348,7 +348,7 @@ impl Default for SlowQueryOptions { record_type: SlowQueriesRecordType::SystemTable, threshold: Duration::from_secs(30), sample_ratio: 1.0, - ttl: Duration::from_days(90), + ttl: Duration::from_secs(90 * 86400), } } } diff --git a/src/common/wal/src/config.rs b/src/common/wal/src/config.rs index 343df193ab4..f79a28607cc 100644 --- a/src/common/wal/src/config.rs +++ b/src/common/wal/src/config.rs @@ -234,7 +234,7 @@ mod tests { create_topic_timeout: Duration::from_secs(30), }, auto_create_topics: true, - auto_prune_interval: Duration::from_mins(30), + auto_prune_interval: Duration::from_secs(30 * 60), auto_prune_logical_delete: false, auto_prune_parallelism: 10, flush_trigger_size: ReadableSize::mb(512), diff --git a/src/common/wal/src/config/kafka/common.rs b/src/common/wal/src/config/kafka/common.rs index cf632fbdce3..daa5607e476 100644 --- a/src/common/wal/src/config/kafka/common.rs +++ b/src/common/wal/src/config/kafka/common.rs @@ -38,7 +38,7 @@ pub const DEFAULT_BACKOFF_CONFIG: BackoffConfig = BackoffConfig { }; /// Default interval for auto WAL pruning. -pub const DEFAULT_AUTO_PRUNE_INTERVAL: Duration = Duration::from_mins(30); +pub const DEFAULT_AUTO_PRUNE_INTERVAL: Duration = Duration::from_secs(30 * 60); /// Default mode for auto WAL pruning. pub const DEFAULT_AUTO_PRUNE_LOGICAL_DELETE: bool = false; /// Default limit for concurrent auto pruning tasks. diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index df23ca2352a..07f4f977caf 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/src/index/src/bloom_filter/creator.rs @@ -161,10 +161,10 @@ impl BloomFilterCreator { if let Some(elem) = elem { let old_len = self.cur_seg_distinct_elems.len(); - // A borrowed entry lookup avoids hashing new values twice and only - // allocates when the value is absent from the current segment. - self.cur_seg_distinct_elems - .get_or_insert_with(elem, <[u8]>::to_vec); + // Only allocate when the value is absent from the current segment. + if !self.cur_seg_distinct_elems.contains(elem) { + self.cur_seg_distinct_elems.insert(elem.to_vec()); + } if self.cur_seg_distinct_elems.len() != old_len { self.cur_seg_distinct_elems_mem_usage += elem.len(); self.global_memory_usage diff --git a/src/index/src/inverted_index/search.rs b/src/index/src/inverted_index/search.rs index 8e28440c7e2..581135c2d4e 100644 --- a/src/index/src/inverted_index/search.rs +++ b/src/index/src/inverted_index/search.rs @@ -16,3 +16,13 @@ pub mod fst_apply; pub mod fst_values_mapper; pub mod index_apply; pub mod predicate; + +/// Partitions `slice` in place so that elements matching `pred` come first, +/// preserving the original relative order within each group. Returns the +/// number of matching elements. +/// +/// Stable replacement for the unstable `Iterator::partition_in_place`. +pub(crate) fn partition_in_place(slice: &mut [T], mut pred: impl FnMut(&T) -> bool) -> usize { + slice.sort_by_key(|x| !pred(x)); + slice.iter().filter(|x| pred(x)).count() +} diff --git a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs index d565749452d..b68f58cdfd3 100644 --- a/src/index/src/inverted_index/search/fst_apply/keys_apply.rs +++ b/src/index/src/inverted_index/search/fst_apply/keys_apply.rs @@ -68,16 +68,16 @@ impl KeysFstApplier { } fn split_at_in_lists(predicates: &mut [Predicate]) -> (&mut [Predicate], &mut [Predicate]) { - let in_list_index = predicates - .iter_mut() - .partition_in_place(|p| matches!(p, Predicate::InList(_))); + let in_list_index = crate::inverted_index::search::partition_in_place(predicates, |p| { + matches!(p, Predicate::InList(_)) + }); predicates.split_at_mut(in_list_index) } fn split_at_ranges(predicates: &mut [Predicate]) -> (&mut [Predicate], &mut [Predicate]) { - let range_index = predicates - .iter_mut() - .partition_in_place(|p| matches!(p, Predicate::Range(_))); + let range_index = crate::inverted_index::search::partition_in_place(predicates, |p| { + matches!(p, Predicate::Range(_)) + }); predicates.split_at_mut(range_index) } diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index eaea8cfb95c..5fcc962ca7d 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs @@ -127,9 +127,10 @@ impl PredicatesIndexApplier { let mut fst_appliers = Vec::with_capacity(predicates.len()); // InList predicates are applied first to benefit from higher selectivity. - let in_list_index = predicates - .iter_mut() - .partition_in_place(|(_, ps)| ps.iter().any(|p| matches!(p, Predicate::InList(_)))); + let in_list_index = + crate::inverted_index::search::partition_in_place(&mut predicates, |(_, ps)| { + ps.iter().any(|p| matches!(p, Predicate::InList(_))) + }); let mut iter = predicates.into_iter(); for _ in 0..in_list_index { let (column_name, predicates) = iter.next().unwrap(); diff --git a/src/index/src/lib.rs b/src/index/src/lib.rs index 2c1ceb4e3f6..571866ade90 100644 --- a/src/index/src/lib.rs +++ b/src/index/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(iter_partition_in_place)] -#![feature(hash_set_entry)] - pub mod bitmap; pub mod bloom_filter; pub mod error; diff --git a/src/meta-srv/src/gc/options.rs b/src/meta-srv/src/gc/options.rs index fb43bd6dce1..298fb30bfd8 100644 --- a/src/meta-srv/src/gc/options.rs +++ b/src/meta-srv/src/gc/options.rs @@ -38,7 +38,7 @@ impl Default for SoftDropGcOptions { fn default() -> Self { Self { enable: false, - retention: Duration::from_days(7), + retention: Duration::from_secs(7 * 86400), } } } @@ -242,7 +242,7 @@ mod tests { assert!(!options.experimental_soft_drop.enable); assert_eq!( - Duration::from_days(7), + Duration::from_secs(7 * 86400), options.experimental_soft_drop.retention ); } @@ -254,7 +254,7 @@ mod tests { enable: true, experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; @@ -269,7 +269,7 @@ mod tests { enable: true, experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; @@ -284,7 +284,7 @@ mod tests { let options = GcSchedulerOptions { experimental_soft_drop: SoftDropGcOptions { enable: true, - retention: Duration::from_days(1), + retention: Duration::from_secs(86400), }, ..Default::default() }; diff --git a/src/meta-srv/src/handler/persist_stats_handler.rs b/src/meta-srv/src/handler/persist_stats_handler.rs index 04d63e26e2f..472072aa42b 100644 --- a/src/meta-srv/src/handler/persist_stats_handler.rs +++ b/src/meta-srv/src/handler/persist_stats_handler.rs @@ -154,9 +154,9 @@ fn align_ts(ts: i64, interval: Duration) -> i64 { impl PersistStatsHandler { /// Creates a new [`PersistStatsHandler`]. pub fn new(inserter: Box, mut persist_interval: Duration) -> Self { - if persist_interval < Duration::from_mins(10) { + if persist_interval < Duration::from_secs(10 * 60) { warn!("persist_interval is less than 10 minutes, set to 10 minutes"); - persist_interval = Duration::from_mins(10); + persist_interval = Duration::from_secs(10 * 60); } Self { diff --git a/src/meta-srv/src/lib.rs b/src/meta-srv/src/lib.rs index 88e8466be39..948cab240d4 100644 --- a/src/meta-srv/src/lib.rs +++ b/src/meta-srv/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(hash_set_entry)] -#![feature(duration_constructors)] - pub mod bootstrap; pub mod cache_invalidator; pub mod cluster; diff --git a/src/meta-srv/src/metasrv.rs b/src/meta-srv/src/metasrv.rs index 65a079bec17..86aee9f8b82 100644 --- a/src/meta-srv/src/metasrv.rs +++ b/src/meta-srv/src/metasrv.rs @@ -122,7 +122,7 @@ impl Default for StatsPersistenceOptions { fn default() -> Self { Self { ttl: Duration::ZERO, - interval: Duration::from_mins(10), + interval: Duration::from_secs(10 * 60), } } } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index cbc859716a4..57e7deb3db6 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -92,7 +92,7 @@ use crate::utils::database::DatabaseOperator; use crate::utils::insert_forwarder::InsertForwarder; /// The time window for twcs compaction of the region stats table. -const REGION_STATS_TABLE_TWCS_COMPACTION_TIME_WINDOW: Duration = Duration::from_days(1); +const REGION_STATS_TABLE_TWCS_COMPACTION_TIME_WINDOW: Duration = Duration::from_secs(86400); // TODO(fys): try use derive_builder macro pub struct MetasrvBuilder { diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index f42d8aacd0c..008c0ce05bb 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -107,7 +107,7 @@ pub struct PersistentContext { } fn default_timeout() -> Duration { - Duration::from_mins(2) + Duration::from_secs(2 * 60) } impl PersistentContext { diff --git a/src/meta-srv/src/procedure/wal_prune/manager.rs b/src/meta-srv/src/procedure/wal_prune/manager.rs index e398be93e45..4de21086906 100644 --- a/src/meta-srv/src/procedure/wal_prune/manager.rs +++ b/src/meta-srv/src/procedure/wal_prune/manager.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::collections::HashSet; -use std::collections::hash_set::Entry; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, RwLock}; @@ -46,15 +45,13 @@ impl WalPruneProcedureTracker { /// consume acquire a semaphore permit for the given topic name. pub fn insert_running_procedure(&self, topic_name: String) -> Option { let mut running_procedures = self.running_procedures.write().unwrap(); - match running_procedures.entry(topic_name.clone()) { - Entry::Occupied(_) => None, - Entry::Vacant(entry) => { - entry.insert(); - Some(WalPruneProcedureGuard { - topic_name, - running_procedures: self.running_procedures.clone(), - }) - } + if running_procedures.insert(topic_name.clone()) { + Some(WalPruneProcedureGuard { + topic_name, + running_procedures: self.running_procedures.clone(), + }) + } else { + None } } diff --git a/src/mito2/src/compaction/run.rs b/src/mito2/src/compaction/run.rs index 2ad9f2f1eba..664005b9af3 100644 --- a/src/mito2/src/compaction/run.rs +++ b/src/mito2/src/compaction/run.rs @@ -475,7 +475,11 @@ where for item in items { let (start, _) = item.range(); - while let Some(run) = runs_sort_by_end.pop_if(|x| x.run.end.unwrap() <= start) { + while runs_sort_by_end + .peek() + .is_some_and(|x| x.run.end.unwrap() <= start) + { + let run = runs_sort_by_end.pop().unwrap(); runs_sort_by_index.push(Wrapper(run)); } diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index 4f4409afd8a..dcbfe60dea5 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -16,10 +16,6 @@ //! //! Mito is the a region engine to store timeseries data. -#![feature(debug_closure_helpers)] -#![feature(duration_constructors)] -#![feature(binary_heap_pop_if)] - #[cfg(any(test, feature = "test"))] #[cfg_attr(feature = "test", allow(unused))] pub mod test_util; diff --git a/src/mito2/src/memtable/time_partition.rs b/src/mito2/src/memtable/time_partition.rs index a6b27cdd0f3..0ca88ca9a70 100644 --- a/src/mito2/src/memtable/time_partition.rs +++ b/src/mito2/src/memtable/time_partition.rs @@ -43,7 +43,7 @@ use crate::memtable::{KeyValues, MemtableBuilderRef, MemtableId, MemtableRef}; use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; /// Initial time window if not specified. -const INITIAL_TIME_WINDOW: Duration = Duration::from_days(1); +const INITIAL_TIME_WINDOW: Duration = Duration::from_secs(86400); /// A partition holds rows with timestamps between `[min, max)`. #[derive(Debug, Clone)] diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index c18d6195f19..87684586fbf 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -298,20 +298,33 @@ fn is_false(value: &bool) -> bool { !*value } +/// Formats a debug field with a custom closure, as a stable replacement for +/// the unstable `DebugStruct::field_with`. +struct DebugFmt) -> fmt::Result>(F); + +impl) -> fmt::Result> std::fmt::Debug for DebugFmt { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + (self.0)(f) + } +} + impl Debug for FileMeta { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let mut debug_struct = f.debug_struct("FileMeta"); debug_struct .field("region_id", &self.region_id) - .field_with("file_id", |f| write!(f, "{} ", self.file_id)) - .field_with("time_range", |f| { - write!( - f, - "({}, {}) ", - self.time_range.0.to_iso8601_string(), - self.time_range.1.to_iso8601_string() - ) - }) + .field("file_id", &DebugFmt(|f| write!(f, "{} ", self.file_id))) + .field( + "time_range", + &DebugFmt(|f| { + write!( + f, + "({}, {}) ", + self.time_range.0.to_iso8601_string(), + self.time_range.1.to_iso8601_string() + ) + }), + ) .field("level", &self.level) .field("file_size", &ReadableSize(self.file_size)) .field( @@ -327,14 +340,17 @@ impl Debug for FileMeta { debug_struct .field("num_rows", &self.num_rows) .field("num_row_groups", &self.num_row_groups) - .field_with("sequence", |f| match self.sequence { - None => { - write!(f, "None") - } - Some(seq) => { - write!(f, "{}", seq) - } - }) + .field( + "sequence", + &DebugFmt(|f| match self.sequence { + None => { + write!(f, "None") + } + Some(seq) => { + write!(f, "{}", seq) + } + }), + ) .field("partition_expr", &self.partition_expr) .field("num_series", &self.num_series); if self.primary_key_min.is_some() || self.primary_key_max.is_some() { diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index bbe3949a540..eb25071c91c 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -741,9 +741,9 @@ fn resolve_value( &ConcreteDataType::string_datatype(), schema_info, )?; - Some(ValueData::StringValue(String::from_utf8_lossy_owned( - v.to_vec(), - ))) + Some(ValueData::StringValue( + String::from_utf8_lossy(&v).into_owned(), + )) } VrlValue::Regex(v) => { diff --git a/src/pipeline/src/lib.rs b/src/pipeline/src/lib.rs index 5d1acb7462e..15954b18f98 100644 --- a/src/pipeline/src/lib.rs +++ b/src/pipeline/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(string_from_utf8_lossy_owned)] - mod dispatcher; pub mod error; mod etl; diff --git a/src/pipeline/src/tablesuffix.rs b/src/pipeline/src/tablesuffix.rs index 01ce4202cfc..4a72f49bbc6 100644 --- a/src/pipeline/src/tablesuffix.rs +++ b/src/pipeline/src/tablesuffix.rs @@ -56,7 +56,7 @@ impl TableSuffixTemplate { let v = val.get(key.as_str())?; match v { VrlValue::Integer(v) => Some(v.to_string()), - VrlValue::Bytes(v) => Some(String::from_utf8_lossy_owned(v.to_vec())), + VrlValue::Bytes(v) => Some(String::from_utf8_lossy(v).into_owned()), _ => None, } }) diff --git a/src/puffin/src/lib.rs b/src/puffin/src/lib.rs index 96a8421f98b..2be956e43dc 100644 --- a/src/puffin/src/lib.rs +++ b/src/puffin/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(trait_alias)] - pub mod blob_metadata; pub mod error; pub mod file_format; diff --git a/src/puffin/src/puffin_manager/stager.rs b/src/puffin/src/puffin_manager/stager.rs index 512e94f4e85..8a0adb34bca 100644 --- a/src/puffin/src/puffin_manager/stager.rs +++ b/src/puffin/src/puffin_manager/stager.rs @@ -43,13 +43,13 @@ pub type DirWriterProviderRef = Box; /// /// `Stager` will provide a `BoxWriter` that the caller of `get_blob` /// can use to write the blob into the staging area. -pub trait InitBlobFn = FnOnce(BoxWriter) -> WriteResult; +pub type InitBlobFn<'a> = dyn FnOnce(BoxWriter) -> WriteResult + Send + Sync + 'a; /// Function that initializes a directory. /// /// `Stager` will provide a `DirWriterProvider` that the caller of `get_dir` /// can use to write files inside the directory into the staging area. -pub trait InitDirFn = FnOnce(DirWriterProviderRef) -> WriteResult; +pub type InitDirFn<'a> = dyn FnOnce(DirWriterProviderRef) -> WriteResult + Send + Sync + 'a; /// `Stager` manages the staging area for the puffin files. #[async_trait] @@ -67,7 +67,7 @@ pub trait Stager: Send + Sync { &self, handle: &Self::FileHandle, key: &str, - init_factory: Box, + init_factory: Box>, ) -> Result; /// Retrieves a directory, initializing it if necessary using the provided `init_fn`. @@ -79,7 +79,7 @@ pub trait Stager: Send + Sync { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result<(Self::Dir, DirMetrics)>; /// Stores a directory in the staging area. diff --git a/src/puffin/src/puffin_manager/stager/bounded_stager.rs b/src/puffin/src/puffin_manager/stager/bounded_stager.rs index dfb9285452b..e0952b2892e 100644 --- a/src/puffin/src/puffin_manager/stager/bounded_stager.rs +++ b/src/puffin/src/puffin_manager/stager/bounded_stager.rs @@ -146,7 +146,7 @@ impl Stager for BoundedStager { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result { let handle_str = handle.to_string(); let cache_key = Self::encode_cache_key(&handle_str, key); @@ -202,7 +202,7 @@ impl Stager for BoundedStager { &self, handle: &Self::FileHandle, key: &str, - init_fn: Box, + init_fn: Box>, ) -> Result<(Self::Dir, DirMetrics)> { let handle_str = handle.to_string(); @@ -332,10 +332,7 @@ impl BoundedStager { BASE64_URL_SAFE.encode(hash) } - async fn write_blob( - target_path: &PathBuf, - init_fn: Box, - ) -> Result { + async fn write_blob(target_path: &PathBuf, init_fn: Box>) -> Result { // To guarantee the atomicity of writing the file, we need to write // the file to a temporary file first... let tmp_path = target_path.with_extension(TMP_EXTENSION); @@ -354,10 +351,7 @@ impl BoundedStager { Ok(size) } - async fn write_dir( - target_path: &PathBuf, - init_fn: Box, - ) -> Result { + async fn write_dir(target_path: &PathBuf, init_fn: Box>) -> Result { // To guarantee the atomicity of writing the directory, we need to write // the directory to a temporary directory first... let tmp_base = target_path.with_extension(TMP_EXTENSION); diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 4b2a2d5264d..aa83c8b8f0e 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -12,12 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(int_roundings)] -#![feature(try_blocks)] -#![feature(stmt_expr_attributes)] -#![feature(iterator_try_collect)] -#![feature(box_patterns)] - mod analyze; pub mod datafusion; pub mod dist_plan; diff --git a/src/query/src/log_query/planner.rs b/src/query/src/log_query/planner.rs index 5bc8d034a12..0855974241e 100644 --- a/src/query/src/log_query/planner.rs +++ b/src/query/src/log_query/planner.rs @@ -142,7 +142,7 @@ impl LogQueryPlanner { let exprs = filters .iter() .filter_map(|filter| self.build_filters(filter, schema).transpose()) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { Ok(None) } else { @@ -153,7 +153,7 @@ impl LogQueryPlanner { let exprs = filters .iter() .filter_map(|filter| self.build_filters(filter, schema).transpose()) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { Ok(None) } else { @@ -191,7 +191,7 @@ impl LogQueryPlanner { self.build_content_filter_with_expr(col_expr.clone(), filter, &df_schema) .transpose() }) - .try_collect::>()?; + .collect::, _>>()?; if filter_exprs.is_empty() { return Ok(Some(col_expr.is_true())); @@ -285,7 +285,7 @@ impl LogQueryPlanner { self.build_content_filter_with_expr(col_expr.clone(), filter, schema) .transpose() }) - .try_collect::>()?; + .collect::, _>>()?; if exprs.is_empty() { return Ok(None); @@ -326,19 +326,19 @@ impl LogQueryPlanner { let args = args .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; if let Some(alias) = alias { Ok(aggr_fn.call(args).alias(alias)) } else { Ok(aggr_fn.call(args)) } }) - .try_collect::>()?; + .collect::, _>>()?; let group_exprs = by .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; Ok((aggr_expr, group_exprs)) } @@ -380,7 +380,7 @@ impl LogQueryPlanner { let args = args .iter() .map(|expr| self.log_expr_to_df_expr(expr, schema)) - .try_collect::>()?; + .collect::, _>>()?; let func = self.session_state.scalar_functions().get(name).context( UnknownScalarFunctionSnafu { name: name.to_string(), diff --git a/src/query/src/optimizer/windowed_sort.rs b/src/query/src/optimizer/windowed_sort.rs index 3e4ccaf91c4..800bc103d38 100644 --- a/src/query/src/optimizer/windowed_sort.rs +++ b/src/query/src/optimizer/windowed_sort.rs @@ -199,12 +199,12 @@ fn fetch_partition_range(input: Arc) -> DataFusionResult, _>()?; + .collect::, _>>()?; let (min_val, max_val) = (typed_sorted_range[0].clone(), typed_sorted_range[1].clone()); @@ -2682,7 +2682,10 @@ mod test { let exec_stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); let real_output = exec_stream.collect::>().await; - let real_output: Vec<_> = real_output.into_iter().try_collect().unwrap(); + let real_output: Vec<_> = real_output + .into_iter() + .collect::, _>>() + .unwrap(); real_output } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 518d0f65e32..05956716d4b 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -231,7 +231,7 @@ impl PrometheusJsonResponse { ) -> Self { // Hold the collector while a streaming result is consumed. let collector = query_id.and_then(get_promql_annotation_collector); - let response: Result = try { + let response: Result = async { let result = result?; let mut resp = match result.data { @@ -259,8 +259,9 @@ impl PrometheusJsonResponse { resp.resp_metrics = re; } - resp - }; + Ok(resp) + } + .await; let result_type_string = result_type.to_string(); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 1106f89d3e9..2fed3d43de0 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] -#![feature(exclusive_wrapper)] - use datafusion_expr::LogicalPlan; use sql::statements::statement::Statement; // Re-export for use in add_service! macro diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 8a7b7272447..1fa075a8e17 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::fmt::Debug; -use std::sync::Exclusive; use ::auth::{ BEARER_TOKEN_USER, Identity, Password, PgAuthInfo, PgScramSha256Verifier, UserInfoRef, @@ -252,7 +251,7 @@ impl StartupHandler for PostgresServerHandlerInner { auth::save_startup_parameters_to_metadata(client, startup); // check if db is valid - match resolve_db_info(Exclusive::new(client), self.query_handler.clone()).await? { + match resolve_db_info(client, self.query_handler.clone()).await? { DbResolution::Resolved(catalog, schema) => { let metadata = client.metadata_mut(); let _ = metadata.insert(super::METADATA_CATALOG.to_owned(), catalog); @@ -595,13 +594,13 @@ enum DbResolution { /// A function extracted to resolve lifetime and readability issues: async fn resolve_db_info( - client: Exclusive<&mut C>, + client: &mut C, query_handler: ServerSqlQueryHandlerRef, ) -> PgWireResult where C: ClientInfo + Unpin + Send, { - let db_ref = client.into_inner().metadata().get(super::METADATA_DATABASE); + let db_ref = client.metadata().get(super::METADATA_DATABASE); if let Some(db) = db_ref { let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); if query_handler diff --git a/src/sql/src/lib.rs b/src/sql/src/lib.rs index b017f798f3f..e6bfe327902 100644 --- a/src/sql/src/lib.rs +++ b/src/sql/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(box_patterns)] - pub mod ast; pub mod dialect; pub mod error; diff --git a/src/sql/src/statements/insert.rs b/src/sql/src/statements/insert.rs index 52318cddbf4..9c069bde43d 100644 --- a/src/sql/src/statements/insert.rs +++ b/src/sql/src/statements/insert.rs @@ -14,7 +14,7 @@ use serde::Serialize; use sqlparser::ast::{ - Insert as SpInsert, ObjectName, ObjectNamePart, Parens, Query, SetExpr, Statement, TableObject, + Insert as SpInsert, ObjectName, ObjectNamePart, Parens, SetExpr, Statement, TableObject, UnaryOperator, ValueWithSpan, Values, }; use sqlparser::parser::ParserError; @@ -70,14 +70,18 @@ impl Insert { /// Extracts the literal insert statement body if possible pub fn values_body(&self) -> Result>> { match &self.inner { - Statement::Insert(SpInsert { - source: - Some(box Query { - body: box SetExpr::Values(Values { rows, .. }), - .. - }), - .. - }) => sql_exprs_to_values(rows), + Statement::Insert(SpInsert { source, .. }) => { + let rows = source + .as_deref() + .and_then(|query| match query.body.as_ref() { + SetExpr::Values(Values { rows, .. }) => Some(rows), + _ => None, + }); + match rows { + Some(rows) => sql_exprs_to_values(rows), + None => unreachable!(), + } + } _ => unreachable!(), } } @@ -86,36 +90,37 @@ impl Insert { /// The rules is the same as function `values_body()`. pub fn can_extract_values(&self) -> bool { match &self.inner { - Statement::Insert(SpInsert { - source: - Some(box Query { - body: box SetExpr::Values(Values { rows, .. }), - .. - }), - .. - }) => rows.iter().all(|es| { - es.iter().all(|expr| match expr { - Expr::Value(_) => true, - Expr::Identifier(ident) => { - if ident.quote_style.is_none() { - ident.value.to_lowercase() == "default" - } else { - ident.quote_style == Some('"') - } - } - Expr::UnaryOp { op, expr } => { - matches!(op, UnaryOperator::Minus | UnaryOperator::Plus) - && matches!( - &**expr, - Expr::Value(ValueWithSpan { - value: Value::Number(_, _), - .. - }) - ) - } - _ => false, + Statement::Insert(SpInsert { source, .. }) => source + .as_deref() + .and_then(|query| match query.body.as_ref() { + SetExpr::Values(Values { rows, .. }) => Some(rows), + _ => None, }) - }), + .is_some_and(|rows| { + rows.iter().all(|es| { + es.iter().all(|expr| match expr { + Expr::Value(_) => true, + Expr::Identifier(ident) => { + if ident.quote_style.is_none() { + ident.value.to_lowercase() == "default" + } else { + ident.quote_style == Some('"') + } + } + Expr::UnaryOp { op, expr } => { + matches!(op, UnaryOperator::Minus | UnaryOperator::Plus) + && matches!( + &**expr, + Expr::Value(ValueWithSpan { + value: Value::Number(_, _), + .. + }) + ) + } + _ => false, + }) + }) + }), _ => false, } } @@ -124,7 +129,7 @@ impl Insert { pub fn has_non_values_query_source(&self) -> bool { match &self.inner { Statement::Insert(SpInsert { - source: Some(box query), + source: Some(query), .. }) => !matches!(&*query.body, SetExpr::Values(_)), _ => false, @@ -134,9 +139,9 @@ impl Insert { pub fn query_body(&self) -> Result> { Ok(match &self.inner { Statement::Insert(SpInsert { - source: Some(box query), + source: Some(query), .. - }) => Some(query.clone().try_into()?), + }) => Some(query.as_ref().clone().try_into()?), _ => None, }) } @@ -383,11 +388,9 @@ mod tests { let q = insert.query_body().unwrap().unwrap(); assert!(insert.has_non_values_query_source()); assert!(matches!( - q.inner, - Query { - body: box SetExpr::Select { .. }, - .. - } + &q.inner, + sqlparser::ast::Query { body, .. } + if matches!(body.as_ref(), SetExpr::Select { .. }) )); } _ => unreachable!(), diff --git a/src/table/src/lib.rs b/src/table/src/lib.rs index abc7dc83002..9968ee4dc97 100644 --- a/src/table/src/lib.rs +++ b/src/table/src/lib.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![feature(try_blocks)] - pub mod dist_table; pub mod error; pub mod metadata; diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 934f838a33a..2db051a1e9c 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -144,8 +144,8 @@ impl RegionScanExec { )) }) .collect::>(); - let ts_col: Option = try { - PhysicalSortExpr::new( + let ts_col: Option = (|| { + Some(PhysicalSortExpr::new( Arc::new( Column::new_with_schema( &metadata.time_index_column().column_schema.name, @@ -157,8 +157,8 @@ impl RegionScanExec { descending: false, nulls_first: true, }, - ) - }; + )) + })(); let eq_props = match request.distribution { Some(TimeSeriesDistribution::PerSeries) => { From 822042c7bb24728df30a518eaf3e551f42589f0f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 24 Sep 2026 08:41:56 +0000 Subject: [PATCH 035/133] feat(log-store): chain object store WAL objects and recover the latest chain (#9334) * feat(log-store): chain object store WAL objects and recover the latest chain A conditional create can fail with an unknown outcome while its object is stored, or still lands later, and Mito reuses the row sequences of a failed append. Replaying such an object next to a later acknowledged one can let the unacknowledged rows win after a restart. Format version 2 gives every object header its writer's epoch, a link to the object it extends (sequence and writer instance) and a header CRC32, and allows objects without segments. Recovery replays only the chain ending at the complete object with the largest epoch and sequence: a link holds when its predecessor is present with the recorded writer instance, or is missing below every present object. Objects off the chain are orphans that are never replayed but keep their sequences. Each open writes an empty object that starts an epoch above every present object, linked to the recovered tip, before it accepts writes, so a late object of an earlier instance never ends the chain. A start object that meets an object of an earlier epoch moves to the next sequence; one of an equal or later epoch fails the open. Signed-off-by: jeremyhi * fix(log-store): break the WAL chain at every missing predecessor Accepting a missing predecessor below every present object lets a late object that lands below the chain change which links hold. Nothing collects objects yet, so a missing predecessor now always breaks the link, and recovery fails when objects are present but none completes a chain. Signed-off-by: jeremyhi * fix(log-store): keep the object store WAL format at version 1 The object store WAL has not been enabled anywhere, so no object in the previous layout exists and the chained header can stay version 1. Signed-off-by: jeremyhi * refactor(log-store): link object store WAL objects by epoch Only one store instance writes under an epoch, since every open starts an epoch above every present object and a start object that meets the same or a later epoch fails the open. The epoch therefore identifies the instance, and the random writer instance id is dropped from the header. A link now records the sequence and the epoch of the object it extends, and holds when the predecessor carries that epoch. The header shrinks to 46 bytes, and the store logs its epoch when it opens. Signed-off-by: jeremyhi * fix(log-store): fail an open whose start object is already present Without a random writer instance, two opens that recover the same objects encode byte-identical start objects, and a conditional create treats the same bytes as its own retry. A start object that is already present therefore fails the open with a retryable error instead of letting both opens claim the epoch; the next open counts the object and starts a later epoch. Signed-off-by: jeremyhi * fix(log-store): derive the WAL epoch from the claimed start sequence Two opens can recover different listings when a late object lands across a sequence gap between them, pick the same largest epoch plus one, and both create their start objects under different sequences. The epoch of an instance is now one above the sequence its start object claims, so a successful create decides the epoch and no two instances share one. It stays above every epoch recovery listed, and an object that carries an epoch above the next sequence fails the open. Signed-off-by: jeremyhi * test(log-store): cover a WAL start object stored with an unknown outcome Add a fault-injection test in which the create of the start object stores the object but reports an error: the open fails without moving to another sequence, and the next open counts the stored object and claims a later epoch. Rename the test helper that writes a whole object from a given header to put_object_with_header, and drop a needless clone. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/error.rs | 18 +- src/log-store/src/object_store_wal.rs | 27 +- src/log-store/src/object_store_wal/catalog.rs | 20 +- src/log-store/src/object_store_wal/format.rs | 188 ++-- src/log-store/src/object_store_wal/store.rs | 942 +++++++++++++++--- 5 files changed, 968 insertions(+), 227 deletions(-) diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index de304e192d2..8be0b8bd3da 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -374,6 +374,18 @@ pub enum Error { location: Location, }, + #[snafu(display( + "WAL object {} that starts epoch {} was already present, so the open cannot tell whether it wrote it", + path, + epoch + ))] + UnconfirmedWalEpochStart { + path: String, + epoch: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to {} WAL object, path: {}", operation, path))] WalObjectStore { operation: &'static str, @@ -481,6 +493,7 @@ impl ErrorExt for Error { | WriteIndex { .. } | ReadIndex { .. } | WalObjectStore { .. } + | UnconfirmedWalEpochStart { .. } | Io { .. } => StatusCode::StorageUnavailable, // Raft engine FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => { @@ -513,7 +526,10 @@ impl ErrorExt for Error { | WalObjectStore { error, .. } => retry_hint_from_opendal_error(error), ObjectStoreWal { source, .. } => source.retry_hint(), Io { error, .. } => retry_hint_from_io_error(error), - FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } => RetryHint::Retryable, + FetchEntry { .. } + | RaftEngine { .. } + | AddEntryLogBatch { .. } + | UnconfirmedWalEpochStart { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { rskafka::client::producer::Error::Client(error) => { rskafka_client_error_to_retry_hint(error) diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 70bc39dedca..3c1958b5cdb 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -23,7 +23,8 @@ //! ``` //! //! The header carries the magic `GTWALOBJ`, the format version, the object -//! sequence and the instance that wrote the object. Each segment holds the +//! sequence, the epoch of the instance that wrote the object, the link to the +//! object it extends, and a CRC32 of the header. Each segment holds the //! entries of exactly one region, ordered by entry id, and segments are ordered //! by region id. The footer indexes every segment with its region id, entry id //! range, byte range and CRC32. The fixed-size trailer points at the footer and @@ -31,13 +32,25 @@ //! the footer by reading the fixed-length trailer at the end of the object. //! //! Object sequences increase monotonically within one prefix and may leave -//! gaps, so recovery continues after the largest sequence it indexed. An object -//! is created conditionally: rewriting a sequence with the content it already +//! gaps, so recovery continues after the largest sequence present. An object is +//! created conditionally: rewriting a sequence with the content it already //! holds is a no-op at the object store, while different content under a taken -//! sequence is a conflict. Recovery lists the objects, reads and verifies only -//! the header, trailer and footer of each, and indexes the footers in sequence -//! order to rebuild the object catalog, which rejects a sequence it already -//! holds. Segments are read and checksummed only when a read decodes them. +//! sequence is a conflict. A create can fail with an unknown outcome and its +//! object can still appear, so objects form a chain: every object links to the +//! object it extends by sequence and epoch, and recovery replays only the chain +//! that ends at the complete object with the largest epoch and sequence. An +//! object is complete when every link on its chain names a present object of +//! the recorded epoch, back to an object that starts the chain. Objects off the +//! chain are orphans, which are never replayed but keep their sequences. Each +//! open writes an object without segments above every present object before +//! it accepts writes, and its epoch is one above the sequence of that object, +//! so no two instances share an epoch and a late object of an earlier +//! instance never ends the chain. +//! +//! Recovery lists the objects, reads and verifies only the header, trailer and +//! footer of each, and indexes the footers of the chain in sequence order to +//! rebuild the object catalog, which rejects a sequence it already holds. +//! Segments are read and checksummed only when a read decodes them. //! //! Entry ids are object-sequence-major, see [`entry_id`]: the high bits of an //! id name the object that holds the entry, the low bits its position among diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs index 5e9650d14a5..4866da359cf 100644 --- a/src/log-store/src/object_store_wal/catalog.rs +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -38,19 +38,13 @@ impl ObjectCatalog { /// Indexes the footer of the object `object_seq`. Objects may be inserted /// in any order, which lets recovery index them as it discovers them. /// Inserting a sequence that is already indexed is rejected, whether or not - /// the footer matches the indexed one. + /// the footer matches the indexed one. An object without segments holds no + /// entries but still takes its sequence. pub(crate) fn insert_object( &mut self, object_seq: u64, mut footer: Vec, ) -> Result<()> { - ensure!( - !footer.is_empty(), - CorruptedWalObjectSnafu { - reason: format!("object {object_seq} has an empty footer"), - } - ); - footer.sort_unstable_by_key(|entry| entry.region_id); for entries in footer.windows(2) { ensure!( @@ -310,6 +304,11 @@ mod tests { .unwrap(); assert_eq!(5, catalog.next_object_seq().unwrap()); + + // An object without segments takes its sequence all the same. + catalog.insert_object(7, Vec::new()).unwrap(); + assert_eq!(8, catalog.next_object_seq().unwrap()); + assert_eq!(Some(12), catalog.region_max_entry_id(region_id)); } #[test] @@ -416,11 +415,6 @@ mod tests { ), "duplicate footer entries", ); - assert_corrupted( - catalog.insert_object(1, vec![]), - "object 1 has an empty footer", - ); - let mut invalid = footer_entry(region_id, 2, 1); invalid.entry_count = 0; assert_corrupted( diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index 62fe290aa82..b3d46049523 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -27,9 +27,13 @@ use crate::error::{CorruptedWalObjectSnafu, Result}; const HEADER_MAGIC: &[u8; 8] = b"GTWALOBJ"; const TRAILER_MAGIC: &[u8; 8] = b"GTWALTRL"; const FORMAT_VERSION: u16 = 1; +/// Predecessor sequence recorded for an object that starts a chain. Object +/// sequences stay below `OBJECT_SEQ_LIMIT`, so it never names an object. +const NO_PREDECESSOR: u64 = u64::MAX; -/// Length of the object header: magic, version, object sequence, writer instance. -pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 16; +/// Length of the object header: magic, version, object sequence, epoch, +/// predecessor sequence, predecessor epoch and the CRC32 of the header. +pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 8 + 8 + 8 + 4; /// Length of the fixed trailer: footer offset, footer length, footer CRC32, /// object CRC32 and magic. pub(crate) const TRAILER_LEN: usize = 8 + 8 + 4 + 4 + 8; @@ -44,7 +48,20 @@ const FOOTER_COUNT_LEN: usize = 4; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Header { pub(crate) object_seq: u64, - pub(crate) writer_instance: [u8; 16], + /// Epoch of the store instance that wrote the object: one above the + /// sequence of the start object the instance created, so the epoch also + /// identifies the instance. + pub(crate) epoch: u64, + /// The object this object extends, `None` for the first object of a chain. + pub(crate) prev: Option, +} + +/// Names the object another object extends: its sequence and the epoch +/// recorded in its header. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ChainLink { + pub(crate) object_seq: u64, + pub(crate) epoch: u64, } /// A single WAL entry inside an object. @@ -94,15 +111,9 @@ pub(crate) struct DecodedObject { } /// Encodes `records` into one object. Records are grouped into a segment per -/// region; entry ids within a region must be unique. +/// region; entry ids within a region must be unique. An object without records +/// only extends the chain. pub(crate) fn encode_object(header: Header, records: &[Record]) -> Result { - ensure!( - !records.is_empty(), - CorruptedWalObjectSnafu { - reason: "object has no records", - } - ); - let mut grouped = BTreeMap::>::new(); for record in records { grouped.entry(record.region_id).or_default().push(record); @@ -180,7 +191,8 @@ pub(crate) fn object_len(footer: &[FooterEntry]) -> u64 { segments + framing as u64 } -/// Decodes the header from the first [`HEADER_LEN`] bytes of an object. +/// Decodes and verifies the header from the first [`HEADER_LEN`] bytes of an +/// object. pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ ensure!( bytes.len() >= HEADER_LEN, @@ -207,13 +219,25 @@ pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ } ); let object_seq = reader.u64("header")?; - let writer_instance = reader - .take(16, "header")? - .try_into() - .expect("writer instance has a fixed length"); + let epoch = reader.u64("header")?; + let prev_object_seq = reader.u64("header")?; + let prev_epoch = reader.u64("header")?; + let header_crc32 = reader.u32("header")?; + let checksum = crc32fast::hash(&bytes[..HEADER_LEN - 4]); + ensure!( + checksum == header_crc32, + CorruptedWalObjectSnafu { + reason: checksum_mismatch("header", header_crc32, checksum), + } + ); + let prev = (prev_object_seq != NO_PREDECESSOR).then_some(ChainLink { + object_seq: prev_object_seq, + epoch: prev_epoch, + }); Ok(Header { object_seq, - writer_instance, + epoch, + prev, }) } @@ -274,12 +298,6 @@ pub(crate) fn decode_footer(bytes: &[u8], trailer: FixedTrailer) -> Result 0, - CorruptedWalObjectSnafu { - reason: "object has no records", - } - ); let expected_len = count .checked_mul(FOOTER_ENTRY_LEN) .and_then(|len| len.checked_add(4)) @@ -551,10 +569,18 @@ pub(crate) fn decode_object(bytes: &[u8]) -> Result { } fn encode_header(header: &Header, output: &mut BytesMut) { + let start = output.len(); output.put_slice(HEADER_MAGIC); output.put_u16(FORMAT_VERSION); output.put_u64(header.object_seq); - output.put_slice(&header.writer_instance); + output.put_u64(header.epoch); + let prev = header.prev.unwrap_or(ChainLink { + object_seq: NO_PREDECESSOR, + epoch: 0, + }); + output.put_u64(prev.object_seq); + output.put_u64(prev.epoch); + output.put_u32(crc32fast::hash(&output[start..])); } fn encode_footer(directory: &[FooterEntry], output: &mut BytesMut) -> Result<()> { @@ -697,7 +723,11 @@ mod tests { fn header() -> Header { Header { object_seq: 42, - writer_instance: *b"writer-instance!", + epoch: 5, + prev: Some(ChainLink { + object_seq: 40, + epoch: 4, + }), } } @@ -804,13 +834,23 @@ mod tests { bad_version[8..10].copy_from_slice(&(FORMAT_VERSION + 1).to_be_bytes()); assert_corrupted(decode_object(&bad_version), "unsupported format version 2"); - let mut bad_object_seq = encoded.bytes.to_vec(); - bad_object_seq[10] ^= 1; - assert_corrupted(decode_object(&bad_object_seq), "object checksum mismatch"); + // Every field after the version is covered by the header checksum, + // which recovery verifies without reading the whole object. + for offset in [10, 18, 26, 34, HEADER_LEN - 1] { + let mut bad_field = encoded.bytes.to_vec(); + bad_field[offset] ^= 1; + assert_corrupted(decode_header(&bad_field), "header checksum mismatch"); + } + } - let mut bad_writer = encoded.bytes.to_vec(); - bad_writer[18] ^= 1; - assert_corrupted(decode_object(&bad_writer), "object checksum mismatch"); + #[test] + fn test_format_round_trips_the_first_object_of_a_chain() { + let first = Header { + prev: None, + ..header() + }; + let encoded = encode_object(first.clone(), &records()).unwrap(); + assert_eq!(first, decode_object(&encoded.bytes).unwrap().header); } #[test] @@ -864,41 +904,19 @@ mod tests { } #[test] - fn test_format_rejects_empty_object() { - assert_corrupted(encode_object(header(), &[]), "object has no records"); - } - - #[test] - fn test_format_rejects_checksum_valid_empty_footer() { - let mut object = BytesMut::new(); - encode_header(&header(), &mut object); - let footer_offset = object.len(); - encode_footer(&[], &mut object).unwrap(); - let footer = object[footer_offset..].to_vec(); - let trailer = FixedTrailer { - footer_offset: footer_offset as u64, - footer_len: footer.len() as u64, - footer_crc32: crc32fast::hash(&footer), - object_crc32: 0, - }; - let trailer = FixedTrailer { - object_crc32: object_crc32(&object, trailer), - ..trailer - }; - encode_trailer(trailer, &mut object); - - // Every structural check that recovery runs before the footer passes. - assert_eq!(header(), decode_header(&object).unwrap()); - let trailer_start = object.len() - TRAILER_LEN; - assert_eq!(trailer, decode_trailer(&object[trailer_start..]).unwrap()); + fn test_format_round_trips_an_object_without_records() { + let encoded = encode_object(header(), &[]).unwrap(); + assert!(encoded.footer.is_empty()); assert_eq!( - footer_offset..trailer_start, - footer_range(trailer, object.len()).unwrap() + MIN_OBJECT_LEN, + encoded.bytes.len(), + "an empty object is a header, an empty footer and the trailer" ); - verify_segment_ranges(&[], footer_offset).unwrap(); - - assert_corrupted(decode_footer(&footer, trailer), "object has no records"); - assert_corrupted(decode_object(&object), "object has no records"); + assert_eq!(encoded.bytes.len() as u64, object_len(&encoded.footer)); + let decoded = decode_object(&encoded.bytes).unwrap(); + assert_eq!(header(), decoded.header); + assert!(decoded.records.is_empty()); + assert!(decoded.footer.is_empty()); } #[test] @@ -950,12 +968,16 @@ mod tests { /// to field order, endianness or checksum coverage fails this test even when /// the encoder and the decoder change together. const FIXTURE_V1_HEX: &str = concat!( - // Header: magic, version 1, object sequence 7, writer instance. + // Header: magic, version 1, object sequence 7, epoch 3, predecessor + // sequence 6, predecessor epoch 2, header CRC32. "475457414c4f424a", "0001", "0000000000000007", - "7772697465722d666978747572652d31", - // Segment of region 1 at offset 34: region id, entry count, then + "0000000000000003", + "0000000000000006", + "0000000000000002", + "1ca121b1", + // Segment of region 1 at offset 46: region id, entry count, then // (entry id, payload length, payload) per entry. "0000000100000001", "00000003", @@ -968,7 +990,7 @@ mod tests { "0000000000000003", "00000003", "636363", - // Segment of region 2 at offset 88. + // Segment of region 2 at offset 100. "0000000200000001", "00000002", "000000000000000a", @@ -977,28 +999,28 @@ mod tests { "000000000000000b", "00000002", "7979", - // Footer at offset 127: entry count, then per segment region id, min and + // Footer at offset 139: entry count, then per segment region id, min and // max entry id, entry count, segment offset, length and CRC32. "00000002", "0000000100000001", "0000000000000001", "0000000000000003", "00000003", - "0000000000000022", + "000000000000002e", "0000000000000036", "25ac0486", "0000000200000001", "000000000000000a", "000000000000000b", "00000002", - "0000000000000058", + "0000000000000064", "0000000000000027", "65dc08ec", // Trailer: footer offset, footer length, footer CRC32, object CRC32, magic. - "000000000000007f", + "000000000000008b", "0000000000000064", - "8298e0f7", - "7a009db1", + "6b97e3e4", + "631bb7d1", "475457414c54524c", ); @@ -1031,7 +1053,11 @@ mod tests { let fixture = fixture_bytes(); let header = Header { object_seq: 7, - writer_instance: *b"writer-fixture-1", + epoch: 3, + prev: Some(ChainLink { + object_seq: 6, + epoch: 2, + }), }; let footer = vec![ FooterEntry { @@ -1039,7 +1065,7 @@ mod tests { min_entry_id: 1, max_entry_id: 3, entry_count: 3, - segment_offset: 34, + segment_offset: 46, segment_len: 54, segment_crc32: 0x25ac0486, }, @@ -1048,12 +1074,12 @@ mod tests { min_entry_id: 10, max_entry_id: 11, entry_count: 2, - segment_offset: 88, + segment_offset: 100, segment_len: 39, segment_crc32: 0x65dc08ec, }, ]; - assert_eq!(259, fixture.len()); + assert_eq!(271, fixture.len()); let decoded = decode_object(&fixture).unwrap(); assert_eq!(header, decoded.header); @@ -1061,10 +1087,10 @@ mod tests { assert_eq!(fixture_records(), decoded.records); assert_eq!( FixedTrailer { - footer_offset: 127, + footer_offset: 139, footer_len: 100, - footer_crc32: 0x8298e0f7, - object_crc32: 0x7a009db1, + footer_crc32: 0x6b97e3e4, + object_crc32: 0x631bb7d1, }, decode_trailer(&fixture[fixture.len() - TRAILER_LEN..]).unwrap() ); diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index c0d9cc89363..7b024bb1a80 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -14,7 +14,7 @@ //! Object store WAL construction, recovery and region reads. -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::fmt; use std::ops::Range; use std::sync::atomic::{AtomicBool, Ordering}; @@ -23,6 +23,7 @@ use std::time::Duration; use async_stream::try_stream; use bytes::Bytes; +use common_telemetry::info; use common_wal::config::object_store::ObjectStoreWalConfig; use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStore; @@ -36,12 +37,15 @@ use tokio::sync::{mpsc, oneshot}; use crate::error::{ CorruptedWalObjectSnafu, Error, InvalidProviderSnafu, InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, MismatchedWalRegionSnafu, - ObjectStoreWalSnafu, Result, UnsupportedObjectStoreWalOperationSnafu, + ObjectStoreWalSnafu, Result, UnconfirmedWalEpochStartSnafu, + UnsupportedObjectStoreWalOperationSnafu, WalObjectSequenceExhaustedSnafu, }; +use crate::object_store_wal::batch::OBJECT_SEQ_LIMIT; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ - FixedTrailer, FooterEntry, HEADER_LEN, MIN_OBJECT_LEN, TRAILER_LEN, decode_footer, - decode_header, decode_segment, decode_trailer, footer_range, verify_segment_ranges, + ChainLink, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, TRAILER_LEN, + decode_footer, decode_header, decode_segment, decode_trailer, encode_object, footer_range, + verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; @@ -79,8 +83,9 @@ impl fmt::Debug for ObjectStoreLogStore { impl ObjectStoreLogStore { /// Builds the store under the node and generation prefix derived from `config`, - /// recovering the catalog from the objects that already exist. Recovery - /// fails on the first corrupted or conflicting object. + /// recovering the catalog from the objects that already exist and writing + /// the object that starts the epoch of this instance. Recovery fails on the + /// first corrupted or conflicting object. pub(crate) async fn try_new( object_store: ObjectStore, config: &ObjectStoreWalConfig, @@ -108,7 +113,36 @@ impl ObjectStoreLogStore { ); positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; - let (catalog, next_object_seq, durable_entry_ids) = recover(io.as_ref()).await?; + let Recovered { + mut catalog, + next_object_seq, + durable_entry_ids, + tip, + max_epoch, + } = recover(io.as_ref()).await?; + // Every epoch is one above the sequence of the start object its + // instance created, and every object is at or above its start object, + // so an epoch above the next sequence names no instance that ran. + ensure!( + max_epoch <= next_object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "an object carries epoch {max_epoch}, above the next sequence {next_object_seq}" + ), + } + ); + let start = start_epoch(io.as_ref(), next_object_seq, tip).await?; + let epoch = start.epoch; + // The epoch identifies this instance in every object it writes. + info!( + "Opened object store WAL under {prefix} at epoch {epoch}, start object {}", + start.object_seq + ); + catalog + .insert_object(start.object_seq, Vec::new()) + .with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(start.object_seq), + })?; let catalog = Arc::new(RwLock::new(catalog)); let terminal_error = TerminalError::default(); let stopped = Arc::new(AtomicBool::new(false)); @@ -118,7 +152,12 @@ impl ObjectStoreLogStore { stopped: stopped.clone(), command_rx, issued_entry_ids: durable_entry_ids, - next_object_seq: Some(next_object_seq), + next_object_seq: start + .object_seq + .checked_add(1) + .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT), + epoch, + chain_tip: start, stop: Vec::new(), }; common_runtime::spawn_global(actor.run()); @@ -363,6 +402,9 @@ struct Actor { command_rx: mpsc::Receiver, issued_entry_ids: HashMap, next_object_seq: Option, + epoch: u64, + /// The object a newly sealed batch extends. + chain_tip: ChainLink, stop: Vec>>, } @@ -417,49 +459,228 @@ fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) } -/// Rebuilds the catalog from object footers, +/// What recovery rebuilt from the objects under a prefix. +#[derive(Debug)] +struct Recovered { + catalog: ObjectCatalog, + /// Sequence above every present object and every id the catalog holds. + next_object_seq: u64, + durable_entry_ids: HashMap, + /// The object the next object extends, `None` on an empty prefix. + tip: Option, + /// Largest epoch any present object carries, zero on an empty prefix. + max_epoch: u64, +} + +/// An object as recovery fetched it: its key, header and footer. +struct FetchedObject { + object: ListedObject, + header: Header, + footer: Vec, +} + +/// Rebuilds the catalog from object headers and footers, /// so recovery costs a few small reads per object however large the objects /// are. Segments are not read; a segment checksum is verified by the read that /// decodes it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects /// at a time and indexed in sequence order, so the catalog checks the entry /// ranges of every object against its predecessors like a sequential replay. -async fn recover(io: &dyn WalObjectIo) -> Result<(ObjectCatalog, u64, HashMap)> { +/// Only the objects on the chain [`select_chain`] picks are indexed. +async fn recover(io: &dyn WalObjectIo) -> Result { let objects = io.list().await?; + finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?) +} + +/// Indexes the chain of `objects`, which are ordered by sequence. Objects off +/// the chain are orphans: they are not indexed, but no later object takes +/// their sequences. +fn finish_recovery(objects: Vec) -> Result { + let headers = objects + .iter() + .map(|fetched| (fetched.object.object_seq, &fetched.header)) + .collect::>(); + let chain = select_chain(&headers); + // Every open writes an object that starts a chain or extends a complete + // one, and nothing removes objects, so present objects without any + // complete chain mean objects of the chain are gone. + ensure!( + headers.is_empty() || !chain.is_empty(), + CorruptedWalObjectSnafu { + reason: format!( + "no object among {} present objects completes a chain", + headers.len() + ), + } + ); + let tip = chain.last().map(|object_seq| ChainLink { + object_seq: *object_seq, + epoch: headers[object_seq].epoch, + }); + let max_epoch = headers + .values() + .map(|header| header.epoch) + .max() + .unwrap_or(0); + let after_listed = match headers.last_key_value() { + None => 0, + Some((&last_object_seq, _)) => last_object_seq + .checked_add(1) + .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, + }; + let chain = chain.into_iter().collect::>(); + let mut catalog = ObjectCatalog::default(); - for (object, footer) in fetch_footers(io, objects, RECOVERY_CONCURRENCY).await? { - catalog - .insert_object(object.object_seq, footer) - .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + for FetchedObject { object, footer, .. } in objects { + if chain.contains(&object.object_seq) { + catalog + .insert_object(object.object_seq, footer) + .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + } } - finish_recovery(catalog) -} - -fn finish_recovery( - catalog: ObjectCatalog, -) -> Result<(ObjectCatalog, u64, HashMap)> { - let next_object_seq = catalog.next_object_seq()?; + let next_object_seq = catalog.next_object_seq()?.max(after_listed); + ensure!( + next_object_seq < OBJECT_SEQ_LIMIT, + WalObjectSequenceExhaustedSnafu { + last_object_seq: next_object_seq - 1, + } + ); let durable_entry_ids = durable_entry_ids(&catalog); - Ok((catalog, next_object_seq, durable_entry_ids)) + Ok(Recovered { + catalog, + next_object_seq, + durable_entry_ids, + tip, + max_epoch, + }) } -/// Fetches and verifies the footers of `objects`, up to `concurrency` objects -/// at a time, and returns them ordered by object sequence whatever the order -/// the fetches complete in. The first failure abandons the remaining fetches. +/// Returns, in sequence order, the objects on the chain that ends at the tip: +/// the complete object with the largest epoch, then the largest sequence. +/// +/// An object is complete when every link on its chain holds: the chain starts +/// at an object without a predecessor, and every other link names a present +/// object that carries the epoch the link records. Only one instance writes +/// under an epoch, so the epoch tells an object of the linking instance from +/// another object under the same sequence. A create that +/// was reported as failed may still leave its object, but every object +/// written after that failure links past it, and every instance writes under +/// an epoch above every object present when it opened, so neither such an +/// object nor a late object of an earlier instance ends the chosen chain. +fn select_chain(headers: &BTreeMap) -> Vec { + // A predecessor precedes its successor, so one pass in sequence order + // settles every object. + let mut complete = HashMap::with_capacity(headers.len()); + for (&object_seq, header) in headers { + let holds = match header.prev { + None => true, + Some(link) if link.object_seq >= object_seq => false, + Some(link) => match headers.get(&link.object_seq) { + Some(prev) => prev.epoch == link.epoch && complete[&link.object_seq], + None => false, + }, + }; + complete.insert(object_seq, holds); + } + let Some(tip) = headers + .iter() + .filter(|(object_seq, _)| complete[*object_seq]) + .max_by_key(|(object_seq, header)| (header.epoch, **object_seq)) + .map(|(object_seq, _)| *object_seq) + else { + return Vec::new(); + }; + let mut chain = vec![tip]; + while let Some(link) = headers[chain.last().expect("the chain holds the tip")].prev + && headers.contains_key(&link.object_seq) + { + chain.push(link.object_seq); + } + chain.reverse(); + chain +} + +/// Writes the empty object that starts the epoch of this instance at +/// `object_seq`, linked to the recovered `tip`, and returns it as the tip +/// later objects extend. The epoch is one above the sequence the create +/// claims, so no two instances share one however their opens interleave, and +/// it is above the epoch of every object recovery listed. +/// +/// An object an earlier instance left at that sequence after recovery listed +/// the prefix has a lower epoch and never ends a chain, so the start object +/// moves to the next sequence and epoch. An object of an equal or later epoch +/// belongs to another writer of the prefix, and the conflict fails the open, +/// as does a create whose outcome is unknown: the next open counts the object +/// in either case. So does a create that finds the same bytes present: +/// another open that recovered the same objects writes an identical start +/// object, so this open cannot claim the epoch even if the object is its own. +async fn start_epoch( + io: &dyn WalObjectIo, + mut object_seq: u64, + tip: Option, +) -> Result { + loop { + ensure!( + object_seq < OBJECT_SEQ_LIMIT, + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + ); + let epoch = object_seq + 1; + let header = Header { + object_seq, + epoch, + prev: tip, + }; + match io + .put_if_absent(object_seq, encode_object(header, &[])?.bytes) + .await + { + Ok(PutResult::Created) => return Ok(ChainLink { object_seq, epoch }), + Ok(PutResult::AlreadyPresent) => { + return UnconfirmedWalEpochStartSnafu { + path: io.object_path(object_seq), + epoch, + } + .fail(); + } + Err(error @ Error::WalObjectConflict { .. }) => { + let head = io.get_range(object_seq, 0, HEADER_LEN as u64).await?; + let existing = decode_header(&head).with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + })?; + if existing.epoch >= epoch { + return Err(error); + } + object_seq += 1; + } + Err(error) => return Err(error), + } + } +} + +/// Fetches and verifies the headers and footers of `objects`, up to +/// `concurrency` objects at a time, and returns them ordered by object sequence +/// whatever the order the fetches complete in. The first failure abandons the +/// remaining fetches. async fn fetch_footers( io: &dyn WalObjectIo, objects: Vec, concurrency: usize, -) -> Result)>> { - let mut footers = futures::stream::iter(objects) +) -> Result> { + let mut fetched = futures::stream::iter(objects) .map(|object| async move { - let footer = fetch_footer(io, &object).await?; - Ok((object, footer)) + let (header, footer) = fetch_footer(io, &object).await?; + Ok(FetchedObject { + object, + header, + footer, + }) }) .buffer_unordered(concurrency) .try_collect::>() .await?; - footers.sort_unstable_by_key(|(object, _)| object.object_seq); - Ok(footers) + fetched.sort_unstable_by_key(|fetched| fetched.object.object_seq); + Ok(fetched) } /// Reads the header, trailer and footer of `object` and verifies them: the @@ -470,7 +691,10 @@ async fn fetch_footers( /// A short object is read whole. Otherwise the header and a window at the end /// of the object are read concurrently, and the footer is read separately only /// when it starts before the window. -async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result> { +async fn fetch_footer( + io: &dyn WalObjectIo, + object: &ListedObject, +) -> Result<(Header, Vec)> { let ListedObject { object_seq, size, .. } = *object; @@ -514,7 +738,7 @@ async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result= tail_start { tail.slice(footer_range.start - tail_start..footer_range.end - tail_start) @@ -528,18 +752,18 @@ async fn fetch_footer(io: &dyn WalObjectIo, object: &ListedObject) -> Result Result<(FixedTrailer, Range)> { +) -> Result<(Header, FixedTrailer, Range)> { let header = decode_header(head)?; ensure!( header.object_seq == object_seq, @@ -552,7 +776,7 @@ fn locate_footer( ); let trailer = decode_trailer(&tail[tail.len() - TRAILER_LEN..])?; let footer_range = footer_range(trailer, object_len)?; - Ok((trailer, footer_range)) + Ok((header, trailer, footer_range)) } fn durable_entry_ids(catalog: &ObjectCatalog) -> HashMap { @@ -693,26 +917,21 @@ mod tests { let config = ObjectStoreWalConfig::default(); let identities = [(1, 2), (3, 2), (1, 4)]; for (index, (node_id, generation)) in identities.iter().enumerate() { - let encoded = encode_object( - Header { - object_seq: 0, - writer_instance: [0; 16], - }, + let io = ObjectStoreIo::new( + object_store.clone(), + config.node_prefix(*node_id, *generation), + ) + .unwrap(); + put_fixture( + &io, + 0, &[Record { region_id: region(1), entry_id: index as u64 + 1, payload: Bytes::from_static(b"entry"), }], ) - .unwrap(); - ObjectStoreIo::new( - object_store.clone(), - config.node_prefix(*node_id, *generation), - ) - .unwrap() - .put_if_absent(0, encoded.bytes) - .await - .unwrap(); + .await; } // A root-level object must not be part of any node's recovery. let root_io = ObjectStoreIo::new(object_store.clone(), &config.prefix).unwrap(); @@ -749,31 +968,44 @@ mod tests { } /// Rebuilds the catalog by decoding whole objects as a recovery oracle. - async fn recover_by_decoding( - io: &dyn WalObjectIo, - ) -> Result<(ObjectCatalog, u64, HashMap)> { - let mut catalog = ObjectCatalog::default(); - for ListedObject { - object_seq, path, .. - } in io.list().await? - { - let bytes = io.get(object_seq).await?; - decode_object(&bytes) + async fn recover_by_decoding(io: &dyn WalObjectIo) -> Result { + let mut objects = Vec::new(); + for object in io.list().await? { + let bytes = io.get(object.object_seq).await?; + let decoded = decode_object(&bytes) .and_then(|decoded| { ensure!( - decoded.header.object_seq == object_seq, + decoded.header.object_seq == object.object_seq, CorruptedWalObjectSnafu { reason: format!( - "header sequence {} does not match key sequence {object_seq}", - decoded.header.object_seq + "header sequence {} does not match key sequence {}", + decoded.header.object_seq, object.object_seq ), } ); - catalog.insert_object(object_seq, decoded.footer) + Ok(decoded) }) - .with_context(|_| InvalidWalObjectSnafu { path })?; + .with_context(|_| InvalidWalObjectSnafu { + path: object.path.clone(), + })?; + objects.push(FetchedObject { + object, + header: decoded.header, + footer: decoded.footer, + }); } - finish_recovery(catalog) + finish_recovery(objects) + } + + fn assert_same_recovery(expected: &Recovered, actual: &Recovered) { + assert_eq!( + catalog_contents(&expected.catalog), + catalog_contents(&actual.catalog) + ); + assert_eq!(expected.next_object_seq, actual.next_object_seq); + assert_eq!(expected.durable_entry_ids, actual.durable_entry_ids); + assert_eq!(expected.tip, actual.tip); + assert_eq!(expected.max_epoch, actual.max_epoch); } fn catalog_contents(catalog: &ObjectCatalog) -> Vec<(u64, Vec)> { @@ -798,19 +1030,45 @@ mod tests { } async fn put_records(object_store: &ObjectStore, object_seq: u64, records: &[Record]) { - let encoded = encode_object( - Header { - object_seq, - writer_instance: [0; 16], - }, - records, - ) - .unwrap(); - ObjectStoreIo::new(object_store.clone(), PREFIX) - .unwrap() - .put_if_absent(object_seq, encoded.bytes) + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&io, object_seq, records).await; + } + + /// The header of a fixture object at `object_seq`, which extends the + /// present object right below it under the epoch of that object, so the + /// fixtures a test writes in sequence order form one chain. + async fn fixture_header(io: &ObjectStoreIo, object_seq: u64) -> Header { + let below = io + .list() .await - .unwrap(); + .unwrap() + .into_iter() + .map(|object| object.object_seq) + .filter(|seq| *seq < object_seq) + .max(); + let prev = match below { + Some(seq) => Some(decode_header(&io.get(seq).await.unwrap()).unwrap()), + None => None, + }; + Header { + object_seq, + epoch: prev.as_ref().map_or(1, |prev| prev.epoch), + prev: prev.map(|prev| ChainLink { + object_seq: prev.object_seq, + epoch: prev.epoch, + }), + } + } + + async fn put_fixture(io: &ObjectStoreIo, object_seq: u64, records: &[Record]) { + let header = fixture_header(io, object_seq).await; + put_object_with_header(io, header, records).await; + } + + async fn put_object_with_header(io: &ObjectStoreIo, header: Header, records: &[Record]) { + let object_seq = header.object_seq; + let encoded = encode_object(header, records).unwrap(); + io.put_if_absent(object_seq, encoded.bytes).await.unwrap(); } fn object_path(object_store: &ObjectStore, object_seq: u64) -> String { @@ -862,18 +1120,13 @@ mod tests { populate(&object_store, 40, 5).await; let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); - let (catalog, next_object_seq, durable) = recover(&io).await.unwrap(); - let (expected_catalog, expected_next_object_seq, expected_durable) = - recover_by_decoding(&io).await.unwrap(); + let recovered = recover(&io).await.unwrap(); + let expected = recover_by_decoding(&io).await.unwrap(); - assert_eq!(40, catalog_contents(&catalog).len()); - assert_eq!( - catalog_contents(&expected_catalog), - catalog_contents(&catalog) - ); - assert_eq!(expected_next_object_seq, next_object_seq); - assert_eq!(40, next_object_seq); - assert_eq!(expected_durable, durable); + assert_eq!(40, catalog_contents(&recovered.catalog).len()); + assert_same_recovery(&expected, &recovered); + assert_eq!(40, recovered.next_object_seq); + let durable = recovered.durable_entry_ids; assert_eq!(5, durable.len()); let store = open(object_store, &eager()).await; @@ -934,10 +1187,7 @@ mod tests { let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let encoded = encode_object( - Header { - object_seq: 0, - writer_instance: [0; 16], - }, + fixture_header(&io, 0).await, &[Record { region_id: region(1), entry_id: 1, @@ -1036,16 +1286,10 @@ mod tests { let (trailer, _) = footer_of(&object_store.read(&wide.path).await.unwrap().to_vec()); assert!(trailer.footer_len > RECOVERY_TAIL_WINDOW as u64); - let (catalog, next_object_seq, durable) = recover(io.as_ref()).await.unwrap(); - let (expected_catalog, expected_next_object_seq, expected_durable) = - recover_by_decoding(io.as_ref()).await.unwrap(); - assert_eq!( - catalog_contents(&expected_catalog), - catalog_contents(&catalog) - ); - assert_eq!(expected_next_object_seq, next_object_seq); - assert_eq!(expected_durable, durable); - assert_eq!(regions as usize, durable.len()); + let recovered = recover(io.as_ref()).await.unwrap(); + let expected = recover_by_decoding(io.as_ref()).await.unwrap(); + assert_same_recovery(&expected, &recovered); + assert_eq!(regions as usize, recovered.durable_entry_ids.len()); // The wide object took the header, the tail window and the footer; // the narrow one was read whole. @@ -1096,36 +1340,24 @@ mod tests { payload: Bytes::from(format!("e{entry_id}")), }) .collect::>(); - let encoded = encode_object( - Header { - object_seq, - writer_instance: [0; 16], - }, - &records, - ) - .unwrap(); - ObjectStoreIo::new(object_store.clone(), PREFIX) - .unwrap() - .put_if_absent(object_seq, encoded.bytes) - .await - .unwrap(); + put_records(object_store, object_seq, &records).await; } #[tokio::test] async fn test_store_resumes_sequence_and_durable_ids() { let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); - let (_, next, durable) = recover(&io).await.unwrap(); - assert_eq!(0, next); - assert!(durable.is_empty()); + let recovered = recover(&io).await.unwrap(); + assert_eq!(0, recovered.next_object_seq); + assert!(recovered.durable_entry_ids.is_empty()); put_object(&object_store, 2, region(1), &[1, id(5, 7)]).await; put_object(&object_store, 4, region(2), &[8]).await; - let (_, next, durable) = recover(&io).await.unwrap(); - assert_eq!(6, next); + let recovered = recover(&io).await.unwrap(); + assert_eq!(6, recovered.next_object_seq); assert_eq!( HashMap::from([(region(1), id(5, 7)), (region(2), 8)]), - durable + recovered.durable_entry_ids ); let store = open(object_store, &eager()).await; assert_eq!(id(5, 7), latest(&store, region(1))); @@ -1288,7 +1520,7 @@ mod tests { .await .unwrap() .into_iter() - .map(|(object, _)| object.object_seq) + .map(|fetched| fetched.object.object_seq) .collect::>() }) }; @@ -1356,7 +1588,8 @@ mod tests { assert!(parked.try_recv().is_err()); let store = open(object_store, &eager()).await; assert_eq!(id(15, 1), latest(&store, region(1))); - assert_eq!(16, store.catalog.read().unwrap().next_object_seq().unwrap()); + // Sequence 16 holds the object that started the epoch of the store. + assert_eq!(17, store.catalog.read().unwrap().next_object_seq().unwrap()); store.stop().await.unwrap(); } @@ -1676,6 +1909,465 @@ mod tests { store.stop().await.unwrap(); } + fn chain_header(object_seq: u64, epoch: u64, prev: Option<(u64, u64)>) -> Header { + Header { + object_seq, + epoch, + prev: prev.map(|(object_seq, epoch)| ChainLink { object_seq, epoch }), + } + } + + fn chain_of(headers: &[Header]) -> Vec { + select_chain( + &headers + .iter() + .map(|header| (header.object_seq, header)) + .collect(), + ) + } + + #[test] + fn test_store_selects_the_chain_of_the_latest_complete_object() { + assert!(chain_of(&[]).is_empty()); + // A linear chain is replayed whole. + assert_eq!( + vec![0, 1, 2], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(2, 1, Some((1, 1))), + ]) + ); + // Object 1 was reported as failed and object 2 links past it. + assert_eq!( + vec![0, 2], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(2, 1, Some((0, 1))), + ]) + ); + // Object 3 extends object 2, which never landed, so object 1 is the + // tip although object 3 has a higher sequence. + assert_eq!( + vec![0, 1], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 1, Some((0, 1))), + chain_header(3, 1, Some((2, 1))), + ]) + ); + // Object 1 links to an object of epoch 2 at sequence 0, where an + // object of epoch 1 is. + assert_eq!( + vec![0], + chain_of(&[chain_header(0, 1, None), chain_header(1, 2, Some((0, 2))),]) + ); + // Instance A left object 2 behind after instance B, a later epoch, + // acknowledged object 1: the later epoch ends the chain. + assert_eq!( + vec![0, 1], + chain_of(&[ + chain_header(0, 1, None), + chain_header(1, 2, Some((0, 1))), + chain_header(2, 1, Some((0, 1))), + ]) + ); + // A missing predecessor breaks the link, whether or not it lies below + // every present object, so an object that lands late below it cannot + // change which links hold. + assert!( + chain_of(&[ + chain_header(3, 1, Some((2, 1))), + chain_header(4, 1, Some((3, 1))), + ]) + .is_empty() + ); + assert!( + chain_of(&[ + chain_header(2, 1, Some((1, 1))), + chain_header(4, 2, Some((3, 1))), + chain_header(5, 2, Some((4, 2))), + ]) + .is_empty() + ); + assert_eq!( + vec![1], + chain_of(&[chain_header(1, 1, None), chain_header(3, 1, Some((2, 1))),]) + ); + // A predecessor at or above the object itself never holds. + assert!(chain_of(&[chain_header(5, 1, Some((5, 1)))]).is_empty()); + } + + #[tokio::test] + async fn test_store_recovery_indexes_only_the_chain() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = |entry_id: EntryId, data: &'static str| Record { + region_id: region(1), + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + put_fixture(&io, 0, &[record(1, "a")]).await; + let first = fixture_header(&io, 1).await; + // Object 1 was reported as failed; object 2 reused its row sequences + // and links past it. + put_object_with_header(&io, first.clone(), &[record(id(1, 1), "old")]).await; + let second = Header { + object_seq: 2, + ..first + }; + put_object_with_header(&io, second, &[record(id(2, 1), "new")]).await; + + for restart in 0..2 { + let store = open(object_store.clone(), &eager()).await; + assert_eq!( + expected_entries(region(1), &[(1, "a"), (id(2, 1), "new")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!(id(2, 1), latest(&store, region(1))); + store.stop().await.unwrap(); + // Every open starts the epoch one above the sequence of its start + // object, and the first restart's start object extends the chain + // the second restart replays. + let start = decode_header(&io.get(3 + restart).await.unwrap()).unwrap(); + assert_eq!(4 + restart, start.epoch); + assert_eq!(Some(2 + restart), start.prev.map(|prev| prev.object_seq)); + } + assert_eq!(vec![0, 1, 2, 3, 4], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_recovery_never_reuses_the_sequence_of_an_orphan() { + let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); + put_fixture(&io, 0, &[]).await; + // Object 2 extends object 1, which never landed. + let orphan = Header { + prev: Some(ChainLink { + object_seq: 1, + epoch: 1, + }), + ..fixture_header(&io, 2).await + }; + put_object_with_header(&io, orphan, &[]).await; + let recovered = recover(&io).await.unwrap(); + assert_eq!(Some(0), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(3, recovered.next_object_seq); + } + + #[tokio::test] + async fn test_store_recovery_rejects_objects_without_a_complete_chain() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // Object 1 extends object 0, which is gone. + put_object_with_header(&io, chain_header(1, 1, Some((0, 1))), &[]).await; + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::CorruptedWalObject { reason, .. } + if reason.contains("no object among 1 present objects completes a chain")), + "unexpected error: {error:?}" + ); + assert_eq!(vec![1], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_recovery_keeps_the_later_epoch_over_a_late_object() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = |entry_id: EntryId, data: &'static str| Record { + region_id: region(1), + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + put_object_with_header(&io, chain_header(0, 1, None), &[record(1, "a")]).await; + // Instance B acknowledged object 1 before the create instance A issued + // for object 2 landed. + put_object_with_header( + &io, + chain_header(1, 2, Some((0, 1))), + &[record(id(1, 1), "b")], + ) + .await; + put_object_with_header( + &io, + chain_header(2, 1, Some((0, 1))), + &[record(id(2, 1), "late")], + ) + .await; + + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region(1), &[(1, "a"), (id(1, 1), "b")]), + read_entries(&store, region(1), 0).await + ); + store.stop().await.unwrap(); + let start = decode_header(&io.get(3).await.unwrap()).unwrap(); + assert_eq!(4, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: 1, + epoch: 2, + }), + start.prev + ); + } + + #[tokio::test] + async fn test_store_open_starts_the_first_epoch_on_an_empty_prefix() { + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + store.stop().await.unwrap(); + let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); + let start = decode_header(&io.get(0).await.unwrap()).unwrap(); + assert_eq!(1, start.epoch); + assert_eq!(None, start.prev); + assert_eq!(vec![0], object_seqs(&io).await); + } + + #[tokio::test] + async fn test_store_start_object_moves_past_an_earlier_epoch_only() { + for (epoch, moves) in [(1, true), (2, false)] { + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&fixtures, 0, &[]).await; + // An object lands at the sequence of the start object after + // recovery listed the prefix. + let late = encode_object(chain_header(1, epoch, None), &[]) + .unwrap() + .bytes; + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((1, late))), + }); + let result = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await; + if moves { + result.unwrap().stop().await.unwrap(); + let start = decode_header(&fixtures.get(2).await.unwrap()).unwrap(); + assert_eq!(3, start.epoch); + assert_eq!(Some(0), start.prev.map(|prev| prev.object_seq)); + } else { + let error = result.unwrap_err(); + assert!( + matches!(&error, Error::WalObjectConflict { path, .. } if path == &fixtures.object_path(1)), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(&fixtures).await); + } + } + } + + #[tokio::test] + async fn test_store_open_fails_when_its_start_object_is_already_present() { + // On an empty prefix and after object 0 of epoch 1, another open that + // recovered the same objects wrote the very start object this open + // writes. + for present in [false, true] { + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let (start_seq, epoch, tip) = if present { + put_fixture(&fixtures, 0, &[]).await; + (1, 2, Some((0, 1))) + } else { + (0, 1, None) + }; + let same = encode_object(chain_header(start_seq, epoch, tip), &[]) + .unwrap() + .bytes; + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((start_seq, same))), + }); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::UnconfirmedWalEpochStart { epoch: actual, path, .. } + if *actual == epoch && path == &fixtures.object_path(start_seq)), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + + // The next open counts that object and starts a later epoch. + open(object_store, &eager()).await.stop().await.unwrap(); + let start = decode_header(&fixtures.get(start_seq + 1).await.unwrap()).unwrap(); + assert_eq!(epoch + 1, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: start_seq, + epoch, + }), + start.prev + ); + } + } + + #[tokio::test] + async fn test_store_opens_racing_across_a_late_object_take_distinct_epochs() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&io, 0, &[]).await; + // Open A recovers object 0 and pauses before its start object. + let recovered = recover(&io).await.unwrap(); + // A late object of epoch 1 lands at sequence 2, across the gap at + // sequence 1, and open B, which lists it, completes. + put_object_with_header(&io, chain_header(2, 1, Some((0, 1))), &[]).await; + open(object_store, &eager()).await.stop().await.unwrap(); + let b = decode_header(&io.get(3).await.unwrap()).unwrap(); + assert_eq!(4, b.epoch); + + // Open A resumes at the sequence it recovered: the epoch follows from + // the sequence it claims, so it differs from the epoch of open B. + let a = start_epoch(&io, recovered.next_object_seq, recovered.tip) + .await + .unwrap(); + assert_eq!( + ChainLink { + object_seq: 1, + epoch: 2, + }, + a + ); + } + + #[tokio::test] + async fn test_store_open_rejects_an_epoch_above_the_next_sequence() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // No instance that starts at or below sequence 0 writes epoch 5. + put_object_with_header(&io, chain_header(0, 5, None), &[]).await; + let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::CorruptedWalObject { reason, .. } + if reason.contains("carries epoch 5, above the next sequence 1")), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(&io).await); + } + + async fn object_seqs(io: &ObjectStoreIo) -> Vec { + io.list() + .await + .unwrap() + .into_iter() + .map(|object| object.object_seq) + .collect() + } + + #[tokio::test] + async fn test_store_open_fails_when_its_start_object_lands_with_an_unknown_outcome() { + let object_store = memory_store(); + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + put_fixture(&fixtures, 0, &[]).await; + // The start object is stored, but its create reports an error. + let io = Arc::new(LostResponseIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + lose_next: AtomicBool::new(true), + }); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + // The open does not move to another sequence. + assert_eq!(vec![0, 1], object_seqs(&fixtures).await); + let stored = decode_header(&fixtures.get(1).await.unwrap()).unwrap(); + assert_eq!(2, stored.epoch); + + // The next open counts the stored start object and claims a later + // epoch. + open(object_store, &eager()).await.stop().await.unwrap(); + let start = decode_header(&fixtures.get(2).await.unwrap()).unwrap(); + assert_eq!(3, start.epoch); + assert_eq!( + Some(ChainLink { + object_seq: 1, + epoch: 2, + }), + start.prev + ); + } + + /// Object access whose next create stores the object but reports a + /// transient failure, as a create whose response is lost. + struct LostResponseIo { + inner: ObjectStoreIo, + lose_next: AtomicBool, + } + + #[async_trait::async_trait] + impl WalObjectIo for LostResponseIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + let result = self.inner.put_if_absent(object_seq, content).await?; + if !self.lose_next.swap(false, Ordering::SeqCst) { + return Ok(result); + } + Err( + object_store::Error::new(object_store::ErrorKind::Unexpected, "lost response") + .set_temporary(), + ) + .context(WalObjectStoreSnafu { + operation: "write", + path: self.object_path(object_seq), + }) + } + + async fn get(&self, object_seq: u64) -> Result { + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.inner.get_range(object_seq, offset, len).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } + + /// Object access whose first create lets another object land at a given + /// sequence first, as a create issued before recovery would. + struct RacingIo { + inner: ObjectStoreIo, + late: Mutex>, + } + + #[async_trait::async_trait] + impl WalObjectIo for RacingIo { + async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + let late = self.late.lock().unwrap().take(); + if let Some((late_seq, late)) = late { + self.inner.put_if_absent(late_seq, late).await?; + } + self.inner.put_if_absent(object_seq, content).await + } + + async fn get(&self, object_seq: u64) -> Result { + self.inner.get(object_seq).await + } + + async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { + self.inner.get_range(object_seq, offset, len).await + } + + async fn list(&self) -> Result> { + self.inner.list().await + } + + fn object_path(&self, object_seq: u64) -> String { + self.inner.object_path(object_seq) + } + } + struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, From 357c9358043bb89e21f52845b97b9d3f1a8a1970 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 24 Sep 2026 08:45:40 +0000 Subject: [PATCH 036/133] fix: make vector aggregates work with GROUP BY and partial aggregation (#9338) * fix: make vector aggregates work with GROUP BY and partial aggregation Signed-off-by: Dennis Zhuang * test: make partitioned vec_avg case distinguish weighted averages Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/aggrs/vector/avg.rs | 6 +- .../function/src/aggrs/vector/product.rs | 123 ++++++++++++---- src/common/function/src/aggrs/vector/sum.rs | 133 ++++++++++++------ .../common/function/vector/vector.result | 86 +++++++++++ .../common/function/vector/vector.sql | 45 ++++++ 5 files changed, 321 insertions(+), 72 deletions(-) diff --git a/src/common/function/src/aggrs/vector/avg.rs b/src/common/function/src/aggrs/vector/avg.rs index 822a7c16efe..82e6a28059f 100644 --- a/src/common/function/src/aggrs/vector/avg.rs +++ b/src/common/function/src/aggrs/vector/avg.rs @@ -62,14 +62,14 @@ impl VectorAvg { } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_AVG` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_AVG`" diff --git a/src/common/function/src/aggrs/vector/product.rs b/src/common/function/src/aggrs/vector/product.rs index 11cfa4cb1a8..8178b74f369 100644 --- a/src/common/function/src/aggrs/vector/product.rs +++ b/src/common/function/src/aggrs/vector/product.rs @@ -28,6 +28,9 @@ use crate::scalars::vector::impl_conv::{ }; /// Aggregates by multiplying elements across the same dimension, returns a vector. +/// +/// The result is NULL if any input vector is NULL, so the partial state carries a +/// `has_null` flag: a NULL `product` alone can't tell a NULL input from an empty partition. #[derive(Debug, Default)] pub struct VectorProduct { product: Option>, @@ -49,20 +52,23 @@ impl VectorProduct { signature, DataType::Binary, Arc::new(Self::accumulator), - vec![Arc::new(Field::new("x", DataType::Binary, true))], + vec![ + Arc::new(Field::new("product", DataType::Binary, true)), + Arc::new(Field::new("has_null", DataType::Boolean, true)), + ], ); AggregateUDF::from(udaf) } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_PRODUCT` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_PRODUCT`" @@ -72,13 +78,33 @@ impl VectorProduct { Ok(Box::new(VectorProduct::default())) } - fn inner(&mut self, len: usize) -> &mut OVector { - self.product.get_or_insert_with(|| { - OVector::from_iterator_generic(Dyn(len), Const::<1>, (0..len).map(|_| 1.0)) - }) + fn mul(&mut self, vector: &[f32]) { + let vector = DVectorView::from_slice(vector, vector.len()); + let product = self.product.get_or_insert_with(|| { + OVector::from_iterator_generic( + Dyn(vector.len()), + Const::<1>, + (0..vector.len()).map(|_| 1.0), + ) + }); + *product = product.component_mul(&vector); } - fn update(&mut self, values: &[ArrayRef], is_update: bool) -> Result<()> { + fn set_null(&mut self) { + self.has_null = true; + self.product = None; + } +} + +impl Accumulator for VectorProduct { + fn state(&mut self) -> Result> { + Ok(vec![ + self.evaluate()?, + ScalarValue::Boolean(Some(self.has_null)), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { if values.is_empty() || self.has_null { return Ok(()); }; @@ -112,33 +138,34 @@ impl VectorProduct { } }; if vectors.len() != values[0].len() { - if is_update { - self.has_null = true; - self.product = None; - } + self.set_null(); return Ok(()); } - vectors.iter().for_each(|v| { - let v = DVectorView::from_slice(v, v.len()); - let inner = self.inner(v.len()); - *inner = inner.component_mul(&v); - }); + vectors.iter().for_each(|v| self.mul(v)); Ok(()) } -} - -impl Accumulator for VectorProduct { - fn state(&mut self) -> Result> { - self.evaluate().map(|v| vec![v]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.update(values, true) - } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.update(states, false) + let [products, has_nulls] = states else { + return Err(datafusion_common::DataFusionError::Internal(format!( + "expect 2 states for `VEC_PRODUCT`, actual {}", + states.len() + ))); + }; + if self.has_null { + return Ok(()); + } + if has_nulls.as_boolean().true_count() > 0 { + self.set_null(); + return Ok(()); + } + + // A NULL product without `has_null` comes from a partition without input rows. + for b in products.as_binary::().iter().flatten() { + self.mul(&binlit_as_veclit(b)?); + } + Ok(()) } fn evaluate(&mut self) -> Result { @@ -225,4 +252,42 @@ mod tests { vec_product.evaluate().unwrap() ); } + + #[test] + fn test_merge_batch() { + let partial = |v: Option<&str>| { + let mut acc = VectorProduct::default(); + let v: ArrayRef = Arc::new(StringArray::from(vec![v])); + acc.update_batch(&[v]).unwrap(); + acc.state().unwrap() + }; + let states = |states: Vec>| -> Vec { + (0..2) + .map(|i| ScalarValue::iter_to_array(states.iter().map(|s| s[i].clone())).unwrap()) + .collect() + }; + + // An empty partition in the middle of the batch must not stop the merge. + let mut merged = VectorProduct::default(); + merged + .merge_batch(&states(vec![ + partial(Some("[1.0,2.0]")), + VectorProduct::default().state().unwrap(), + partial(Some("[3.0,4.0]")), + ])) + .unwrap(); + assert_eq!( + ScalarValue::Binary(Some(veclit_to_binlit(&[3.0, 8.0]))), + merged.evaluate().unwrap() + ); + + // A NULL input in any partition makes the result NULL. + merged + .merge_batch(&states(vec![partial(Some("[1.0,2.0]")), partial(None)])) + .unwrap(); + merged + .merge_batch(&states(vec![partial(Some("[3.0,4.0]"))])) + .unwrap(); + assert_eq!(ScalarValue::Binary(None), merged.evaluate().unwrap()); + } } diff --git a/src/common/function/src/aggrs/vector/sum.rs b/src/common/function/src/aggrs/vector/sum.rs index 0ee8f0df7bf..b1b68906ded 100644 --- a/src/common/function/src/aggrs/vector/sum.rs +++ b/src/common/function/src/aggrs/vector/sum.rs @@ -28,6 +28,9 @@ use crate::scalars::vector::impl_conv::{ }; /// The accumulator for the `vec_sum` aggregate function. +/// +/// The result is NULL if any input vector is NULL, so the partial state carries a +/// `has_null` flag: a NULL `sum` alone can't tell a NULL input from an empty partition. #[derive(Debug, Default)] pub struct VectorSum { sum: Option>, @@ -49,20 +52,23 @@ impl VectorSum { signature, DataType::Binary, Arc::new(Self::accumulator), - vec![Arc::new(Field::new("x", DataType::Binary, true))], + vec![ + Arc::new(Field::new("sum", DataType::Binary, true)), + Arc::new(Field::new("has_null", DataType::Boolean, true)), + ], ); AggregateUDF::from(udaf) } fn accumulator(args: AccumulatorArgs) -> Result> { - if args.schema.fields().len() != 1 { + if args.exprs.len() != 1 { return Err(datafusion_common::DataFusionError::Internal(format!( "expect creating `VEC_SUM` with only one input field, actual {}", - args.schema.fields().len() + args.exprs.len() ))); } - let t = args.schema.field(0).data_type(); + let t = args.expr_fields[0].data_type(); if !matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary) { return Err(datafusion_common::DataFusionError::Internal(format!( "unexpected input datatype {t} when creating `VEC_SUM`" @@ -72,12 +78,28 @@ impl VectorSum { Ok(Box::new(VectorSum::default())) } - fn inner(&mut self, len: usize) -> &mut OVector { - self.sum - .get_or_insert_with(|| OVector::zeros_generic(Dyn(len), Const::<1>)) + fn add(&mut self, vector: &[f32]) { + let vector = DVectorView::from_slice(vector, vector.len()); + *self + .sum + .get_or_insert_with(|| OVector::zeros_generic(Dyn(vector.len()), Const::<1>)) += vector; } - fn update(&mut self, values: &[ArrayRef], is_update: bool) -> Result<()> { + fn set_null(&mut self) { + self.has_null = true; + self.sum = None; + } +} + +impl Accumulator for VectorSum { + fn state(&mut self) -> Result> { + Ok(vec![ + self.evaluate()?, + ScalarValue::Boolean(Some(self.has_null)), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { if values.is_empty() || self.has_null { return Ok(()); }; @@ -87,45 +109,30 @@ impl VectorSum { let arr: &StringArray = values[0].as_string(); for s in arr.iter() { let Some(s) = s else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = parse_veclit_from_strlit(s)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&parse_veclit_from_strlit(s)?); } } DataType::LargeUtf8 => { let arr: &LargeStringArray = values[0].as_string(); for s in arr.iter() { let Some(s) = s else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = parse_veclit_from_strlit(s)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&parse_veclit_from_strlit(s)?); } } DataType::Binary => { let arr: &BinaryArray = values[0].as_binary(); for b in arr.iter() { let Some(b) = b else { - if is_update { - self.has_null = true; - self.sum = None; - } + self.set_null(); return Ok(()); }; - let values = binlit_as_veclit(b)?; - let vec_column = DVectorView::from_slice(&values, values.len()); - *self.inner(vec_column.len()) += vec_column; + self.add(&binlit_as_veclit(b)?); } } _ => { @@ -137,19 +144,27 @@ impl VectorSum { } Ok(()) } -} - -impl Accumulator for VectorSum { - fn state(&mut self) -> Result> { - self.evaluate().map(|v| vec![v]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.update(values, true) - } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.update(states, false) + let [sums, has_nulls] = states else { + return Err(datafusion_common::DataFusionError::Internal(format!( + "expect 2 states for `VEC_SUM`, actual {}", + states.len() + ))); + }; + if self.has_null { + return Ok(()); + } + if has_nulls.as_boolean().true_count() > 0 { + self.set_null(); + return Ok(()); + } + + // A NULL sum without `has_null` comes from a partition without input rows. + for b in sums.as_binary::().iter().flatten() { + self.add(&binlit_as_veclit(b)?); + } + Ok(()) } fn evaluate(&mut self) -> Result { @@ -236,4 +251,42 @@ mod tests { vec_sum.evaluate().unwrap() ); } + + #[test] + fn test_merge_batch() { + let partial = |v: Option<&str>| { + let mut acc = VectorSum::default(); + let v: ArrayRef = Arc::new(StringArray::from(vec![v])); + acc.update_batch(&[v]).unwrap(); + acc.state().unwrap() + }; + let states = |states: Vec>| -> Vec { + (0..2) + .map(|i| ScalarValue::iter_to_array(states.iter().map(|s| s[i].clone())).unwrap()) + .collect() + }; + + // An empty partition in the middle of the batch must not stop the merge. + let mut merged = VectorSum::default(); + merged + .merge_batch(&states(vec![ + partial(Some("[1.0,2.0]")), + VectorSum::default().state().unwrap(), + partial(Some("[3.0,4.0]")), + ])) + .unwrap(); + assert_eq!( + ScalarValue::Binary(Some(veclit_to_binlit(&[4.0, 6.0]))), + merged.evaluate().unwrap() + ); + + // A NULL input in any partition makes the result NULL. + merged + .merge_batch(&states(vec![partial(Some("[1.0,2.0]")), partial(None)])) + .unwrap(); + merged + .merge_batch(&states(vec![partial(Some("[3.0,4.0]"))])) + .unwrap(); + assert_eq!(ScalarValue::Binary(None), merged.evaluate().unwrap()); + } } diff --git a/tests/cases/standalone/common/function/vector/vector.result b/tests/cases/standalone/common/function/vector/vector.result index c546f9ad25e..06f7c684fa1 100644 --- a/tests/cases/standalone/common/function/vector/vector.result +++ b/tests/cases/standalone/common/function/vector/vector.result @@ -448,3 +448,89 @@ FROM ( | [4,5,6,10,-8] | +---------------------------------------------------+ +SELECT h, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM ( + SELECT 'a' AS h, '[1.0, 1.0]' AS v + UNION ALL + SELECT 'a' AS h, '[2.0, 2.0]' AS v + UNION ALL + SELECT 'b' AS h, '[3.0, 3.0]' AS v +) GROUP BY h ORDER BY h; + ++---+---------------------------+---------------------------+-------------------------------+ +| h | vec_to_string(vec_sum(v)) | vec_to_string(vec_avg(v)) | vec_to_string(vec_product(v)) | ++---+---------------------------+---------------------------+-------------------------------+ +| a | [3,3] | [1.5,1.5] | [2,2] | +| b | [3,3] | [3,3] | [3,3] | ++---+---------------------------+---------------------------+-------------------------------+ + +-- On partitioned tables the aggregates are split into partial state and merge. +-- Rows only land in two of the three partitions, the third one has an empty state. +-- The two non-empty partitions differ in row count and mean, so vec_avg must weight by count. +CREATE TABLE vector_aggr_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + g STRING, + v VECTOR(2), + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +Affected Rows: 0 + +INSERT INTO vector_aggr_partitioned VALUES + (1000, 1, 'a', '[1.0, 1.0]'), + (2000, 11, 'a', '[8.0, 8.0]'), + (3000, 2, 'a', '[3.0, 3.0]'), + (3500, 3, 'a', '[5.0, 5.0]'), + (4000, 12, 'b', '[4.0, 4.0]'); + +Affected Rows: 5 + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| [21,21] | [4.2,4.2] | [480,480] | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| g | vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| a | [17,17] | [4.25,4.25] | [120,120] | +| b | [4,4] | [4,4] | [4,4] | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +-- A NULL vector makes vec_sum and vec_product NULL, vec_avg skips it. +INSERT INTO vector_aggr_partitioned VALUES (5000, 21, 'b', NULL); + +Affected Rows: 1 + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| | [4.2,4.2] | | ++---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| g | vec_to_string(vec_sum(vector_aggr_partitioned.v)) | vec_to_string(vec_avg(vector_aggr_partitioned.v)) | vec_to_string(vec_product(vector_aggr_partitioned.v)) | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ +| a | [17,17] | [4.25,4.25] | [120,120] | +| b | | [4,4] | | ++---+---------------------------------------------------+---------------------------------------------------+-------------------------------------------------------+ + +DROP TABLE vector_aggr_partitioned; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/vector/vector.sql b/tests/cases/standalone/common/function/vector/vector.sql index 9bbf1583f53..c53438c854d 100644 --- a/tests/cases/standalone/common/function/vector/vector.sql +++ b/tests/cases/standalone/common/function/vector/vector.sql @@ -150,3 +150,48 @@ FROM ( UNION ALL SELECT '[4.0, 5.0, 6.0, 10, -8, 100]' AS v ) ORDER BY v; + +SELECT h, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM ( + SELECT 'a' AS h, '[1.0, 1.0]' AS v + UNION ALL + SELECT 'a' AS h, '[2.0, 2.0]' AS v + UNION ALL + SELECT 'b' AS h, '[3.0, 3.0]' AS v +) GROUP BY h ORDER BY h; + +-- On partitioned tables the aggregates are split into partial state and merge. +-- Rows only land in two of the three partitions, the third one has an empty state. +-- The two non-empty partitions differ in row count and mean, so vec_avg must weight by count. +CREATE TABLE vector_aggr_partitioned ( + ts TIMESTAMP TIME INDEX, + k INT, + g STRING, + v VECTOR(2), + PRIMARY KEY(k) +) +PARTITION ON COLUMNS (k) (k < 10, k >= 10 AND k < 20, k >= 20); + +INSERT INTO vector_aggr_partitioned VALUES + (1000, 1, 'a', '[1.0, 1.0]'), + (2000, 11, 'a', '[8.0, 8.0]'), + (3000, 2, 'a', '[3.0, 3.0]'), + (3500, 3, 'a', '[5.0, 5.0]'), + (4000, 12, 'b', '[4.0, 4.0]'); + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + +-- A NULL vector makes vec_sum and vec_product NULL, vec_avg skips it. +INSERT INTO vector_aggr_partitioned VALUES (5000, 21, 'b', NULL); + +SELECT vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned; + +SELECT g, vec_to_string(vec_sum(v)), vec_to_string(vec_avg(v)), vec_to_string(vec_product(v)) +FROM vector_aggr_partitioned GROUP BY g ORDER BY g; + +DROP TABLE vector_aggr_partitioned; From b5c0e246254735e01d9507e71970f5a1e8afcfe4 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:45:45 +0000 Subject: [PATCH 037/133] feat(otlp): add histogram rejection metrics and compact zero buckets (#9321) * feat(otlp): add histogram rejection metrics and compact zero buckets Signed-off-by: shuiyisong * fix: test isolation Signed-off-by: shuiyisong * chore: fix minor naming Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- docs/rfcs/2026-05-28-table-semantic-layer.md | 16 +- docs/rfcs/2026-08-04-native-histograms.md | 56 ++++- src/servers/src/metrics.rs | 7 + src/servers/src/otlp/metrics.rs | 239 ++++++++++++++++++- 4 files changed, 299 insertions(+), 19 deletions(-) diff --git a/docs/rfcs/2026-05-28-table-semantic-layer.md b/docs/rfcs/2026-05-28-table-semantic-layer.md index 24ed96dc81b..cb5f6e84dc4 100644 --- a/docs/rfcs/2026-05-28-table-semantic-layer.md +++ b/docs/rfcs/2026-05-28-table-semantic-layer.md @@ -93,13 +93,25 @@ Two design decisions worth pinning down up front, because they constrain everyth - **Conflict.** Some table-level keys (`trace.conventions` lifted from `schema_url`, `metric.temporality`, ...) cannot represent the truth when a long-lived table sees rows from multiple sources. v1 records `mixed` or `unknown` rather than a fictitious single value. Downstream consumers must treat any single-valued semantic key as best-effort, not strong evidence. - **Update.** Semantic options are stamped at table creation. v1 does not specify an update path; promoting `metadata_quality` from `inferred` to `declared`, refreshing `resource.attributes_preserved`, or revising `trace.conventions` on later writes is deferred. If real usage shows update is needed, it lands as a separate RFC. +OTLP delta sums and explicit histograms store raw interval values without +accumulating across timestamps. Delta-to-cumulative conversion and its +per-series ingestion state are intentionally out of scope. Explicit histogram +bucket counts are prefix-summed only within each point to produce the classic +histogram representation. Delta exponential histograms are currently rejected; +[planned raw-delta support](2026-08-04-native-histograms.md#planned-raw-delta-exponential-histograms) +will store their interval values in the canonical native-histogram Struct using +the same temporality tag and a gauge reset hint. The tag will select histogram +summation for `increase()` and summation divided by range duration for `rate()`, +without cumulative reset correction or extrapolation. + OTLP delta sums and explicit histograms additionally store the query-visible String tag `otlp_aggregation_temporality="delta"` on each generated row. Its name is fixed and does not follow `default_column_prefix`. The tag is part of series identity and is authoritative for per-series float `rate()` and `increase()` behavior; the table option is never used as a row-level -discriminator. Native histograms retain their native algorithms. Prometheus -metadata reports `unknown` for counter, histogram, and up/down-counter tables +discriminator. Native histograms currently retain their native algorithms; +the planned delta extension will change only tagged delta histogram series. +Prometheus metadata reports `unknown` for counter, histogram, and up/down-counter tables whose catalog temporality is `delta` or `mixed`. A same-request conflict can create `mixed`; a later write does not update an existing table option, so the catalog value can also be stale while the row tag remains authoritative. diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index 6260e79f3d5..3aef4eb7ffb 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -1,6 +1,6 @@ --- Feature Name: Native Histogram Support and Compatibility Decisions -Tracking Issue: TBD +Tracking Issue: https://github.com/GreptimeTeam/greptimedb/issues/8887 Date: 2026-08-04 Author: codex --- @@ -34,6 +34,9 @@ continues to return scalar samples only. - Returning native histograms through Prometheus Remote Read. - Persisting Remote Write metric metadata. - Persisting exemplars. +- Converting OTLP delta values to cumulative values during ingestion. Supported + delta metrics are stored as interval values without per-series accumulation + state; this also applies to the planned exponential-histogram support. - Providing an OTLP-specific histogram query surface or reconstructing and re-exporting the original OTLP point after persistence. - Removing mixed float/histogram handling from PromQL expressions. @@ -123,10 +126,20 @@ explicit if sampled responses are implemented. OTLP/HTTP exponential histograms are accepted by default. OTel Arrow exponential histograms are rejected because the current Arrow wire format omits `zero_threshold`; accepting them would -silently change the distribution. Cumulative temporality is required; delta and -unspecified exponential histograms are rejected before their points are +silently change the distribution. Currently, cumulative temporality is required; +delta and unspecified exponential histograms are rejected before their points are converted. Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and -`_count` representation, including their existing delta behavior. +`_count` representation. OTLP delta sums and explicit histograms are stored as +raw interval values with `otlp_aggregation_temporality="delta"`. Explicit +histogram bucket counts are prefix-summed within each point for the classic +histogram representation, but values are never accumulated across timestamps. +For these tagged float series, PromQL `increase()` sums the interval values +in the query range, and `rate()` divides that sum by the range duration. +Delta-to-cumulative conversion is intentionally out of scope, not deferred work. +Raw-delta exponential-histogram support is planned below; delta exponential +points remain rejected until it is implemented. See +[Table Semantic Layer](2026-05-28-table-semantic-layer.md#conflict-and-update-semantics) +for the per-series temporality contract. OTLP scales `-4` through `8` map directly to Prometheus schemas. Higher scales are downscaled to schema `8`: dense counts that collide are merged before the @@ -156,8 +169,36 @@ for mixed batches and `INVALID_ARGUMENT` when all points are rejected. Rejection details are bounded, and metric metadata is emitted only for a metric that produced an accepted row. -Minimum, maximum, and exemplars are not persisted. Delta accumulation, -zero-run span compaction, and a dedicated rejection metric remain deferred. +Zero-count buckets are omitted after downscaling, splitting nonzero runs into +sparse spans without changing integer counts. Rejected data points increment +`greptime_servers_otlp_exponential_histogram_rejected_data_points_total`, labeled +by `reason`: `delta_temporality`, `unspecified_temporality`, or +`invalid_data_point`. Each rejected point is counted once, including when +resource descriptors are enabled. Minimum, maximum, and exemplars are not persisted. + +### Planned raw-delta exponential histograms + +Accept OTLP delta exponential histograms and store each point's interval counts +and sum in the existing canonical native-histogram Struct, without accumulating +across timestamps. Reuse the integer-count encoding, scale conversion, +zero-threshold preservation, stale-marker handling, and partial-success path. +OTel Arrow support remains blocked independently by its missing `zero_threshold`. + +Store `otlp_aggregation_temporality="delta"` as part of series identity, as for +delta sums and explicit histograms. Use the gauge reset hint for delta histogram +samples so decreasing interval counts are not interpreted as counter resets. +Prometheus uses the same gauge hint for +[raw-delta exponential-histogram ingestion](https://github.com/prometheus/prometheus/blob/9f159f6df198280b757163ec555ffde0c8bfe8e0/storage/remote/otlptranslator/prometheusremotewrite/histograms.go#L125). +The temporality tag, rather than the gauge hint alone or table metadata, +determines raw-delta query behavior. + +Extend the per-series `increase()` and `rate()` selection to native histogram +inputs: sum interval histograms over the query range, then divide by the range +duration for `rate()`. Do not apply cumulative reset correction, synthetic-zero +insertion, or counter extrapolation to these delta samples. Preserve cumulative +histogram behavior and mixed float/histogram range warnings. Acceptance coverage +must include multiple intervals whose counts decrease, exact integer ingestion, +stale points, partial rejection, and pure-native and mixed-range query paths. # PromQL Compatibility Decisions @@ -273,7 +314,8 @@ native-histogram data remains readable. - Native-histogram Remote Read, including exact integer round-trips and streamed chunks with start timestamps. -- OTLP exponential-histogram delta accumulation. +- Raw-delta OTLP exponential-histogram storage and temporality-aware native + histogram queries, as described above; no delta-to-cumulative conversion. - Persistent Remote Write metadata and accurate help/unit updates. - Native-histogram exemplars and exemplar query APIs. - Start-timestamp overlap annotations. diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 0d70540be31..888bdc42cdf 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -130,6 +130,13 @@ lazy_static! { &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); + /// OTLP exponential-histogram data points rejected during conversion, by reason. + pub static ref OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS: IntCounterVec = register_int_counter_vec!( + "greptime_servers_otlp_exponential_histogram_rejected_data_points_total", + "Number of OTLP exponential histogram data points rejected during conversion", + &["reason"] + ) + .unwrap(); pub static ref PENDING_BATCHES: IntGauge = register_int_gauge!( "greptime_prom_store_pending_batches", "Number of pending batches waiting to be flushed" diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 7862af28db3..e518d732479 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -41,6 +41,7 @@ use table::requests::{ }; use crate::error::{self, Result}; +use crate::metrics::OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; @@ -634,6 +635,9 @@ fn encode_exponential_histogram( reject_data_points(outcome, histogram.data_points.len(), || { rejection.message(name) })?; + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&[rejection.reason_label()]) + .inc_by(histogram.data_points.len() as u64); return Ok(false); } @@ -651,6 +655,9 @@ fn encode_exponential_histogram( reject_data_points(outcome, 1, || { format!("metric `{name}` data point {index}: {reason}") })?; + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS + .with_label_values(&["invalid_data_point"]) + .inc(); continue; } }; @@ -689,6 +696,13 @@ pub(crate) enum ExponentialHistogramRejection { } impl ExponentialHistogramRejection { + fn reason_label(&self) -> &'static str { + match self { + Self::DeltaTemporality => "delta_temporality", + Self::UnspecifiedTemporality => "unspecified_temporality", + } + } + fn message(&self, name: &str) -> String { match self { Self::DeltaTemporality => format!( @@ -704,6 +718,8 @@ impl ExponentialHistogramRejection { /// Whole-metric acceptance, decided once for the encoder and for the resource /// descriptor, which must not describe a resource whose data was rejected. /// Individual points can still fail [`exponential_histogram_value`]. +/// Raw-delta storage currently supports sums and explicit histograms; exponential +/// histograms still require cumulative input. Delta-to-cumulative conversion is out of scope. pub(crate) fn exponential_histogram_gate( histogram: &ExponentialHistogram, ) -> std::result::Result<(), ExponentialHistogramRejection> { @@ -870,25 +886,44 @@ fn convert_bucket_range( merged.push((target_index, count)); } - let length = u32::try_from(merged.len()) - .map_err(|_| format!("{name} bucket span length exceeds u32"))?; - let span = BucketSpan { - offset: merged[0].0, - length, - }; + let mut spans = Vec::::new(); let mut deltas = Vec::with_capacity(merged.len()); + let mut previous_index = None::; let mut previous = 0i64; - for (_, count) in merged { + for (index, count) in merged { + if count == 0 { + continue; + } + match (spans.last_mut(), previous_index) { + (Some(span), Some(prev_index)) if prev_index.checked_add(1) == Some(index) => { + span.length = span + .length + .checked_add(1) + .ok_or_else(|| format!("{name} bucket span length exceeds u32"))?; + } + (_, prev_index) => { + let offset = match prev_index { + Some(prev_index) => index + .checked_sub(prev_index) + .and_then(|gap| gap.checked_sub(1)) + .ok_or_else(|| format!("{name} bucket span offset overflows i32"))?, + None => index, + }; + spans.push(BucketSpan { offset, length: 1 }); + } + } let count = i64::try_from(count) .map_err(|_| format!("{name} bucket count {count} overflows i64"))?; let delta = count .checked_sub(previous) .ok_or_else(|| format!("{name} bucket delta overflows i64"))?; deltas.push(delta); + previous_index = Some(index); + // Bucket deltas continue across span gaps; omitted zero buckets do not reset them. previous = count; } - Ok((vec![span], deltas, total)) + Ok((spans, deltas, total)) } fn downscale_bucket_index(index: i32, downscale_shift: u32) -> std::result::Result { @@ -1109,8 +1144,8 @@ fn encode_gauge( Ok(()) } -/// encode this sum metric -/// +/// Encodes sums, preserving delta points as interval values with a temporality tag. +/// Ingestion is stateless: values are never accumulated across timestamps. fn encode_sum( table_writer: &mut MultiTableData, name: &str, @@ -1170,6 +1205,8 @@ const HISTOGRAM_LE_COLUMN: &str = "le"; /// /// By its Prometheus compatibility, we hope to be able to use prometheus /// quantile functions on this table. +/// Delta points retain their temporality tag and interval values. Bucket counts +/// are prefix-summed within each point, never accumulated across timestamps. fn encode_histogram( table_writer: &mut MultiTableData, name: &str, @@ -2286,6 +2323,105 @@ mod tests { assert_eq!(total, 21); } + #[test] + fn test_convert_bucket_range_compacts_zero_runs() { + let large = (1u64 << 53) + 1; + for (offset, counts, shift, expected_spans, expected_deltas, expected_total) in [ + ( + -3, + vec![0, 3, 0, 0, 7, 9, 0], + 0, + vec![(-1, 1), (2, 2)], + vec![3, 4, 2], + 19, + ), + ( + -4, + vec![0, 0, 2, 3, 0, 0, 0, 0, 7, 0], + 1, + vec![(0, 1), (2, 1)], + vec![5, 2], + 12, + ), + (0, vec![0, 0, 0], 0, vec![], vec![], 0), + (0, vec![], 0, vec![], vec![], 0), + (i32::MIN, vec![0, 1], 0, vec![(i32::MIN + 2, 1)], vec![1], 1), + (i32::MAX - 1, vec![1], 0, vec![(i32::MAX, 1)], vec![1], 1), + ( + 0, + vec![large, 0, large + 2], + 0, + vec![(1, 1), (1, 1)], + vec![large as i64, 2], + large * 2 + 2, + ), + ] { + let buckets = exponential_buckets(offset, counts); + let (spans, deltas, total) = + convert_bucket_range("positive", Some(&buckets), shift).unwrap(); + assert_eq!( + spans + .iter() + .map(|span| (span.offset, span.length)) + .collect::>(), + expected_spans, + "{buckets:?}, shift={shift}" + ); + assert_eq!(deltas, expected_deltas, "{buckets:?}, shift={shift}"); + assert_eq!(total, expected_total); + } + + // Compaction must not hide invalid source indexes, even for empty buckets. + let buckets = exponential_buckets(i32::MAX, vec![0]); + assert!( + convert_bucket_range("positive", Some(&buckets), 0) + .unwrap_err() + .contains("shifted bucket index overflows") + ); + } + + #[test] + fn test_exponential_histogram_compaction_preserves_integer_counts() { + use common_query::native_histogram::{ + COUNT_I64_FIELD, NEGATIVE_BUCKETS_I64_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, + NEGATIVE_SPAN_OFFSETS_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_LENGTHS_FIELD, + POSITIVE_SPAN_OFFSETS_FIELD, + }; + + let large = (1u64 << 53) + 1; + let buckets = exponential_buckets(-2, vec![0, large, 0, 3, 0]); + let point = ExponentialHistogramDataPoint { + count: 2 * (large + 3), + positive: Some(buckets.clone()), + negative: Some(buckets), + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(point.count as i64)) + ); + for (offsets, lengths, counts) in [ + ( + POSITIVE_SPAN_OFFSETS_FIELD, + POSITIVE_SPAN_LENGTHS_FIELD, + POSITIVE_BUCKETS_I64_FIELD, + ), + ( + NEGATIVE_SPAN_OFFSETS_FIELD, + NEGATIVE_SPAN_LENGTHS_FIELD, + NEGATIVE_BUCKETS_I64_FIELD, + ), + ] { + assert_eq!(i32_list(native_field(&value, offsets)), vec![0, 1]); + assert_eq!(i32_list(native_field(&value, lengths)), vec![1, 1]); + assert_eq!( + i64_list(native_field(&value, counts)), + vec![large as i64, 3] + ); + } + } + #[test] fn test_exponential_histogram_value_uses_integer_family() { use common_query::native_histogram::{ @@ -2489,6 +2625,89 @@ mod tests { } } + #[test] + fn test_exponential_histogram_rejection_metrics() { + // Other conversion tests update the same process-global counters. + const ISOLATED_ENV: &str = "GREPTIME_TEST_OTLP_REJECTION_METRICS_ISOLATED"; + if std::env::var_os(ISOLATED_ENV).is_none() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "otlp::metrics::tests::test_exponential_histogram_rejection_metrics", + ]) + .env(ISOLATED_ENV, "1") + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "isolated rejection metrics test failed\nstdout:\n{stdout}\nstderr:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + return; + } + + let mut invalid = exponential_point(); + invalid.scale = -5; + for (temporality, points, reason, accepted, rejected) in [ + ( + AggregationTemporality::Delta, + vec![exponential_point(); 2], + "delta_temporality", + 0, + 2, + ), + ( + AggregationTemporality::Unspecified, + vec![exponential_point(); 2], + "unspecified_temporality", + 0, + 2, + ), + ( + AggregationTemporality::Cumulative, + vec![exponential_point(), invalid], + "invalid_data_point", + 1, + 1, + ), + ( + AggregationTemporality::Cumulative, + vec![exponential_point()], + "invalid_data_point", + 1, + 0, + ), + ] { + let counter = + OTLP_EXPONENTIAL_HISTOGRAM_REJECTED_DATA_POINTS.with_label_values(&[reason]); + let before = counter.get(); + let mut request = + metrics_request(vec![exponential_metric("latency", points, temporality)]); + request.resource_metrics[0].resource = Some(Resource { + attributes: vec![keyvalue("service.name", "api")], + ..Default::default() + }); + let mut ctx = descriptor_ctx(); + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert_eq!(conversion.outcome.accepted_data_points, accepted); + assert_eq!(conversion.outcome.rejected_data_points, rejected); + assert_eq!(counter.get() - before, rejected as u64, "{reason}"); + assert_eq!(conversion.resource_info.is_some(), accepted > 0); + + let before = counter.get(); + let empty = metrics_request(vec![exponential_metric("empty", vec![], temporality)]); + assert_eq!( + to_grpc_insert_requests(empty, &mut ctx) + .unwrap() + .outcome + .rejected_data_points, + 0 + ); + assert_eq!(counter.get(), before, "empty metric: {reason}"); + } + } + #[test] fn test_exponential_histogram_default_context_and_partial_outcome() { let request = metrics_request(vec![ From b547a3e36957b615886ae6a38305c80a277af775 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:33:45 +0000 Subject: [PATCH 038/133] ci: update dev-builder image tag (#9355) Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 1379d86683c..9144432b66d 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-04f07a91-20260921031255 +DEV_BUILDER_IMAGE_TAG ?= 1.96.1-822042c7-20260924091315 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu From 3c2aac0a55900721ee3d9d93937fb2f3d15a8ee4 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 24 Sep 2026 09:39:21 +0000 Subject: [PATCH 039/133] feat: add manual series index reconciliation (#9323) * feat(mito): support manual series index reconciliation Signed-off-by: evenyag * feat(storage): route series index build requests Signed-off-by: evenyag * feat(admin): add BUILD_SERIES_INDEX Signed-off-by: evenyag * test: specify compaction type in series index fixtures Signed-off-by: evenyag * test: correct series index SQL fixtures and error assertions Signed-off-by: evenyag * test(compat): preserve legacy index rebuild across upgrades Signed-off-by: evenyag * test(sql): cover series index admin validation Signed-off-by: evenyag * refactor(mito): bound series index maintenance queue Signed-off-by: evenyag * docs: remove series index how-to guide Signed-off-by: evenyag * test: remove index build upgrade compatibility case Signed-off-by: evenyag * fix: address series index reconciliation review feedback Signed-off-by: evenyag * fix: bound manual series index reconciliation admission Signed-off-by: evenyag * chore: update greptime-proto to merged index build options Signed-off-by: evenyag --------- Signed-off-by: evenyag --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/common/function/src/admin.rs | 3 +- .../function/src/admin/build_index_table.rs | 85 ++- src/metric-engine/src/engine.rs | 31 + src/mito2/src/engine/index_build_test.rs | 57 +- src/mito2/src/series_index/task.rs | 568 ++++++++++++++++-- src/mito2/src/series_index/tests.rs | 137 ++++- src/mito2/src/worker.rs | 16 +- src/mito2/src/worker/handle_rebuild_index.rs | 61 +- src/operator/src/request.rs | 45 +- src/store-api/src/region_request.rs | 39 +- src/table/src/requests.rs | 4 +- tests-integration/src/standalone.rs | 19 +- tests-integration/tests/build_series_index.rs | 148 +++++ tests-integration/tests/main.rs | 1 + 16 files changed, 1111 insertions(+), 107 deletions(-) create mode 100644 tests-integration/tests/build_series_index.rs diff --git a/Cargo.lock b/Cargo.lock index fca6847c3db..2cae7d03322 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5996,7 +5996,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=3d232f406fad574e8c630fb1c304fe2ebb013c9c#3d232f406fad574e8c630fb1c304fe2ebb013c9c" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=9533078d6fd4bdaddbca0750cb9bf7ad64b2e568#9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 5cdc3cf63a2..5dc37afceba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -160,7 +160,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 = "3d232f406fad574e8c630fb1c304fe2ebb013c9c" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" } hex = "0.4" hostname = "0.4.0" http = "1" diff --git a/src/common/function/src/admin.rs b/src/common/function/src/admin.rs index ff8172ebb5a..79182dd74e9 100644 --- a/src/common/function/src/admin.rs +++ b/src/common/function/src/admin.rs @@ -34,7 +34,7 @@ use reconcile_catalog::ReconcileCatalogFunction; use reconcile_database::ReconcileDatabaseFunction; use reconcile_table::ReconcileTableFunction; -use crate::admin::build_index_table::BuildIndexFunction; +use crate::admin::build_index_table::{BuildIndexFunction, BuildSeriesIndexFunction}; use crate::admin::discard_unflushed_data::DiscardUnflushedDataFunction; use crate::flush_flow::FlushFlowFunction; use crate::function_registry::FunctionRegistry; @@ -53,6 +53,7 @@ impl AdminFunction { registry.register(GcRegionsFunction::factory()); registry.register(GcTableFunction::factory()); registry.register(BuildIndexFunction::factory()); + registry.register(BuildSeriesIndexFunction::factory()); registry.register(FlushFlowFunction::factory()); registry.register(ReconcileCatalogFunction::factory()); registry.register(ReconcileDatabaseFunction::factory()); diff --git a/src/common/function/src/admin/build_index_table.rs b/src/common/function/src/admin/build_index_table.rs index 155f198c792..24026aead8d 100644 --- a/src/common/function/src/admin/build_index_table.rs +++ b/src/common/function/src/admin/build_index_table.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use api::v1::region::build_index_request; use arrow::datatypes::DataType as ArrowDataType; use common_error::ext::BoxedError; use common_macro::admin_fn; @@ -38,6 +39,45 @@ pub(crate) async fn build_index( table_mutation_handler: &TableMutationHandlerRef, query_ctx: &QueryContextRef, params: &[ValueRef<'_>], +) -> Result { + build_index_impl( + table_mutation_handler, + query_ctx, + params, + "build_index", + None, + ) + .await +} + +/// Reconciles series indexes on every physical data-region leader and waits for publication. +#[admin_fn( + name = BuildSeriesIndexFunction, + display_name = build_series_index, + sig_fn = build_index_signature, + ret = uint64 +)] +pub(crate) async fn build_series_index( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], +) -> Result { + build_index_impl( + table_mutation_handler, + query_ctx, + params, + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())), + ) + .await +} + +async fn build_index_impl( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], + function: &str, + options: Option, ) -> Result { ensure!( params.len() == 1, @@ -51,7 +91,7 @@ pub(crate) async fn build_index( let ValueRef::String(table_name) = params[0] else { return UnsupportedInputDataTypeSnafu { - function: "build_index", + function, datatypes: params.iter().map(|v| v.data_type()).collect::>(), } .fail(); @@ -64,6 +104,7 @@ pub(crate) async fn build_index( let affected_rows = table_mutation_handler .build_index( BuildIndexTableRequest { + options, catalog_name, schema_name, table_name, @@ -78,3 +119,45 @@ pub(crate) async fn build_index( fn build_index_signature() -> Signature { Signature::uniform(1, vec![ArrowDataType::Utf8], Volatility::Immutable) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::function::FunctionContext; + + #[tokio::test] + async fn test_build_series_index_arguments() { + let ctx = FunctionContext::mock(); + let handler = ctx.state.table_mutation_handler.as_ref().unwrap(); + for params in [ + vec![], + vec![ValueRef::Int32(1)], + vec![ValueRef::Null], + vec![ValueRef::String("a"), ValueRef::String("b")], + ] { + assert!( + build_index_impl( + handler, + &ctx.query_ctx, + ¶ms, + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())) + ) + .await + .is_err() + ); + } + assert_eq!( + Value::UInt64(42), + build_index_impl( + handler, + &ctx.query_ctx, + &[ValueRef::String("greptime.public.test")], + "build_series_index", + Some(build_index_request::Options::SeriesIndex(Default::default())), + ) + .await + .unwrap() + ); + } +} diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 18d6ff6f197..ec180839526 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -614,6 +614,37 @@ mod test { TestEnv, build_rows, create_logical_region_request, row_schema_with_tags, }; + #[tokio::test] + async fn test_build_series_index_forwarding() { + use api::v1::region::build_index_request; + use store_api::region_request::RegionBuildIndexRequest; + + let env = TestEnv::new().await; + env.init_metric_region().await; + let request = || { + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(build_index_request::Options::SeriesIndex(Default::default())), + }) + }; + // Default Mito config disables series indexes: reaching that error proves forwarding. + let error = env + .metric() + .handle_request(env.default_physical_region_id(), request()) + .await + .unwrap_err(); + let error = format!("{error:?}"); + assert!(error.contains("series index is disabled"), "{error}"); + let error = env + .metric() + .handle_request(env.default_logical_region_id(), request()) + .await + .unwrap_err(); + assert_eq!( + common_error::status_code::StatusCode::Unsupported, + error.status_code() + ); + } + #[tokio::test] async fn test_discard_unflushed_data_only() { let env = TestEnv::new().await; diff --git a/src/mito2/src/engine/index_build_test.rs b/src/mito2/src/engine/index_build_test.rs index d2d59b038ff..6d2241846b4 100644 --- a/src/mito2/src/engine/index_build_test.rs +++ b/src/mito2/src/engine/index_build_test.rs @@ -563,7 +563,7 @@ async fn run_index_publication_compaction_race(phase: IndexPublicationPhase, gc_ engine_for_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -705,7 +705,7 @@ async fn test_index_build_uses_physical_file_region_and_logical_manifest_region( engine .handle_request( target_region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await .unwrap(); @@ -1391,7 +1391,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 0); // Trigger manual index build task and make sure index file is built without flush or compaction. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); listener.wait_finish(1).await; let scanner = engine @@ -1403,7 +1403,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Test idempotency: Second manual index build request on the same file. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; let scanner = engine @@ -1416,7 +1416,7 @@ async fn test_index_build_type_manual_basic() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Test idempotency again: Third manual index build request to further verify. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; let scanner = engine @@ -1475,7 +1475,7 @@ async fn test_index_build_type_manual_consistency() { assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); // Check index build task for consistent file will be skipped. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); // Reopen the region to ensure the task wasn't skipped due to insufficient time. reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; @@ -1505,7 +1505,7 @@ async fn test_index_build_type_manual_consistency() { // SyncColumns won't trigger index build. assert_listener_counts(&listener, 1, 1); - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); listener.wait_finish(2).await; // previous 1 + new 1 // Because the file is inconsistent, new index build task is triggered. @@ -1581,7 +1581,7 @@ async fn test_index_build_type_manual_duplicate_in_flight() { // so we must spawn it in a separate task. let engine_clone = engine.clone(); let first_handle = tokio::spawn(async move { - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine_clone.handle_request(region_id, request).await }); @@ -1596,7 +1596,7 @@ async fn test_index_build_type_manual_duplicate_in_flight() { // Issue the second manual BuildIndex for the same region/file. // Since the file is already in building_files (from the first manual build), // schedule_build detects the duplicate and calls on_index_build_abort. - let request = RegionRequest::BuildIndex(RegionBuildIndexRequest {}); + let request = RegionRequest::BuildIndex(RegionBuildIndexRequest::default()); engine.handle_request(region_id, request).await.unwrap(); // The second request should have been aborted as duplicate. @@ -1679,7 +1679,7 @@ async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { engine_for_old_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -1707,7 +1707,7 @@ async fn test_reopen_waits_for_active_index_build_of_previous_incarnation() { engine_for_new_build .handle_request( region_id, - RegionRequest::BuildIndex(RegionBuildIndexRequest {}), + RegionRequest::BuildIndex(RegionBuildIndexRequest::default()), ) .await }); @@ -1879,3 +1879,38 @@ async fn test_index_build_type_compact_abort_race() { assert_eq!(scanner.num_files(), 1); assert_eq!(num_of_index_files(&engine, &scanner, region_id).await, 1); } + +#[tokio::test] +async fn test_build_series_index_rejects_non_metric_region() { + let mut env = TestEnv::with_prefix("series-unsupported-schema").await; + let engine = env + .create_engine(MitoConfig { + experimental_enable_series_index: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let error = engine + .handle_request( + region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(api::v1::region::build_index_request::Options::SeriesIndex( + Default::default(), + )), + }), + ) + .await + .unwrap_err(); + assert!( + error.to_string().contains("sparse metric metadata"), + "{error}" + ); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/series_index/task.rs b/src/mito2/src/series_index/task.rs index 5579cf87ce4..6ce340827a1 100644 --- a/src/mito2/src/series_index/task.rs +++ b/src/mito2/src/series_index/task.rs @@ -20,32 +20,52 @@ use std::time::Duration; use common_telemetry::{info, warn}; use object_store::ObjectStore; -use tokio::sync::Notify; -use tokio::sync::mpsc::UnboundedReceiver; +use snafu::ensure; +use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver}; +use tokio::sync::oneshot; use tokio::task::JoinHandle; -use tokio::time::{Instant, MissedTickBehavior}; +use tokio::time::Instant; +use crate::error::{InvalidRequestSnafu, RegionBusySnafu, Result, WorkerStoppedSnafu}; use crate::metrics::{ SERIES_INDEX_CAPACITY_DEFERRED, SERIES_INDEX_DISK_BYTES, SERIES_INDEX_RECONCILE_TOTAL, }; -use crate::region::{RegionLeaderState, RegionMapRef, RegionRoleState}; -use crate::series_index::maintenance::reconcile_series_indexes; +use crate::read::series_candidate::is_sparse_metric_metadata; +use crate::region::{MitoRegionRef, RegionLeaderState, RegionMapRef, RegionRoleState}; +use crate::series_index::maintenance::{ReconcileStats, reconcile_series_indexes}; use crate::series_index::purger::{IndexFilePurger, PurgeRequest, run_index_purge_task}; use crate::time_provider::TimeProviderRef; -/// Shared lifecycle state for a worker's series-index task. +/// Commands serialized with periodic series-index maintenance. +#[derive(Debug)] +pub(crate) enum SeriesIndexCommand { + Wake, + Reconcile { + region: MitoRegionRef, + sender: oneshot::Sender>, + }, + Stop, +} + +/// Shared lifecycle state and command sender for a worker's series-index task. #[derive(Debug)] pub(crate) struct SeriesIndexTaskState { + worker_id: u32, running: AtomicBool, - notify: Notify, + sender: Sender, } impl SeriesIndexTaskState { - pub(crate) fn new() -> Self { - Self { - running: AtomicBool::new(true), - notify: Notify::new(), - } + pub(crate) fn new(worker_id: u32, channel_size: usize) -> (Self, Receiver) { + let (sender, receiver) = mpsc::channel(channel_size); + ( + Self { + worker_id, + running: AtomicBool::new(true), + sender, + }, + receiver, + ) } pub(crate) fn is_running(&self) -> bool { @@ -53,17 +73,37 @@ impl SeriesIndexTaskState { } pub(crate) fn wake(&self) { - self.notify.notify_one(); + // Queued work already wakes the task. Periodic maintenance still runs even + // if the queued command only reconciles one region. + if self.is_running() && self.sender.capacity() == self.sender.max_capacity() { + let _ = self.sender.try_send(SeriesIndexCommand::Wake); + } + } + + /// Returns a receiver for build completion, or RegionBusy if the queue is full. + pub(crate) fn try_reconcile( + &self, + region: MitoRegionRef, + ) -> Result>> { + ensure!(self.is_running(), WorkerStoppedSnafu { id: self.worker_id }); + let region_id = region.region_id; + let (sender, receiver) = oneshot::channel(); + self.sender + .try_send(SeriesIndexCommand::Reconcile { region, sender }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => RegionBusySnafu { region_id }.build(), + mpsc::error::TrySendError::Closed(_) => { + WorkerStoppedSnafu { id: self.worker_id }.build() + } + })?; + Ok(receiver) } pub(crate) fn stop(&self) { - self.running.store(false, Ordering::Release); - // Retain a permit if maintenance has not started waiting yet. - self.notify.notify_one(); - } - - pub(crate) async fn notified(&self) { - self.notify.notified().await; + if self.running.swap(false, Ordering::AcqRel) { + // A full queue already wakes the task, which checks the running flag. + let _ = self.sender.try_send(SeriesIndexCommand::Stop); + } } } @@ -74,6 +114,7 @@ pub(crate) fn spawn_series_index_tasks( store: ObjectStore, regions: RegionMapRef, state: Arc, + receiver: Receiver, bucket_width: Duration, purger: IndexFilePurger, purge_receiver: UnboundedReceiver, @@ -100,6 +141,7 @@ pub(crate) fn spawn_series_index_tasks( max_size, reported_usage: 0, state, + receiver, interval, time_provider, enable_range_index, @@ -120,6 +162,7 @@ struct SeriesIndexTask { reported_usage: u64, worker_id: u32, state: Arc, + receiver: Receiver, interval: Duration, time_provider: TimeProviderRef, enable_range_index: bool, @@ -131,21 +174,98 @@ impl SeriesIndexTask { let worker_id = self.worker_id; info!("Start series-index background task, worker: {worker_id}"); let interval = self.time_provider.wait_duration(self.interval); - let mut timer = tokio::time::interval_at(Instant::now() + interval, interval); - // Schedule future ticks from a late tick rather than the original cadence. - timer.set_missed_tick_behavior(MissedTickBehavior::Delay); + let mut deadline = Instant::now() + interval; while self.state.is_running() { - tokio::select! { - _ = self.state.notified() => {} - _ = timer.tick() => {} - } - if self.state.is_running() { + // Check before receiving so a busy channel cannot starve periodic maintenance. + if Instant::now() >= deadline { self.maintain().await; + deadline = Instant::now() + interval; + continue; + } + match tokio::time::timeout_at(deadline, self.receiver.recv()).await { + Err(_) => continue, + Ok(None | Some(SeriesIndexCommand::Stop)) => break, + Ok(Some(SeriesIndexCommand::Wake)) => { + if self.state.is_running() { + self.maintain().await; + deadline = Instant::now() + interval; + } + } + Ok(Some(SeriesIndexCommand::Reconcile { region, sender })) => { + let result = self.reconcile_manual(region).await; + let _ = sender.send(result); + } + } + } + self.state.stop(); + self.receiver.close(); + while let Some(command) = self.receiver.recv().await { + if let SeriesIndexCommand::Reconcile { sender, .. } = command { + let _ = sender.send(WorkerStoppedSnafu { id: worker_id }.fail()); } } info!("Stop series-index background task, worker: {worker_id}"); } + /// A manual request must still refer to the same writable region instance. + fn validate_manual_region(&self, region: &MitoRegionRef) -> Result<()> { + ensure!( + self.state.is_running(), + WorkerStoppedSnafu { id: self.worker_id } + ); + let current = self.regions.writable_non_staging_region(region.region_id)?; + ensure!( + Arc::ptr_eq(¤t, region), + InvalidRequestSnafu { + region_id: region.region_id, + reason: "region was replaced during series-index reconciliation", + } + ); + ensure!( + is_sparse_metric_metadata(®ion.version().metadata), + InvalidRequestSnafu { + region_id: region.region_id, + reason: "series indexes require sparse metric metadata", + } + ); + Ok(()) + } + + async fn reconcile_manual(&mut self, region: MitoRegionRef) -> Result { + self.validate_manual_region(®ion)?; + self.refresh_usage(); + let result = self.reconcile_region(region.clone()).await; + self.refresh_usage(); + let result = result?; + self.validate_manual_region(®ion)?; + Ok(result) + } + + async fn reconcile_region(&self, region: MitoRegionRef) -> Result { + // Full capacity defers builds, but cleanup must still reclaim published usage. + let allow_builds = self.disk_usage.load(Ordering::Relaxed) < self.max_size; + if !allow_builds { + SERIES_INDEX_CAPACITY_DEFERRED.inc(); + } + let result = reconcile_series_indexes( + self.worker_id, + self.store.clone(), + region, + self.bucket_width, + self.time_provider.current_time_millis(), + self.purger.clone(), + self.enable_range_index, + allow_builds, + ) + .await; + if result.is_err() { + SERIES_INDEX_RECONCILE_TOTAL + .with_label_values(&["failure"]) + .inc(); + } + result + } + /// Each worker contributes only its open regions, refreshed at maintenance boundaries. fn refresh_usage(&mut self) { let usage = self @@ -185,26 +305,7 @@ impl SeriesIndexTask { ) { continue; } - // Full capacity defers builds, but cleanup must still reclaim published usage. - let allow_builds = self.disk_usage.load(Ordering::Relaxed) < self.max_size; - if !allow_builds { - SERIES_INDEX_CAPACITY_DEFERRED.inc(); - } - if let Err(error) = reconcile_series_indexes( - self.worker_id, - self.store.clone(), - region.clone(), - self.bucket_width, - self.time_provider.current_time_millis(), - self.purger.clone(), - self.enable_range_index, - allow_builds, - ) - .await - { - SERIES_INDEX_RECONCILE_TOTAL - .with_label_values(&["failure"]) - .inc(); + if let Err(error) = self.reconcile_region(region.clone()).await { warn!(error; "Failed to reconcile series indexes, worker: {}, region: {}", self.worker_id, region.region_id); } self.refresh_usage(); @@ -220,8 +321,10 @@ impl Drop for SeriesIndexTask { #[cfg(test)] mod tests { + use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::services::Memory; - use store_api::region_engine::RegionEngine; + use store_api::region_engine::{RegionEngine, RegionRole}; + use tokio::sync::Notify; use super::*; use crate::region::RegionMap; @@ -230,6 +333,331 @@ mod tests { use crate::series_index::tests::prepare_region; use crate::test_util::TestEnv; + struct BlockingCatalogWriter { + inner: oio::Writer, + entered: Arc, + release: Arc, + } + + impl mock::Write for BlockingCatalogWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.inner.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + self.entered.notify_one(); + self.release.notified().await; + self.inner.close().await + } + + async fn abort(&mut self) -> mock::Result<()> { + self.inner.abort().await + } + } + + #[derive(Clone, Copy, Debug)] + enum LifecycleChange { + Remove, + Demote, + Stop, + } + + fn new_task(region: MitoRegionRef) -> (SeriesIndexTask, UnboundedReceiver) { + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, purge_receiver) = series_index_channel(store.clone()); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region); + let (state, receiver) = SeriesIndexTaskState::new(0, 1); + ( + SeriesIndexTask { + store, + regions, + bucket_width: Duration::from_secs(100), + purger, + disk_usage: Arc::default(), + max_size: u64::MAX, + reported_usage: 0, + worker_id: 0, + state: Arc::new(state), + receiver, + interval: Duration::from_secs(3600), + time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, + }, + purge_receiver, + ) + } + + #[test] + fn test_coalesce_wakeups_and_stop() { + let (state, mut receiver) = SeriesIndexTaskState::new(0, 2); + state.wake(); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Wake + )); + assert!(receiver.try_recv().is_err()); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Wake + )); + state.stop(); + state.stop(); + state.wake(); + assert!(matches!( + receiver.try_recv().unwrap(), + SeriesIndexCommand::Stop + )); + assert!(receiver.try_recv().is_err()); + } + + #[rstest::rstest] + #[case(true)] + #[case(false)] + #[tokio::test] + async fn test_manual_reconcile_publishes_before_reply(#[case] range_enabled: bool) { + let mut env = TestEnv::with_prefix("series-manual").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + task.enable_range_index = range_enabled; + let state = task.state.clone(); + let store = task.store.clone(); + let first = state.try_reconcile(region.clone()).unwrap(); + state.wake(); + assert_eq!(1, task.receiver.len()); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + let handle = tokio::spawn(task.run()); + let first = tokio::time::timeout(Duration::from_secs(10), first) + .await + .unwrap() + .unwrap() + .unwrap(); + // Admission recovers after completion; both builds reuse the same publication. + let second = state.try_reconcile(region.clone()).unwrap(); + let second = tokio::time::timeout(Duration::from_secs(10), second) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(1, first.built_series); + assert_eq!(if range_enabled { 4 } else { 0 }, first.built_range); + assert_eq!(0, second.built_series); + assert_eq!(0, second.built_range); + assert_eq!(1, region.series_index_version().series_indexes.len()); + assert!( + store + .exists(&series_catalog_path(region.region_id)) + .await + .unwrap() + ); + assert_eq!( + range_enabled, + store + .exists(&range_catalog_path(region.region_id)) + .await + .unwrap() + ); + state.stop(); + handle.await.unwrap(); + assert!(state.try_reconcile(region).is_err()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_manual_reconcile_respects_capacity_and_refreshes_usage() { + let mut env = TestEnv::with_prefix("series-manual-capacity").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + task.max_size = 0; + task.reconcile_manual(region.clone()).await.unwrap(); + assert_eq!(0, region.series_index_version().disk_usage()); + assert_eq!(0, task.disk_usage.load(Ordering::Relaxed)); + + task.max_size = u64::MAX; + task.reconcile_manual(region.clone()).await.unwrap(); + let bytes = region.series_index_version().disk_usage(); + assert!(bytes > 0); + assert_eq!(bytes, task.disk_usage.load(Ordering::Relaxed)); + let usage = task.disk_usage.clone(); + drop(task); + assert_eq!(0, usage.load(Ordering::Relaxed)); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_shutdown_fails_queued_reconciliation() { + let mut env = TestEnv::with_prefix("series-manual-stop").await; + let (engine, region) = prepare_region(&mut env).await; + let (task, _purge_receiver) = new_task(region.clone()); + let state = task.state.clone(); + let result = state.try_reconcile(region.clone()).unwrap(); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + state.stop(); + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::WorkerStopped { .. }) + )); + tokio::time::timeout(Duration::from_secs(10), async { + task.run().await; + assert!(matches!( + result.await.unwrap(), + Err(crate::error::Error::WorkerStopped { .. }) + )); + }) + .await + .unwrap(); + assert!(region.series_index_version().series_indexes.is_empty()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_manual_reconcile_rejects_closed_region() { + let mut env = TestEnv::with_prefix("series-manual-closed").await; + let (engine, region) = prepare_region(&mut env).await; + let (task, _purge_receiver) = new_task(region.clone()); + let result = task.state.try_reconcile(region.clone()).unwrap(); + task.regions.remove_region(region.region_id); + let state = task.state.clone(); + let handle = tokio::spawn(task.run()); + assert!(matches!( + result.await.unwrap(), + Err(crate::error::Error::RegionNotFound { .. }) + )); + assert!(region.series_index_version().series_indexes.is_empty()); + state.stop(); + handle.await.unwrap(); + engine.stop().await.unwrap(); + } + + #[rstest::rstest] + #[case::remove(LifecycleChange::Remove)] + #[case::demote(LifecycleChange::Demote)] + #[case::stop(LifecycleChange::Stop)] + #[tokio::test] + async fn test_manual_reconcile_rejects_lifecycle_change_during_io( + #[case] change: LifecycleChange, + ) { + let mut env = TestEnv::with_prefix("series-manual-lifecycle").await; + let (engine, region) = prepare_region(&mut env).await; + let (mut task, _purge_receiver) = new_task(region.clone()); + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let writer_entered = entered.clone(); + let writer_release = release.clone(); + let catalog_path = series_catalog_path(region.region_id); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + if path == catalog_path { + Box::new(BlockingCatalogWriter { + inner, + entered: writer_entered.clone(), + release: writer_release.clone(), + }) + } else { + inner + } + })) + .build() + .unwrap(); + task.store = task.store.clone().layer(layer); + let state = task.state.clone(); + let regions = task.regions.clone(); + let mut result = state.try_reconcile(region.clone()).unwrap(); + let handle = tokio::spawn(task.run()); + tokio::time::timeout(Duration::from_secs(10), entered.notified()) + .await + .unwrap(); + assert!(matches!( + result.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + + // A canceled queued caller still occupies capacity while catalog I/O is blocked. + drop(state.try_reconcile(region.clone()).unwrap()); + for _ in 0..3 { + assert!(matches!( + state.try_reconcile(region.clone()), + Err(crate::error::Error::RegionBusy { .. }) + )); + } + + // The initial validation passed; invalidate the request during catalog I/O. + match change { + LifecycleChange::Remove => { + regions.remove_region(region.region_id).unwrap(); + } + LifecycleChange::Demote => region.set_role(RegionRole::Follower), + LifecycleChange::Stop => state.stop(), + } + release.notify_one(); + let result = tokio::time::timeout(Duration::from_secs(10), result) + .await + .unwrap() + .unwrap(); + state.stop(); + tokio::time::timeout(Duration::from_secs(10), handle) + .await + .unwrap() + .unwrap(); + engine.stop().await.unwrap(); + + match change { + LifecycleChange::Remove => assert!(matches!( + result, + Err(crate::error::Error::RegionNotFound { .. }) + )), + LifecycleChange::Demote => assert!(matches!( + result, + Err(crate::error::Error::RegionState { .. }) + )), + LifecycleChange::Stop => assert!(matches!( + result, + Err(crate::error::Error::WorkerStopped { .. }) + )), + } + } + + #[tokio::test] + async fn test_manual_traffic_does_not_postpone_periodic_maintenance() { + let mut env = TestEnv::with_prefix("series-manual-deadline").await; + let (engine, region) = prepare_region(&mut env).await; + region.switch_state_to_staging(RegionLeaderState::Writable); + region + .manifest_ctx + .exit_staging(region.region_id, RegionRoleState::Follower) + .unwrap(); + let (mut task, _purge_receiver) = new_task(region.clone()); + task.interval = Duration::from_millis(20); + let state = task.state.clone(); + let handle = tokio::spawn(task.run()); + // Manual builds on followers fail. Only a periodic sweep can publish coverage. + tokio::time::timeout(Duration::from_secs(10), async { + while region.series_index_version().series_indexes.is_empty() { + assert!( + state + .try_reconcile(region.clone()) + .unwrap() + .await + .unwrap() + .is_err() + ); + } + }) + .await + .unwrap(); + state.stop(); + handle.await.unwrap(); + engine.stop().await.unwrap(); + } + #[rstest::rstest] #[case::follower(RegionRoleState::Follower, true)] #[case::writable(RegionRoleState::Leader(RegionLeaderState::Writable), true)] @@ -254,6 +682,7 @@ mod tests { let (purger, _receiver) = series_index_channel(store.clone()); let regions = Arc::new(RegionMap::default()); regions.insert_region(region.clone()); + let (state, receiver) = SeriesIndexTaskState::new(0, 2); let mut task = SeriesIndexTask { store: store.clone(), regions, @@ -263,11 +692,16 @@ mod tests { max_size: u64::MAX, reported_usage: 0, worker_id: 0, - state: Arc::new(SeriesIndexTaskState::new()), + state: Arc::new(state), + receiver, interval: Duration::from_secs(3600), time_provider: Arc::new(crate::time_provider::StdTimeProvider), enable_range_index: true, }; + assert_eq!( + role == RegionRoleState::Leader(RegionLeaderState::Writable), + task.validate_manual_region(®ion).is_ok(), + ); task.maintain().await; assert_eq!( builds, @@ -332,6 +766,7 @@ mod tests { regions.insert_region(region.clone()); let clock = Arc::new(MockTimeProvider::new(0)); let usage = Arc::new(AtomicU64::new(0)); + let (state, command_receiver) = SeriesIndexTaskState::new(0, 1); let mut task = SeriesIndexTask { store: store.clone(), regions, @@ -341,7 +776,8 @@ mod tests { max_size: u64::MAX, reported_usage: 0, worker_id: 0, - state: Arc::new(SeriesIndexTaskState::new()), + state: Arc::new(state), + receiver: command_receiver, interval: Duration::from_secs(3600), time_provider: clock.clone(), enable_range_index, @@ -465,19 +901,23 @@ mod tests { let (engine, region) = prepare_region(&mut env).await; let usage = Arc::new(AtomicU64::new(0)); let store = ObjectStore::new(Memory::default()).unwrap(); - let make_task = || SeriesIndexTask { - store: store.clone(), - regions: Arc::new(RegionMap::default()), - bucket_width: Duration::from_secs(100), - purger: series_index_channel(store.clone()).0, - disk_usage: usage.clone(), - max_size: 1024, - reported_usage: 0, - worker_id: 0, - state: Arc::new(SeriesIndexTaskState::new()), - interval: Duration::from_secs(3600), - time_provider: Arc::new(crate::time_provider::StdTimeProvider), - enable_range_index: true, + let make_task = || { + let (state, command_receiver) = SeriesIndexTaskState::new(0, 1); + SeriesIndexTask { + store: store.clone(), + regions: Arc::new(RegionMap::default()), + bucket_width: Duration::from_secs(100), + purger: series_index_channel(store.clone()).0, + disk_usage: usage.clone(), + max_size: 1024, + reported_usage: 0, + worker_id: 0, + state: Arc::new(state), + receiver: command_receiver, + interval: Duration::from_secs(3600), + time_provider: Arc::new(crate::time_provider::StdTimeProvider), + enable_range_index: true, + } }; let mut task = make_task(); task.regions.insert_region(region.clone()); diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs index 8ee898277bc..631bdc75b87 100644 --- a/src/mito2/src/series_index/tests.rs +++ b/src/mito2/src/series_index/tests.rs @@ -52,12 +52,23 @@ pub(crate) async fn prepare_region_with_timestamps( env: &mut TestEnv, timestamps: &[i64], ) -> (MitoEngine, MitoRegionRef) { - let engine = env - .create_engine(MitoConfig { + prepare_region_with_config( + env, + timestamps, + MitoConfig { min_compaction_interval: Duration::from_secs(3600), ..Default::default() - }) - .await; + }, + ) + .await +} + +async fn prepare_region_with_config( + env: &mut TestEnv, + timestamps: &[i64], + config: MitoConfig, +) -> (MitoEngine, MitoRegionRef) { + let engine = env.create_engine(config).await; let metadata = Arc::new(sst_region_metadata_with_encoding( PrimaryKeyEncoding::Sparse, )); @@ -834,7 +845,8 @@ async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { region.version_control.alter_options(options); let store = ObjectStore::new(Memory::default()).unwrap(); let (purger, receiver) = series_index_channel(store.clone()); - let state = Arc::new(super::task::SeriesIndexTaskState::new()); + let (state, commands) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); let clock = Arc::new(crate::time_provider::mock::MockTimeProvider::new(0)); // A notification issued before the task starts must also trigger maintenance. state.wake(); @@ -845,6 +857,7 @@ async fn test_maintenance_wakeup_and_timer(#[case] enable_range_index: bool) { store, regions, state.clone(), + commands, Duration::from_secs(100), purger, receiver, @@ -917,13 +930,15 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { let control = load_version_control(&store, region_id, &purger).await; let snapshot = control.current(); assert_eq!(&entry, snapshot.series_indexes[&entry.index_uuid].entry()); - let state = Arc::new(super::task::SeriesIndexTaskState::new()); + let (state, commands) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); state.stop(); super::task::spawn_series_index_tasks( 0, store.clone(), Arc::new(RegionMap::default()), state, + commands, Duration::from_secs(100), purger, receiver, @@ -951,3 +966,113 @@ async fn test_drop_catalogs_and_retained_snapshot_after_task_stop() { .await .unwrap(); } + +#[rstest::rstest] +#[case(true)] +#[case(false)] +#[tokio::test] +async fn test_build_series_index_region_request(#[case] enabled: bool) { + use api::v1::region::build_index_request; + use store_api::region_request::RegionBuildIndexRequest; + + let mut env = TestEnv::with_prefix("series-manual-request").await; + let config = MitoConfig { + experimental_enable_series_index: enabled, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (engine, region) = + prepare_region_with_config(&mut env, &[1000, 2000, 3000, 4000], config).await; + for _ in 0..2 { + let result = engine + .handle_request( + region.region_id, + RegionRequest::BuildIndex(RegionBuildIndexRequest { + options: Some(build_index_request::Options::SeriesIndex(Default::default())), + }), + ) + .await; + if enabled { + assert_eq!(0, result.unwrap().affected_rows); + assert!(!region.series_index_version().series_indexes.is_empty()); + } else { + assert!( + result + .unwrap_err() + .to_string() + .contains("series index is disabled") + ); + } + } + engine.stop().await.unwrap(); +} + +#[rstest::rstest] +#[case::build(false)] +#[case::catalog(true)] +#[tokio::test] +async fn test_manual_reconcile_returns_storage_failure(#[case] catalog_failure: bool) { + let mut env = TestEnv::with_prefix("series-manual-failure").await; + let (engine, region) = prepare_region(&mut env).await; + let catalog = series_catalog_path(region.region_id); + let fail_once = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let fail = fail_once.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + let target = if catalog_failure { + path == catalog + } else { + path.contains("/series/") + }; + Box::new(FailingSeriesWriter { + inner, + fail: target && fail.load(std::sync::atomic::Ordering::Relaxed), + }) + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let (purger, purge_receiver) = series_index_channel(store.clone()); + let (state, receiver) = super::task::SeriesIndexTaskState::new(0, 2); + let state = Arc::new(state); + let regions = Arc::new(RegionMap::default()); + regions.insert_region(region.clone()); + let task = super::task::spawn_series_index_tasks( + 0, + store, + regions, + state.clone(), + receiver, + Duration::from_secs(100), + purger, + purge_receiver, + Arc::default(), + u64::MAX, + Duration::from_secs(3600), + Arc::new(StdTimeProvider), + false, + ); + assert!( + state + .try_reconcile(region.clone()) + .unwrap() + .await + .unwrap() + .is_err() + ); + assert!(region.series_index_version().series_indexes.is_empty()); + fail_once.store(false, std::sync::atomic::Ordering::Relaxed); + assert_eq!( + 1, + state + .try_reconcile(region) + .unwrap() + .await + .unwrap() + .unwrap() + .built_series + ); + state.stop(); + task.await.unwrap(); + engine.stop().await.unwrap(); +} diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 0f67d8d2a1b..2eda4fff49f 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -605,16 +605,21 @@ impl WorkerStarter { let (sender, receiver) = mpsc::channel(self.config.worker_channel_size); let running = Arc::new(AtomicBool::new(true)); - let series_index_task_state = self - .series_index_store - .as_ref() - .map(|_| Arc::new(SeriesIndexTaskState::new())); + let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some() + { + let (state, receiver) = + SeriesIndexTaskState::new(self.id, self.config.worker_channel_size); + (Some(Arc::new(state)), Some(receiver)) + } else { + (None, None) + }; let mut series_index_purger = None; let series_index_handle = self .series_index_store .clone() .zip(series_index_task_state.clone()) - .map(|(store, state)| { + .zip(series_index_receiver) + .map(|((store, state), receiver)| { let (purger, purge_receiver) = series_index_channel(store.clone()); series_index_purger = Some(purger.clone()); spawn_series_index_tasks( @@ -622,6 +627,7 @@ impl WorkerStarter { store, regions.clone(), state, + receiver, self.config.experimental_series_index_bucket_width, purger, purge_receiver, diff --git a/src/mito2/src/worker/handle_rebuild_index.rs b/src/mito2/src/worker/handle_rebuild_index.rs index 7b203e9af04..59f42ef4232 100644 --- a/src/mito2/src/worker/handle_rebuild_index.rs +++ b/src/mito2/src/worker/handle_rebuild_index.rs @@ -17,13 +17,14 @@ use std::collections::HashMap; use std::sync::Arc; +use api::v1::region::build_index_request; use common_telemetry::{debug, error, warn}; use store_api::region_request::RegionBuildIndexRequest; use store_api::storage::{FileId, RegionId}; use tokio::sync::mpsc; use crate::cache::CacheStrategy; -use crate::error::Result; +use crate::error::{InvalidRequestSnafu, Result, WorkerStoppedSnafu}; use crate::manifest::action::RegionEdit; use crate::metrics::INDEX_PUBLICATION_STALE_TOTAL; use crate::region::version::VersionRef; @@ -96,18 +97,56 @@ impl RegionWorkerLoop { pub(crate) async fn handle_build_index_request( &mut self, region_id: RegionId, - _req: RegionBuildIndexRequest, + req: RegionBuildIndexRequest, sender: OptionOutputTx, ) { - self.handle_rebuild_index( - BuildIndexRequest { - region_id, - build_type: IndexBuildType::Manual, - file_metas: Vec::new(), - }, - sender, - ) - .await; + match req.options { + Some(build_index_request::Options::SeriesIndex(_)) => { + self.handle_build_series_index_request(region_id, sender); + } + None | Some(build_index_request::Options::SstIndex(_)) => { + self.handle_rebuild_index( + BuildIndexRequest { + region_id, + build_type: IndexBuildType::Manual, + file_metas: Vec::new(), + }, + sender, + ) + .await; + } + } + } + + /// Submits manual reconciliation without blocking the region worker on index I/O. + fn handle_build_series_index_request(&self, region_id: RegionId, mut sender: OptionOutputTx) { + let Some(region) = self.regions.writable_region_or(region_id, &mut sender) else { + return; + }; + let Some(state) = &self.series_index_task_state else { + sender.send( + InvalidRequestSnafu { + region_id, + reason: "series index is disabled; enable experimental_enable_series_index", + } + .fail(), + ); + return; + }; + let receiver = match state.try_reconcile(region) { + Ok(receiver) => receiver, + Err(error) => { + sender.send(Err(error)); + return; + } + }; + let worker_id = self.id; + common_runtime::spawn_global(async move { + let result = receiver + .await + .unwrap_or_else(|_| WorkerStoppedSnafu { id: worker_id }.fail()); + sender.send(result.map(|_| 0)); + }); } pub(crate) async fn handle_rebuild_index( diff --git a/src/operator/src/request.rs b/src/operator/src/request.rs index b46a6b657a6..ce3b3568651 100644 --- a/src/operator/src/request.rs +++ b/src/operator/src/request.rs @@ -18,7 +18,7 @@ use api::helper::to_pb_time_unit; use api::v1::region::region_request::Body as RegionRequestBody; use api::v1::region::{ BuildIndexRequest, CompactRequest, CompactionTimeRange, FlushRequest, RegionRequestHeader, - TruncateRequest, Unflushed, truncate_request, + TruncateRequest, Unflushed, build_index_request, truncate_request, }; use catalog::CatalogManagerRef; use common_catalog::build_db_string; @@ -107,6 +107,30 @@ impl Requester { request: BuildIndexTableRequest, ctx: QueryContextRef, ) -> Result { + if matches!( + request.options, + Some(build_index_request::Options::SeriesIndex(_)) + ) { + let table = self + .catalog_manager + .table( + &request.catalog_name, + &request.schema_name, + &request.table_name, + None, + ) + .await + .context(CatalogSnafu)?; + let table = table.with_context(|| TableNotFoundSnafu { + table_name: common_catalog::format_full_table_name( + &request.catalog_name, + &request.schema_name, + &request.table_name, + ), + })?; + let info = table.table_info(); + ensure_build_series_index_supported(&info.meta.engine, info.is_physical_table())?; + } let partitions = &self .get_table_partition_info( &request.catalog_name, @@ -121,6 +145,7 @@ impl Requester { .map(|partition| { RegionRequestBody::BuildIndex(BuildIndexRequest { region_id: partition.id.into(), + options: request.options, }) }) .collect(); @@ -404,6 +429,16 @@ impl Requester { } } +fn ensure_build_series_index_supported(engine: &str, is_physical_table: bool) -> Result<()> { + ensure!( + engine == METRIC_ENGINE && is_physical_table, + NotSupportedSnafu { + feat: "building series indexes requires a physical metric table", + } + ); + Ok(()) +} + fn ensure_discard_unflushed_supported(engine: &str, is_physical_table: bool) -> Result<()> { ensure!( engine != METRIC_ENGINE || is_physical_table, @@ -422,6 +457,14 @@ mod tests { use super::*; + #[test] + fn test_build_series_index_requires_physical_metric_table() { + ensure_build_series_index_supported(METRIC_ENGINE, true).unwrap(); + assert!(ensure_build_series_index_supported(METRIC_ENGINE, false).is_err()); + assert!(ensure_build_series_index_supported("mito", false).is_err()); + assert!(ensure_build_series_index_supported("mito", true).is_err()); + } + #[test] fn test_to_pb_compaction_time_range_normalizes_mixed_units() { let range = TimestampRange::new( diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 113a2393ccc..762bf80c4b3 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -26,7 +26,8 @@ use api::v1::region::{ AlterRequest, AlterRequests, BuildIndexRequest, BulkInsertRequest, CleanUpRequest as PbCleanUpRequest, CloseRequest, CompactRequest, CreateRequest, CreateRequests, DeleteRequests, DropRequest, DropRequests, FlushRequest, InsertRequests, - OpenRequest, TruncateRequest, alter_request, compact_request, region_request, truncate_request, + OpenRequest, TruncateRequest, alter_request, build_index_request, compact_request, + region_request, truncate_request, }; use api::v1::{ self, Analyzer, ArrowIpc, FulltextBackend as PbFulltextBackend, Option as PbOption, Rows, @@ -442,7 +443,9 @@ fn make_region_build_index(index: BuildIndexRequest) -> Result, +} /// Truncate region request. #[derive(Debug)] @@ -2054,6 +2060,33 @@ mod tests { use super::*; use crate::metadata::RegionMetadataBuilder; + #[test] + fn test_build_index_options_round_trip() { + use prost::Message; + + for options in [ + None, + Some(build_index_request::Options::SstIndex(Default::default())), + Some(build_index_request::Options::SeriesIndex(Default::default())), + ] { + let request = BuildIndexRequest { + region_id: 42, + options, + }; + let decoded = BuildIndexRequest::decode(request.encode_to_vec().as_slice()).unwrap(); + let requests = make_region_build_index(decoded).unwrap(); + assert_eq!(1, requests.len()); + assert_eq!(RegionId::from_u64(42), requests[0].0); + let RegionRequest::BuildIndex(request) = &requests[0].1 else { + panic!("expected build-index request"); + }; + assert_eq!(options, request.options); + } + // The legacy wire message contains only region_id (field 1). + let legacy = BuildIndexRequest::decode(&[0x08, 42][..]).unwrap(); + assert!(legacy.options.is_none()); + } + #[test] fn test_make_region_puts_preserves_skip_wal() { let region_id = RegionId::new(42, 3); diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 2c519b32a5e..899438edf52 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -30,7 +30,7 @@ use datatypes::prelude::VectorRef; use datatypes::schema::{ ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, }; -use greptime_proto::v1::region::compact_request; +use greptime_proto::v1::region::{build_index_request, compact_request}; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; use store_api::metric_engine_consts::{ @@ -805,6 +805,8 @@ pub struct FlushTableRequest { #[derive(Debug, Clone, Default)] pub struct BuildIndexTableRequest { + /// The index build mode. Absent options select SST indexes. + pub options: Option, pub catalog_name: String, pub schema_name: String, pub table_name: String, diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 6fc8ce231af..3fa1d394489 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -91,6 +91,7 @@ pub struct GreptimeDbStandaloneBuilder { experimental_metric_export: bool, logical_batcher: Option, table_batcher: BatcherOptions, + mito_config: Option, } impl GreptimeDbStandaloneBuilder { @@ -113,9 +114,17 @@ impl GreptimeDbStandaloneBuilder { experimental_metric_export: false, logical_batcher: None, table_batcher: BatcherOptions::default(), + mito_config: None, } } + /// Overrides the Mito configuration for this test instance. + #[must_use] + pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self { + self.mito_config = Some(config); + self + } + /// Enables experimental Metric export for the standalone test instance. #[must_use] pub fn with_experimental_metric_export(mut self) -> Self { @@ -386,7 +395,7 @@ impl GreptimeDbStandaloneBuilder { let (procedure_manager, event_recorder_handle) = standalone::build_procedure_manager(kv_backend.clone(), procedure_config); - let standalone_opts = StandaloneOptions { + let mut standalone_opts = StandaloneOptions { storage: opts.storage, procedure: procedure_config, metadata_store: kv_backend_config, @@ -408,6 +417,14 @@ impl GreptimeDbStandaloneBuilder { ..StandaloneOptions::default() }; + if let Some(config) = &self.mito_config { + for engine in &mut standalone_opts.region_engine { + if let datanode::config::RegionEngineConfig::Mito(mito) = engine { + *mito = config.clone(); + } + } + } + self.build_with( kv_backend, guard, diff --git a/tests-integration/tests/build_series_index.rs b/tests-integration/tests/build_series_index.rs new file mode 100644 index 00000000000..8c0e91aa35b --- /dev/null +++ b/tests-integration/tests/build_series_index.rs @@ -0,0 +1,148 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use datanode::config::RegionEngineConfig; +use frontend::instance::Instance; +use mito2::config::MitoConfig; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{ + StorageType, create_tmp_dir_and_datanode_opts, execute_sql, try_execute_sql, +}; + +fn config() -> MitoConfig { + MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + } +} + +async fn exercise_admin(frontend: &Arc) { + execute_sql( + frontend, + r#" + CREATE TABLE series_physical (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) + PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') + ENGINE = metric WITH ( + physical_metric_table = 'true', 'compaction.type' = 'twcs', + 'compaction.twcs.time_window' = '1h', + 'compaction.twcs.active_window.trigger_file_num' = '100', + 'compaction.twcs.inactive_window.trigger_file_num' = '100' + ) + "#, + ) + .await; + execute_sql( + frontend, + r#" + CREATE TABLE series_logical (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) + ENGINE = metric WITH (on_physical_table = 'series_physical') + "#, + ) + .await; + // Reconciliation requires at least four SSTs in each bucket. + for timestamp in [1000, 2000, 3000, 4000] { + execute_sql( + frontend, + &format!( + "INSERT INTO series_logical (ts, host, val) VALUES ({timestamp}, 'a', 1), ({timestamp}, 'z', 2)" + ), + ) + .await; + execute_sql(frontend, "ADMIN FLUSH_TABLE('series_physical')").await; + } + for _ in 0..2 { + let output = execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('series_physical')").await; + assert!(output.data.pretty_print().await.contains("| 0")); + } + // The function must never silently expand a logical table to its shared physical table. + let error = try_execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('series_logical')") + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("physical metric table")); + for sql in [ + "ADMIN BUILD_SERIES_INDEX()", + "ADMIN BUILD_SERIES_INDEX(1)", + "ADMIN BUILD_SERIES_INDEX('missing')", + ] { + assert!(try_execute_sql(frontend, sql).await.is_err(), "{sql}"); + } + execute_sql(frontend, "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX)").await; + assert!( + try_execute_sql(frontend, "ADMIN BUILD_SERIES_INDEX('ordinary')") + .await + .is_err() + ); + execute_sql(frontend, "ADMIN BUILD_INDEX('ordinary')").await; + let output = execute_sql(frontend, "SELECT count(*) AS n FROM series_logical").await; + assert!(output.data.pretty_print().await.contains("| 8")); +} + +fn catalog_count(home: &str) -> usize { + let root = Path::new(home).join("series_index"); + std::fs::read_dir(root) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(|path| path.join("series-index.json").is_file()) + .count() +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_build_series_index_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("build-series-index") + .with_mito_config(config()) + .build() + .await; + exercise_admin(standalone.fe_instance()).await; + assert_eq!(2, catalog_count(&standalone.opts.storage.data_home)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_build_series_index_distributed() { + let builder = GreptimeDbClusterBuilder::new("build-series-index").await; + let mut options = Vec::new(); + let mut guards = Vec::new(); + for id in 1..=2 { + let (mut opts, guard) = create_tmp_dir_and_datanode_opts( + StorageType::File, + vec![], + &format!("build-series-index-{id}"), + Default::default(), + Default::default(), + ); + opts.node_id = Some(id); + for engine in &mut opts.region_engine { + if let RegionEngineConfig::Mito(mito) = engine { + *mito = config(); + } + } + options.push(opts); + guards.push(guard); + } + let cluster = builder.build_with(options, false, guards).await; + exercise_admin(cluster.fe_instance()).await; + assert_eq!( + 2, + cluster + .datanode_options + .iter() + .map(|opts| catalog_count(&opts.storage.data_home)) + .sum::() + ); +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 278a7ba3851..f84b954e23c 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -15,6 +15,7 @@ #![recursion_limit = "256"] mod admin_function_event; +mod build_series_index; mod database_ddl_event; mod dict_groupby_sst; mod event_recorder_test_util; From 194bc2fb3c48ea1c55cff8fdd5886daa8b60d8ad Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 24 Sep 2026 11:07:29 +0000 Subject: [PATCH 040/133] feat(log-store): add the object store WAL durable write path (#9320) * feat(log-store): add the object store WAL durable write path ObjectStoreLogStore can now write. append_batch admits entries into one open batch and assigns object-sequence-major entry ids at admission. The batch is sealed by size, by the flush interval, or before a region would run past the position range, and sealed batches are uploaded with at most four conditional creates in flight, started in sequence order. Created objects are indexed in sequence order, and an append is acknowledged only once its object is durable and indexed. A transient create failure rolls back the failed batch and every later batch unless a later object is already durable, in which case the store poisons itself with a history-gap error. A conflicting object, an encoding or catalog error, or taking the last representable sequence poisons the store. obsolete now goes through the actor and raises the sequence floor together with the watermark, refusing with a retryable error while the next sequence is not settled. stop drops the open batch and the batches whose create has not started, and lets creates in flight finish. Only the durable acknowledgement mode exists. A testing feature exposes hooks to wait for admissions, seal the open batch, and hold or fail creates. Signed-off-by: jeremyhi * test(log-store): synchronize object store WAL write tests with the actor Tests that assert nothing happened round-trip a command through the actor instead of yielding the test task, the conflict test waits for the object it expects, and the obsolete-behind-stop test holds the command channel itself so the unanswered command is deterministic. Signed-off-by: jeremyhi * fix(log-store): bound admitted WAL appends and never reuse a failed sequence A create that reports an error may still have written its object, so the sequences of failed batches are no longer handed out again: the next batch keeps the sequence after the last sealed one and retries are assigned new ids. A retry batched differently can no longer conflict with that object. As the next sequence never moves back, the sequence floor of obsolete only waits for an open batch that has handed out ids. Appends now arrive on their own bounded channel, which the actor stops reading while MAX_SEALED_BATCHES batches wait to become durable, so a stalled object store holds callers back instead of growing the backlog, while stop and obsolete are still handled. Signed-off-by: jeremyhi * fix(log-store): reserve room for two seals before admitting a WAL append One admission can seal the open batch before an append that would exhaust its positions and then the append's own batch, so the actor takes an append only while two more sealed batches fit under MAX_SEALED_BATCHES. Signed-off-by: jeremyhi * feat(log-store): chain object store WAL objects and recover the latest chain A conditional create can fail with an unknown outcome while its object is stored, or still lands later, and Mito reuses the row sequences of a failed append. Replaying such an object next to a later acknowledged one can let the unacknowledged rows win after a restart. Format version 2 gives every object header its writer's epoch, a link to the object it extends (sequence and writer instance) and a header CRC32, and allows objects without segments. Recovery replays only the chain ending at the complete object with the largest epoch and sequence: a link holds when its predecessor is present with the recorded writer instance, or is missing below every present object. Objects off the chain are orphans that are never replayed but keep their sequences. Each open writes an empty object that starts an epoch above every present object, linked to the recovered tip, before it accepts writes, so a late object of an earlier instance never ends the chain. A start object that meets an object of an earlier epoch moves to the next sequence; one of an equal or later epoch fails the open. Signed-off-by: jeremyhi * fix(log-store): break the WAL chain at every missing predecessor Accepting a missing predecessor below every present object lets a late object that lands below the chain change which links hold. Nothing collects objects yet, so a missing predecessor now always breaks the link, and recovery fails when objects are present but none completes a chain. Signed-off-by: jeremyhi * fix(log-store): keep the object store WAL format at version 1 The object store WAL has not been enabled anywhere, so no object in the previous layout exists and the chained header can stay version 1. Signed-off-by: jeremyhi * refactor(log-store): link object store WAL objects by epoch Only one store instance writes under an epoch, since every open starts an epoch above every present object and a start object that meets the same or a later epoch fails the open. The epoch therefore identifies the instance, and the random writer instance id is dropped from the header. A link now records the sequence and the epoch of the object it extends, and holds when the predecessor carries that epoch. The header shrinks to 46 bytes, and the store logs its epoch when it opens. Signed-off-by: jeremyhi * fix(log-store): fail an open whose start object is already present Without a random writer instance, two opens that recover the same objects encode byte-identical start objects, and a conditional create treats the same bytes as its own retry. A start object that is already present therefore fails the open with a retryable error instead of letting both opens claim the epoch; the next open counts the object and starts a later epoch. Signed-off-by: jeremyhi * fix(log-store): derive the WAL epoch from the claimed start sequence Two opens can recover different listings when a late object lands across a sequence gap between them, pick the same largest epoch plus one, and both create their start objects under different sequences. The epoch of an instance is now one above the sequence its start object claims, so a successful create decides the epoch and no two instances share one. It stays above every epoch recovery listed, and an object that carries an epoch above the next sequence fails the open. Signed-off-by: jeremyhi * fix(log-store): never run a held WAL create after the store is dropped The create test hook ignored the closed hold channel, so a create parked when the store was dropped could still run. It now returns without creating. Drop the per-admission bookkeeping of issued entry ids, which nothing reads, and move the parked I/O documentation to its helper. Signed-off-by: jeremyhi * test(log-store): cover a WAL start object stored with an unknown outcome Add a fault-injection test in which the create of the start object stores the object but reports an error: the open fails without moving to another sequence, and the next open counts the stored object and claims a later epoch. Rename the test helper that writes a whole object from a given header to put_object_with_header, and drop a needless clone. Signed-off-by: jeremyhi * test(log-store): make the held-create drop test deterministic Keep the actor running while the store drops the hold sender, so the parked create always completes on the closed channel instead of racing the actor's exit. Drop a comment that restates epoch_of. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/Cargo.toml | 3 + src/log-store/src/error.rs | 51 +- src/log-store/src/object_store_wal/store.rs | 2191 ++++++++++++++++++- 3 files changed, 2149 insertions(+), 96 deletions(-) diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 5d068ca344d..7e4da1557e1 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -11,6 +11,9 @@ protoc-rust.workspace = true [lints] workspace = true +[features] +testing = [] + [dependencies] async-stream.workspace = true async-trait.workspace = true diff --git a/src/log-store/src/error.rs b/src/log-store/src/error.rs index 8be0b8bd3da..85000dbea3c 100644 --- a/src/log-store/src/error.rs +++ b/src/log-store/src/error.rs @@ -350,6 +350,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Incomplete multipart WAL entry of region {}", region_id))] + IncompleteWalEntry { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("WAL object sequence is exhausted, last sequence: {}", last_object_seq))] WalObjectSequenceExhausted { last_object_seq: u64, @@ -357,6 +364,16 @@ pub enum Error { location: Location, }, + #[snafu(display( + "WAL object sequence {} is not settled: the open batch has assigned entry ids under it", + object_seq + ))] + WalObjectSequenceUnsettled { + object_seq: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "WAL entry positions of region {} in one object are exhausted", region_id @@ -374,6 +391,20 @@ pub enum Error { location: Location, }, + #[snafu(display( + "WAL object {} was written by the earlier epoch {}, this store writes epoch {}", + path, + existing_epoch, + epoch + ))] + StaleWalObject { + path: String, + existing_epoch: u64, + epoch: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "WAL object {} that starts epoch {} was already present, so the open cannot tell whether it wrote it", path, @@ -417,16 +448,15 @@ pub enum Error { location: Location, }, - #[snafu(display("Object store WAL operation failed"))] - ObjectStoreWal { - source: Arc, + #[snafu(display("Object store WAL log store is stopped"))] + ObjectStoreWalStopped { #[snafu(implicit)] location: Location, }, - /// Appending to the object store WAL is unsupported. - #[snafu(display("Object store WAL operation is not supported"))] - UnsupportedObjectStoreWalOperation { + #[snafu(display("Object store WAL operation failed"))] + ObjectStoreWal { + source: Arc, #[snafu(implicit)] location: Location, }, @@ -464,6 +494,7 @@ impl ErrorExt for Error { | InvalidWalObjectStore { .. } | MismatchedWalPrefix { .. } | MismatchedWalRegion { .. } + | IncompleteWalEntry { .. } | InvalidWalEntryRange { .. } => StatusCode::InvalidArguments, StartWalTask { .. } | StopWalTask { .. } @@ -477,14 +508,15 @@ impl ErrorExt for Error { | OrderedBatchProducerStopped { .. } | WaitProduceResultReceiver { .. } | WaitDumpIndex { .. } - | MetaLengthExceededLimit { .. } => StatusCode::Internal, + | MetaLengthExceededLimit { .. } + | ObjectStoreWalStopped { .. } => StatusCode::Internal, CorruptedWalObject { .. } | WalObjectConflict { .. } | WalObjectSequenceExhausted { .. } | WalEntryPositionExhausted { .. } => StatusCode::Unexpected, + WalObjectSequenceUnsettled { .. } => StatusCode::IllegalState, - UnsupportedObjectStoreWalOperation { .. } => StatusCode::Unsupported, InvalidWalObject { source, .. } => source.status_code(), ObjectStoreWal { source, .. } => source.status_code(), @@ -493,6 +525,7 @@ impl ErrorExt for Error { | WriteIndex { .. } | ReadIndex { .. } | WalObjectStore { .. } + | StaleWalObject { .. } | UnconfirmedWalEpochStart { .. } | Io { .. } => StatusCode::StorageUnavailable, // Raft engine @@ -529,6 +562,8 @@ impl ErrorExt for Error { FetchEntry { .. } | RaftEngine { .. } | AddEntryLogBatch { .. } + | WalObjectSequenceUnsettled { .. } + | StaleWalObject { .. } | UnconfirmedWalEpochStart { .. } => RetryHint::Retryable, ProduceRecord { error, .. } => match error { rskafka::client::producer::Error::Client(error) => { diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 7b024bb1a80..ee925ae0a38 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Object store WAL construction, recovery and region reads. +//! Object store WAL construction, recovery, writes and region reads. -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt; use std::ops::Range; use std::sync::atomic::{AtomicBool, Ordering}; @@ -25,6 +25,8 @@ use async_stream::try_stream; use bytes::Bytes; use common_telemetry::info; use common_wal::config::object_store::ObjectStoreWalConfig; +use futures::future::BoxFuture; +use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStore; use snafu::{IntoError, OptionExt, ResultExt, ensure}; @@ -32,25 +34,38 @@ use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; use store_api::storage::RegionId; +#[cfg(any(test, feature = "testing"))] +use tokio::sync::watch; use tokio::sync::{mpsc, oneshot}; +use tokio::time::MissedTickBehavior; use crate::error::{ - CorruptedWalObjectSnafu, Error, InvalidProviderSnafu, InvalidWalObjectSnafu, - InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, MismatchedWalRegionSnafu, - ObjectStoreWalSnafu, Result, UnconfirmedWalEpochStartSnafu, - UnsupportedObjectStoreWalOperationSnafu, WalObjectSequenceExhaustedSnafu, + CorruptedWalObjectSnafu, Error, IncompleteWalEntrySnafu, InvalidProviderSnafu, + InvalidWalObjectSnafu, InvalidWalObjectStoreSnafu, MismatchedWalPrefixSnafu, + MismatchedWalRegionSnafu, ObjectStoreWalSnafu, ObjectStoreWalStoppedSnafu, Result, + StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, + WalObjectSequenceUnsettledSnafu, }; -use crate::object_store_wal::batch::OBJECT_SEQ_LIMIT; +use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ - ChainLink, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, TRAILER_LEN, - decode_footer, decode_header, decode_segment, decode_trailer, encode_object, footer_range, - verify_segment_ranges, + ChainLink, EncodedObject, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, + Record, TRAILER_LEN, decode_footer, decode_header, decode_segment, decode_trailer, + encode_object, footer_range, verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; const COMMAND_BUFFER: usize = 1024; +/// Appends queued for the actor; a caller beyond them waits to send. +const APPEND_BUFFER: usize = 16; const MIN_FLUSH_INTERVAL: Duration = Duration::from_millis(10); +const MAX_IN_FLIGHT_CREATES: usize = 4; +/// Number of sealed batches that may wait to be created or indexed. One +/// admission can seal two batches, the open batch before an append that would +/// exhaust its positions and the append's own batch, so the actor takes an +/// append from its channel only while two more fit: a stalled object store +/// holds callers back instead of growing the backlog. +const MAX_SEALED_BATCHES: usize = 2 * MAX_IN_FLIGHT_CREATES; /// Number of objects whose footers recovery fetches at a time. const RECOVERY_CONCURRENCY: usize = 8; /// Bytes recovery reads from the end of an object in one request. The window @@ -59,6 +74,16 @@ const RECOVERY_CONCURRENCY: usize = 8; const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// A log store over the immutable WAL objects under one prefix. +/// +/// Appends are admitted into an open batch. A background actor seals the batch +/// when it reaches the size limit or the flush interval elapses and creates +/// the object under the next sequence while it keeps admitting entries into +/// the next batch; up to [`MAX_IN_FLIGHT_CREATES`] creates run at a time. +/// Objects are indexed in the catalog and their batches acknowledged in +/// sequence order, so an acknowledged entry never has a missing predecessor. +/// An append returns once its object is durable and indexed. While +/// [`MAX_SEALED_BATCHES`] batches wait to become durable no append is +/// admitted. pub(crate) struct ObjectStoreLogStore { prefix: String, io: Arc, @@ -67,6 +92,13 @@ pub(crate) struct ObjectStoreLogStore { terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, + append_tx: mpsc::Sender, + #[cfg(any(test, feature = "testing"))] + admitted_appends: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + creates_held: watch::Sender, + #[cfg(any(test, feature = "testing"))] + creates_fail: Arc, } type ObsoleteEntryIds = Arc>>; @@ -112,7 +144,7 @@ impl ObjectStoreLogStore { } ); - positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; + let max_batch_bytes = positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; let Recovered { mut catalog, next_object_seq, @@ -144,31 +176,62 @@ impl ObjectStoreLogStore { path: io.object_path(start.object_seq), })?; let catalog = Arc::new(RwLock::new(catalog)); + let obsolete_entry_ids = ObsoleteEntryIds::default(); let terminal_error = TerminalError::default(); let stopped = Arc::new(AtomicBool::new(false)); let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); + let (append_tx, append_rx) = mpsc::channel(APPEND_BUFFER); + #[cfg(any(test, feature = "testing"))] + let (admitted_appends_tx, admitted_appends_rx) = watch::channel(0); + #[cfg(any(test, feature = "testing"))] + let (creates_held_tx, creates_held_rx) = watch::channel(false); + #[cfg(any(test, feature = "testing"))] + let creates_fail = Arc::new(AtomicBool::new(false)); + let actor = Actor { + io: io.clone(), + catalog: catalog.clone(), + obsolete_entry_ids: obsolete_entry_ids.clone(), terminal_error: terminal_error.clone(), stopped: stopped.clone(), command_rx, + append_rx, + open_batch: OpenBatch::new(max_batch_bytes), issued_entry_ids: durable_entry_ids, + pending: Vec::new(), + sealed: VecDeque::new(), + creates: FuturesUnordered::new(), + stop: Vec::new(), next_object_seq: start .object_seq .checked_add(1) .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT), epoch, - chain_tip: start, - stop: Vec::new(), + last_indexed: start, + flush_interval: config.flush_interval, + #[cfg(any(test, feature = "testing"))] + admitted_appends: admitted_appends_tx, + #[cfg(any(test, feature = "testing"))] + creates_held: creates_held_rx, + #[cfg(any(test, feature = "testing"))] + creates_fail: creates_fail.clone(), }; common_runtime::spawn_global(actor.run()); Ok(Arc::new(Self { prefix, io, catalog, - obsolete_entry_ids: ObsoleteEntryIds::default(), + obsolete_entry_ids, terminal_error, stopped, command_tx, + append_tx, + #[cfg(any(test, feature = "testing"))] + admitted_appends: admitted_appends_rx, + #[cfg(any(test, feature = "testing"))] + creates_held: creates_held_tx, + #[cfg(any(test, feature = "testing"))] + creates_fail, })) } @@ -236,10 +299,70 @@ fn positive_bytes(bytes: u64, name: &str) -> Result { }) } +#[cfg(any(test, feature = "testing"))] +impl ObjectStoreLogStore { + /// Waits until the actor has admitted at least `expected` append calls + /// since the store was built. + pub(crate) async fn wait_for_admitted_appends(&self, expected: usize) -> Result<()> { + self.admitted_appends + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Seals the open batch regardless of its size and age and returns once + /// its object is durable and indexed, or with the error that failed it. + pub(crate) async fn seal_open_batch(&self) -> Result<()> { + ensure!( + !self.stopped.load(Ordering::Acquire), + ObjectStoreWalStoppedSnafu + ); + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(Command::Seal { + response: response_tx, + }) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? + } + + /// Parks every conditional create that starts from now on until + /// [`release_creates`](Self::release_creates), so a test can observe + /// entries that are admitted but not durable. A create that is parked when + /// the store is dropped never runs. + pub(crate) fn hold_creates(&self) { + self.creates_held.send_replace(true); + } + + /// Lets the creates parked by [`hold_creates`](Self::hold_creates) run. + pub(crate) fn release_creates(&self) { + self.creates_held.send_replace(false); + } + + /// Makes every create that runs from now on fail with a transient object + /// store error instead of writing. + pub(crate) fn fail_creates(&self) { + self.creates_fail.store(true, Ordering::Release); + } + + /// Sets the stopped flag without sending the stop command, which is the + /// state a store is in between the two steps of [`stop`](LogStore::stop). + pub(crate) fn begin_stop(&self) { + self.stopped.store(true, Ordering::Release); + } +} + #[async_trait::async_trait] impl LogStore for ObjectStoreLogStore { type Error = Error; + /// Stops the store. Creates in flight run to completion and acknowledge + /// their entries if they succeed. async fn stop(&self) -> Result<()> { self.stopped.store(true, Ordering::Release); let (response_tx, response_rx) = oneshot::channel(); @@ -256,8 +379,34 @@ impl LogStore for ObjectStoreLogStore { Ok(()) } - async fn append_batch(&self, _entries: Vec) -> Result { - UnsupportedObjectStoreWalOperationSnafu.fail() + async fn append_batch(&self, entries: Vec) -> Result { + ensure!( + !self.stopped.load(Ordering::Acquire), + ObjectStoreWalStoppedSnafu + ); + self.check_terminal()?; + if entries.is_empty() { + return Ok(AppendBatchResponse::default()); + } + for entry in &entries { + let region_id = self.region_of(entry.provider())?; + ensure!( + region_id == entry.region_id(), + MismatchedWalRegionSnafu { + region_id: entry.region_id(), + reason: format!("provider belongs to region {region_id}"), + } + ); + ensure!(entry.is_complete(), IncompleteWalEntrySnafu { region_id }); + } + + let (response_tx, response_rx) = oneshot::channel(); + self.append_tx + .send((entries, response_tx)) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? } /// Reads the provider's region from `entry_id`, hiding obsolete entries. @@ -347,7 +496,20 @@ impl LogStore for ObjectStoreLogStore { .collect()) } - /// Records a memory-only watermark. The caller re-establishes it after a restart. + /// Moves the obsolete watermark of the region up to `entry_id` and makes + /// every id the region is assigned from now on greater than `entry_id`: + /// the sequence of the next object is raised above the object `entry_id` + /// names unless it is there already, so a watermark assigned under + /// another prefix, or one whose object this prefix no longer holds, is + /// never passed by a new id. Zero names no object and needs no floor. + /// + /// The watermark and the floor are applied together by the actor: a call + /// cancelled before its command is queued publishes neither, while a + /// command already queued still applies both. While the open batch holds + /// ids under the next sequence the floor cannot move, and the call fails + /// with [`Error::WalObjectSequenceUnsettled`] and records neither. The + /// watermark is kept in memory; the caller re-establishes it after a + /// restart. async fn obsolete( &self, provider: &Provider, @@ -355,8 +517,26 @@ impl LogStore for ObjectStoreLogStore { entry_id: EntryId, ) -> Result<()> { self.check_region(provider, region_id)?; - record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); - Ok(()) + let (response_tx, response_rx) = oneshot::channel(); + let command = Command::Obsolete { + region_id, + entry_id, + response: response_tx, + }; + let answered = match self.command_tx.send(command).await { + Ok(()) => response_rx.await.ok(), + Err(_) => None, + }; + match answered { + Some(result) => result, + // The actor exited, before the command was queued or with it + // still queued behind the stop: nothing is assigned an id any + // more, so the watermark alone is consistent. + None => { + record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + Ok(()) + } + } } /// Hides every entry of the region in memory. The caller re-establishes @@ -390,40 +570,569 @@ impl LogStore for ObjectStoreLogStore { } } +type AppendResponse = oneshot::Sender>; + +type QueuedAppend = (Vec, AppendResponse); + enum Command { + /// Answered once the watermark is recorded and no id of the region at + /// or below `entry_id` can be assigned, or with the reason neither was + /// done. + Obsolete { + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + }, Stop { response: oneshot::Sender>, }, + #[cfg(any(test, feature = "testing"))] + Seal { + response: oneshot::Sender>, + }, } +/// An append waiting for the object that holds its entries. +struct PendingAppend { + last_entry_ids: HashMap, + response: AppendResponse, +} + +/// Where the conditional create of a sealed batch stands. +enum CreateState { + /// The create has not started because no slot was free. + Pending, + InFlight, + Created, +} + +/// A sealed and encoded batch that holds its object sequence and waits to be +/// created, indexed and acknowledged. +struct SealedBatch { + object_seq: u64, + bytes: Bytes, + footer: Vec, + waiters: Vec, + #[cfg(any(test, feature = "testing"))] + seal_waiters: Vec>>, + state: CreateState, +} + +impl SealedBatch { + fn fail(self, error: impl Fn() -> Error) { + for waiter in self.waiters { + let _ = waiter.response.send(Err(error())); + } + #[cfg(any(test, feature = "testing"))] + for waiter in self.seal_waiters { + let _ = waiter.send(Err(error())); + } + } +} + +type CreateOutcome = (u64, Result); + +/// The actor that owns the open batch and the sealed batches until they are +/// durable. +/// +/// Sealed batches form a pipeline in sequence order. A batch is created under +/// its sequence as soon as one of [`MAX_IN_FLIGHT_CREATES`] slots is free, but +/// it is indexed and acknowledged only once every earlier batch is, so the +/// acknowledged history of a region never has a missing predecessor. Every +/// batch links to its pipeline predecessor, the batch sealed before it, or to +/// the last indexed object when no earlier batch waits, so recovery replays a +/// batch only on the chain of a later object. The outcomes of a create are: +/// +/// | Situation | Sequence | Waiters | Store | +/// | --- | --- | --- | --- | +/// | created, or identical retry | advances | acknowledged in order | healthy | +/// | transient error | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | +/// | an object of an earlier epoch holds the sequence | never reused | as for a transient error | healthy: that object can never be on the chain | +/// | an object of the same or a later epoch holds the sequence, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | +/// | created at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence | +/// +/// A create still in flight when its batch failed runs to completion and its +/// outcome is ignored: whatever it stored is off the chain. After `stop` began +/// nothing is admitted and no create starts; creates in flight run to +/// completion and acknowledge if they succeed. +/// +/// Appends arrive on their own channel, which the actor does not read while +/// [`MAX_SEALED_BATCHES`] batches wait, so commands such as `stop` are handled +/// while admission is held back. struct Actor { + io: Arc, + catalog: Arc>, + obsolete_entry_ids: ObsoleteEntryIds, terminal_error: TerminalError, stopped: Arc, command_rx: mpsc::Receiver, + append_rx: mpsc::Receiver, + open_batch: OpenBatch, issued_entry_ids: HashMap, - next_object_seq: Option, - epoch: u64, - /// The object a newly sealed batch extends. - chain_tip: ChainLink, + /// Waiters of the open batch. + pending: Vec, + /// Batches that are not durable yet, in sequence order. + sealed: VecDeque, + /// Creates in flight, including those of batches that already failed. + creates: FuturesUnordered>, + /// Callers of `stop`, answered once nothing is in flight. stop: Vec>>, + /// Sequence of the next sealed batch, `None` once the sequence is + /// exhausted. The open batch assigns its entry ids from it. It is above + /// every sealed batch and never moves back while the store runs. + next_object_seq: Option, + /// Epoch of this instance, carried by every object it writes. + epoch: u64, + /// The last indexed object, which is the start object until a batch is + /// indexed. A batch sealed while no earlier batch waits extends it. + last_indexed: ChainLink, + flush_interval: Duration, + #[cfg(any(test, feature = "testing"))] + admitted_appends: watch::Sender, + #[cfg(any(test, feature = "testing"))] + creates_held: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + creates_fail: Arc, } impl Actor { async fn run(mut self) { - while let Some(Command::Stop { response }) = self.command_rx.recv().await { - self.handle_stop(response); + let mut interval = tokio::time::interval(self.flush_interval); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + // The first tick completes immediately. + interval.tick().await; + + loop { + tokio::select! { + _ = interval.tick() => { + // Nothing starts after stop began. + if !self.is_stopped() { + self.flush_open_batch(); + } + } + Some((object_seq, result)) = self.creates.next(), if !self.creates.is_empty() => { + self.on_create_completed(object_seq, result); + } + Some((entries, response)) = self.append_rx.recv(), if self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { + self.handle_append(entries, response); + } + command = self.command_rx.recv() => match command { + Some(Command::Obsolete { region_id, entry_id, response }) => { + self.handle_obsolete(region_id, entry_id, response); + } + Some(Command::Stop { response }) => { + self.handle_stop(response); + } + #[cfg(any(test, feature = "testing"))] + Some(Command::Seal { response }) => { + self.handle_seal(response); + } + // Every sender is gone: the store was dropped without + // `stop`. The creates in flight are dropped with the actor. + None => return, + }, + } if self.finish_stop() { return; } } } - fn handle_stop(&mut self, response: oneshot::Sender>) { - self.stop.push(response); + fn handle_append(&mut self, entries: Vec, response: AppendResponse) { + // The append was queued before `stop` set the flag; nothing that is + // not durable yet gets admitted once it is set. + if self.is_stopped() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + return; + } + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + self.admit(entries, response); } + /// Admits `entries` into the open batch, which assigns their ids under + /// the next object sequence. The caller waits for the object. + fn admit(&mut self, entries: Vec, response: AppendResponse) { + // A region would run past the position range of the open batch: the + // batch is sealed and the entries open the next one. The size limit + // seals a batch long before a million entries of one region, so this + // is a theoretical bound. + if !self.open_batch.is_empty() && self.open_batch.would_exhaust_positions(&entries) { + self.flush_open_batch(); + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + } + let Some(object_seq) = self.next_object_seq else { + let error = self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + let _ = response.send(Err(shared(&error))); + return; + }; + let last_entry_ids = match self.open_batch.admit(object_seq, entries) { + Ok(last_entry_ids) => last_entry_ids, + // Nothing was admitted: the append alone runs past the position + // range, which no object can hold. + Err(error) => { + let _ = response.send(Err(error)); + return; + } + }; + self.pending.push(PendingAppend { + last_entry_ids, + response, + }); + #[cfg(any(test, feature = "testing"))] + self.admitted_appends.send_modify(|count| *count += 1); + if self.open_batch.should_seal() { + self.flush_open_batch(); + } + } + + /// Seals the open batch as the object `next_object_seq` and starts its + /// create when a slot is free. Returns whether a batch was sealed. + fn flush_open_batch(&mut self) -> bool { + if self.open_batch.is_empty() { + return false; + } + if let Some(error) = terminal(&self.terminal_error) { + self.fail_unacknowledged(|| shared(&error)); + return false; + } + let Some(object_seq) = self.next_object_seq else { + self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + return false; + }; + + let (entries, _) = self.open_batch.seal(); + let header = Header { + object_seq, + epoch: self.epoch, + prev: Some( + self.sealed + .back() + .map_or(self.last_indexed, |batch| ChainLink { + object_seq: batch.object_seq, + epoch: self.epoch, + }), + ), + }; + let encoded = match encode_batch(header, entries) { + Ok(encoded) => encoded, + Err(error) => { + self.poison(error); + return false; + } + }; + self.next_object_seq = object_seq + .checked_add(1) + .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT); + if self.next_object_seq.is_none() { + // The batch takes the last representable sequence: it is created + // and acknowledged, but no later batch can be allocated one. + set_terminal( + &self.terminal_error, + WalObjectSequenceExhaustedSnafu { + last_object_seq: object_seq, + } + .build(), + ); + } + self.sealed.push_back(SealedBatch { + object_seq, + bytes: encoded.bytes, + footer: encoded.footer, + waiters: std::mem::take(&mut self.pending), + #[cfg(any(test, feature = "testing"))] + seal_waiters: Vec::new(), + state: CreateState::Pending, + }); + self.start_creates(); + true + } + + /// Starts the creates of pending batches in sequence order while fewer + /// than [`MAX_IN_FLIGHT_CREATES`] are in flight, counting the creates of + /// batches that already failed. Nothing starts once stop began. + fn start_creates(&mut self) { + if self.is_stopped() { + return; + } + let mut in_flight = self.creates.len(); + for batch in self.sealed.iter_mut() { + if in_flight >= MAX_IN_FLIGHT_CREATES { + break; + } + if !matches!(batch.state, CreateState::Pending) { + continue; + } + batch.state = CreateState::InFlight; + in_flight += 1; + let io = self.io.clone(); + let object_seq = batch.object_seq; + let bytes = batch.bytes.clone(); + let epoch = self.epoch; + #[cfg(any(test, feature = "testing"))] + let mut creates_held = self.creates_held.clone(); + #[cfg(any(test, feature = "testing"))] + let creates_fail = self.creates_fail.clone(); + self.creates.push(Box::pin(async move { + // The store was dropped while the create was parked: it never + // runs. + #[cfg(any(test, feature = "testing"))] + if creates_held.wait_for(|held| !*held).await.is_err() { + return (object_seq, Err(ObjectStoreWalStoppedSnafu.build())); + } + #[cfg(any(test, feature = "testing"))] + if creates_fail.load(Ordering::Acquire) { + let error = object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected create failure", + ) + .set_temporary(); + let result = Err(error).context(crate::error::WalObjectStoreSnafu { + operation: "write", + path: io.object_path(object_seq), + }); + return (object_seq, result); + } + let result = match io.put_if_absent(object_seq, bytes).await { + Err(error @ Error::WalObjectConflict { .. }) => { + stale_conflict(io.as_ref(), object_seq, epoch, error).await + } + result => result, + }; + (object_seq, result) + })); + } + } + + fn on_create_completed(&mut self, object_seq: u64, result: Result) { + // A batch that already failed, or that the store gave up on when it + // poisoned itself: the object may exist, but it is off the chain. + let Some(index) = self + .sealed + .iter() + .position(|batch| batch.object_seq == object_seq) + else { + self.start_creates(); + return; + }; + match result { + Ok(_) => self.sealed[index].state = CreateState::Created, + // The object store did not confirm the object, which may still + // exist or land later, or an earlier epoch holds the sequence and + // can never be on the chain. The caller retries the append itself. + Err(error @ (Error::WalObjectStore { .. } | Error::StaleWalObject { .. })) => { + self.roll_back(Arc::new(error)) + } + Err(error) => { + self.poison(error); + return; + } + } + self.settle(); + } + + /// Indexes and acknowledges the sealed batches from the front as far as + /// they are created, and starts the creates that a free slot allows. + fn settle(&mut self) { + while let Some(front) = self.sealed.front() { + if !matches!(front.state, CreateState::Created) || !self.index_front() { + break; + } + } + self.start_creates(); + } + + /// Indexes the created object at the front and acknowledges its waiters. + /// Returns false when the catalog rejected it, which poisons the store. + fn index_front(&mut self) -> bool { + let Some(front) = self.sealed.front() else { + return false; + }; + let indexed = self + .catalog + .write() + .unwrap_or_else(PoisonError::into_inner) + .insert_object(front.object_seq, front.footer.clone()); + if let Err(error) = indexed { + self.poison(error); + return false; + } + let Some(batch) = self.sealed.pop_front() else { + return false; + }; + self.last_indexed = ChainLink { + object_seq: batch.object_seq, + epoch: self.epoch, + }; + for waiter in batch.waiters { + let _ = waiter.response.send(Ok(AppendBatchResponse { + last_entry_ids: waiter.last_entry_ids, + })); + } + #[cfg(any(test, feature = "testing"))] + for waiter in batch.seal_waiters { + let _ = waiter.send(Ok(())); + } + true + } + + /// Fails every batch that is not indexed after a create failed + /// transiently, including batches already created and batches whose + /// create is in flight. Their sequences are not reused: a create that + /// reported an error may still have written its object. The next batch + /// links to the last indexed object, so no later chain reaches a failed + /// batch. Waiters of a store that was stopped meanwhile learn that + /// instead of the I/O error, like every other entry that never became + /// durable. + fn roll_back(&mut self, error: Arc) { + let stopped = self.is_stopped(); + let failure = || { + if stopped { + ObjectStoreWalStoppedSnafu.build() + } else { + shared(&error) + } + }; + for batch in self.sealed.drain(..) { + batch.fail(failure); + } + self.reset_open_batch(); + self.fail_unacknowledged(failure); + } + + /// Records `error` as terminal and fails every waiter that is not + /// acknowledged with it, or with the stopped error if the store was + /// stopped meanwhile. Creates in flight run to completion, but their + /// outcome is ignored: the entries of an object they create were never + /// acknowledged, like those of a crash between creation and + /// acknowledgement. Returns the recorded error. + fn poison(&mut self, error: Error) -> Arc { + let error = set_terminal(&self.terminal_error, error); + let stopped = self.is_stopped(); + let failure = || { + if stopped { + ObjectStoreWalStoppedSnafu.build() + } else { + shared(&error) + } + }; + for batch in self.sealed.drain(..) { + batch.fail(failure); + } + self.reset_open_batch(); + self.fail_unacknowledged(failure); + error + } + + /// Drops the entries of the open batch; the next admission hands out the + /// same ids again under the sequence the batch is at. + fn reset_open_batch(&mut self) { + self.open_batch.reset(); + } + + /// Fails the waiters of the open batch. + fn fail_unacknowledged(&mut self, error: impl Fn() -> Error) { + for pending in self.pending.drain(..) { + let _ = pending.response.send(Err(error())); + } + } + + /// Makes sure no id of the region at or below `entry_id` is assigned + /// from now on, then records the obsolete watermark of the region; + /// neither is done when the sequence cannot be raised. + fn handle_obsolete( + &mut self, + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + ) { + let result = self.raise_sequence_floor(entry_id).map(|_| { + record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + }); + let _ = response.send(result); + } + + /// Moves the sequence of the next object above the object that holds + /// `entry_id` unless it is there already. The move fails while the open + /// batch has handed out ids under the next sequence. A floor that does + /// not fit an entry id poisons the store like an exhausted sequence. + fn raise_sequence_floor(&mut self, entry_id: EntryId) -> Result<()> { + // Nothing is assigned an id after stop began. + if self.is_stopped() { + return Ok(()); + } + if let Some(error) = terminal(&self.terminal_error) { + return Err(shared(&error)); + } + let sequence_floor = sequence_floor(entry_id); + let Some(next_object_seq) = self.next_object_seq else { + return Ok(()); + }; + if next_object_seq >= sequence_floor { + return Ok(()); + } + ensure!( + self.open_batch.is_empty(), + WalObjectSequenceUnsettledSnafu { + object_seq: next_object_seq, + } + ); + if sequence_floor < OBJECT_SEQ_LIMIT { + self.next_object_seq = Some(sequence_floor); + Ok(()) + } else { + self.next_object_seq = None; + let error = self.poison( + WalObjectSequenceExhaustedSnafu { + last_object_seq: OBJECT_SEQ_LIMIT - 1, + } + .build(), + ); + Err(shared(&error)) + } + } + + /// Begins stopping. Nothing is admitted from now on; the open batch and + /// the batches whose create has not started are dropped. Stop is answered by [`finish_stop`](Self::finish_stop) + /// once nothing is in flight. + fn handle_stop(&mut self, response: oneshot::Sender>) { + self.stop.push(response); + self.reset_open_batch(); + for pending in self.pending.drain(..) { + let _ = pending + .response + .send(Err(ObjectStoreWalStoppedSnafu.build())); + } + if let Some(index) = self + .sealed + .iter() + .position(|batch| matches!(batch.state, CreateState::Pending)) + { + for batch in self.sealed.drain(index..) { + batch.fail(|| ObjectStoreWalStoppedSnafu.build()); + } + } + } + + /// Answers the callers of `stop` once every sealed batch is settled and + /// every create has completed. Returns true when the actor is done. fn finish_stop(&mut self) -> bool { - if !self.is_stopped() || self.stop.is_empty() { + if self.stop.is_empty() || !self.sealed.is_empty() || !self.creates.is_empty() { return false; } for response in self.stop.drain(..) { @@ -432,6 +1141,22 @@ impl Actor { true } + #[cfg(any(test, feature = "testing"))] + fn handle_seal(&mut self, response: oneshot::Sender>) { + if self.is_stopped() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + return; + } + if self.flush_open_batch() + && let Some(batch) = self.sealed.back_mut() + { + batch.seal_waiters.push(response); + return; + } + let result = terminal(&self.terminal_error).map_or(Ok(()), |error| Err(shared(&error))); + let _ = response.send(result); + } + fn is_stopped(&self) -> bool { self.stopped.load(Ordering::Acquire) } @@ -459,6 +1184,18 @@ fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) } +fn encode_batch(header: Header, entries: Vec) -> Result { + let records = entries + .into_iter() + .map(|entry| Record { + region_id: entry.region_id(), + entry_id: entry.entry_id(), + payload: Bytes::from(entry.into_bytes()), + }) + .collect::>(); + encode_object(header, &records) +} + /// What recovery rebuilt from the objects under a prefix. #[derive(Debug)] struct Recovered { @@ -644,11 +1381,7 @@ async fn start_epoch( .fail(); } Err(error @ Error::WalObjectConflict { .. }) => { - let head = io.get_range(object_seq, 0, HEADER_LEN as u64).await?; - let existing = decode_header(&head).with_context(|_| InvalidWalObjectSnafu { - path: io.object_path(object_seq), - })?; - if existing.epoch >= epoch { + if epoch_of(io, object_seq).await? >= epoch { return Err(error); } object_seq += 1; @@ -658,6 +1391,36 @@ async fn start_epoch( } } +async fn epoch_of(io: &dyn WalObjectIo, object_seq: u64) -> Result { + let head = io.get_range(object_seq, 0, HEADER_LEN as u64).await?; + let header = decode_header(&head).with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + })?; + Ok(header.epoch) +} + +/// Resolves a create of this store's `epoch` that met a different object: +/// an object of an earlier epoch can never be on the chain, so the create +/// fails like a transient error; an object of the same or a later epoch +/// belongs to another writer and `conflict` is returned. +async fn stale_conflict( + io: &dyn WalObjectIo, + object_seq: u64, + epoch: u64, + conflict: Error, +) -> Result { + let existing_epoch = epoch_of(io, object_seq).await?; + if existing_epoch >= epoch { + return Err(conflict); + } + StaleWalObjectSnafu { + path: io.object_path(object_seq), + existing_epoch, + epoch, + } + .fail() +} + /// Fetches and verifies the headers and footers of `objects`, up to /// `concurrency` objects at a time, and returns them ordered by object sequence /// whatever the order the fetches complete in. The first failure abandons the @@ -836,14 +1599,13 @@ mod tests { use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, RetryHint}; use object_store::services::Memory; + use store_api::logstore::entry::{MultiplePartEntry, MultiplePartHeader}; use tokio::time::timeout; use super::*; use crate::error::WalObjectStoreSnafu; - use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, entry_id}; - use crate::object_store_wal::format::{ - FOOTER_ENTRY_LEN, Header, Record, decode_object, encode_object, - }; + use crate::object_store_wal::batch::{POSITION_LIMIT, entry_id}; + use crate::object_store_wal::format::{FOOTER_ENTRY_LEN, decode_object}; const PREFIX: &str = "wal/datanodes/1/epochs/2"; const WAIT: Duration = Duration::from_secs(30); @@ -861,8 +1623,14 @@ mod tests { } } + /// Every append reaches the size limit, so it is persisted on its own. fn eager() -> ObjectStoreWalConfig { - config(Duration::from_millis(10), 1) + config(Duration::from_secs(3600), 1) + } + + /// Nothing is persisted until a test seals the open batch. + fn manual() -> ObjectStoreWalConfig { + config(Duration::from_secs(3600), u64::MAX) } async fn open( @@ -890,6 +1658,126 @@ mod tests { fn id(object_seq: u64, position: u64) -> EntryId { entry_id(object_seq, position) } + + fn entry(store: &ObjectStoreLogStore, region_id: RegionId, data: &str) -> Entry { + store + .entry(data.as_bytes().to_vec(), 0, region_id, &provider(region_id)) + .unwrap() + } + + async fn append( + store: &ObjectStoreLogStore, + region_id: RegionId, + data: &str, + ) -> Result { + store + .append_batch(vec![entry(store, region_id, data)]) + .await + } + + fn spawn_append_batch( + store: &Arc, + entries: Vec, + ) -> tokio::task::JoinHandle> { + let store = store.clone(); + tokio::spawn(async move { store.append_batch(entries).await }) + } + + /// Appends `count` single-entry batches, waiting for each to be admitted + /// before the next is sent, so they are admitted in order. + async fn spawn_appends( + store: &Arc, + region_id: RegionId, + count: usize, + ) -> Vec>> { + let admitted = *store.admitted_appends.borrow(); + let mut handles = Vec::with_capacity(count); + for index in 1..=count { + let entries = vec![entry(store, region_id, &format!("a{index}"))]; + handles.push(spawn_append_batch(store, entries)); + store + .wait_for_admitted_appends(admitted + index) + .await + .unwrap(); + } + handles + } + + async fn object_seqs(io: &dyn WalObjectIo) -> Vec { + io.list() + .await + .unwrap() + .into_iter() + .map(|object| object.object_seq) + .collect() + } + + fn unwrap_shared(error: &Error) -> &Error { + match error { + Error::ObjectStoreWal { source, .. } => source, + other => panic!("expected a shared error, actual {other:?}"), + } + } + + fn assert_stopped(error: &Error) { + assert!( + matches!(error, Error::ObjectStoreWalStopped { .. }), + "unexpected error: {error:?}" + ); + } + + /// Waits for the next create that `ParkedIo` parked. + async fn next_create( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) -> (u64, oneshot::Sender) { + timeout(WAIT, parked.recv()).await.unwrap().unwrap() + } + + /// Waits for `count` parked creates and returns their releases by + /// sequence. + async fn parked_creates( + parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + count: usize, + ) -> HashMap> { + let mut releases = HashMap::new(); + for _ in 0..count { + let (object_seq, release) = next_create(parked).await; + releases.insert(object_seq, release); + } + releases + } + + async fn open_over( + io: Arc, + config: &ObjectStoreWalConfig, + ) -> Arc { + ObjectStoreLogStore::open(io, config, PREFIX.to_string()) + .await + .unwrap() + } + + async fn open_parking_creates( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + ) -> ( + Arc, + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + let (io, parked) = ParkedIo::parking_creates(object_store); + let store = open_over(io.clone(), config).await; + // The start object of the open is created without parking. + io.creates_parked.store(true, Ordering::SeqCst); + (store, io, parked) + } + + /// Round-trips a command through the actor, so an assertion that + /// something did not happen runs after the actor handled every command + /// sent before. The open batch must be empty, as under the eager config. + async fn round_trip_actor(store: &ObjectStoreLogStore) { + store.seal_open_batch().await.unwrap(); + } + #[tokio::test] async fn test_store_rejects_invalid_config() { for config in [ @@ -1071,6 +1959,18 @@ mod tests { io.put_if_absent(object_seq, encoded.bytes).await.unwrap(); } + /// Puts an object of another writer with `epoch` under `object_seq`. A + /// store opened on an empty prefix writes epoch 1. + async fn put_foreign(object_store: &ObjectStore, object_seq: u64, epoch: u64) { + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let header = Header { + object_seq, + epoch, + prev: None, + }; + put_object_with_header(&io, header, &[]).await; + } + fn object_path(object_store: &ObjectStore, object_seq: u64) -> String { ObjectStoreIo::new(object_store.clone(), PREFIX) .unwrap() @@ -1359,10 +2259,51 @@ mod tests { HashMap::from([(region(1), id(5, 7)), (region(2), 8)]), recovered.durable_entry_ids ); - let store = open(object_store, &eager()).await; + let store = open(object_store.clone(), &eager()).await; assert_eq!(id(5, 7), latest(&store, region(1))); assert_eq!(8, latest(&store, region(2))); assert_eq!(0, latest(&store, region(3))); + + // The sequence resumes above the object the largest id names, where + // the start object of the store takes object 6, so the first new id + // of every region is greater than every old one. + let response = store + .append_batch(vec![ + entry(&store, region(1), "a"), + entry(&store, region(2), "b"), + ]) + .await + .unwrap(); + assert_eq!( + HashMap::from([(region(1), id(7, 1)), (region(2), id(7, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![2, 4, 6, 7], object_seqs(store.io.as_ref()).await); + assert_eq!( + expected_entries( + region(1), + &[ + (1, "e1"), + (id(5, 7), &format!("e{}", id(5, 7))), + (id(7, 1), "a") + ] + ), + read_entries(&store, region(1), 0).await + ); + store.stop().await.unwrap(); + + // The sequence continues after the start object of the restart. + let store = open(object_store, &eager()).await; + assert_eq!(id(7, 1), latest(&store, region(1))); + let response = append(&store, region(2), "b2").await.unwrap(); + assert_eq!( + HashMap::from([(region(2), id(9, 1))]), + response.last_entry_ids + ); + assert_eq!( + expected_entries(region(2), &[(8, "e8"), (id(7, 1), "b"), (id(9, 1), "b2")]), + read_entries(&store, region(2), 0).await + ); store.stop().await.unwrap(); } @@ -1401,7 +2342,14 @@ mod tests { let other = Provider::object_store_provider(region_id, "other/prefix".to_string()); let raft = Provider::raft_engine_provider(region_id.as_u64()); for foreign in [&other, &raft] { + let entry = Entry::Naive(NaiveEntry { + provider: foreign.clone(), + region_id, + entry_id: 0, + data: b"a1".to_vec(), + }); let errors = [ + store.append_batch(vec![entry]).await.unwrap_err(), store.latest_entry_id(foreign).unwrap_err(), store.read(foreign, 0, None).await.err().unwrap(), store.create_namespace(foreign).await.unwrap_err(), @@ -1435,12 +2383,14 @@ mod tests { timeout(WAIT, store.command_tx.closed()).await.unwrap(); store.stop().await.unwrap(); assert!(store.stopped.load(Ordering::Acquire)); + // Even an empty append fails once stop began. + assert_stopped(&store.append_batch(Vec::new()).await.unwrap_err()); } - #[tokio::test] - async fn test_store_stop_with_actor_gone() { + /// Builds a store whose commands the test receives instead of an actor. + fn store_without_actor() -> (ObjectStoreLogStore, mpsc::Receiver) { let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); - drop(command_rx); + let (append_tx, _) = mpsc::channel(APPEND_BUFFER); let store = ObjectStoreLogStore { prefix: PREFIX.to_string(), io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), @@ -1449,7 +2399,18 @@ mod tests { terminal_error: Arc::default(), stopped: Arc::new(AtomicBool::new(false)), command_tx, + append_tx, + admitted_appends: watch::channel(0).1, + creates_held: watch::channel(false).0, + creates_fail: Arc::default(), }; + (store, command_rx) + } + + #[tokio::test] + async fn test_store_stop_with_actor_gone() { + let (store, command_rx) = store_without_actor(); + drop(command_rx); store.stop().await.unwrap(); store.stop().await.unwrap(); } @@ -1485,6 +2446,7 @@ mod tests { let region_id = region(1); let provider = provider(region_id); let errors = [ + store.append_batch(Vec::new()).await.unwrap_err(), store.latest_entry_id(&provider).unwrap_err(), store.read(&provider, 0, None).await.err().unwrap(), store.create_namespace(&provider).await.unwrap_err(), @@ -1641,21 +2603,6 @@ mod tests { assert_invalid_object(&error, &object.path, "fewer bytes than the listed"); } - #[tokio::test] - async fn test_store_append_is_unsupported() { - let store = open(memory_store(), &eager()).await; - let error = store.append_batch(Vec::new()).await.unwrap_err(); - assert!(matches!( - error, - Error::UnsupportedObjectStoreWalOperation { .. } - )); - assert_eq!( - common_error::status_code::StatusCode::Unsupported, - error.status_code() - ); - store.stop().await.unwrap(); - } - async fn read_entries( store: &ObjectStoreLogStore, region_id: RegionId, @@ -1799,10 +2746,7 @@ mod tests { store.stop().await.unwrap(); let reopened = open(object_store, &eager()).await; assert_eq!(all, read_entries(&reopened, region(1), 0).await); - reopened - .obsolete(&p, region(1), EntryId::MAX) - .await - .unwrap(); + reopened.obsolete(&p, region(1), 20).await.unwrap(); assert_eq!( Vec::::new(), read_entries(&reopened, region(1), 0).await @@ -1909,6 +2853,1046 @@ mod tests { store.stop().await.unwrap(); } + #[tokio::test] + async fn test_store_assigns_ids_from_the_object_sequence() { + let store = open(memory_store(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + + // Appends admitted into one batch share its object; the regions take + // their own positions in admission order, and every append learns + // the last id of each of its regions. + let first = spawn_append_batch( + &store, + vec![ + entry(&store, region_one, "a1"), + entry(&store, region_two, "b1"), + entry(&store, region_one, "a2"), + ], + ); + store.wait_for_admitted_appends(1).await.unwrap(); + let second = spawn_append_batch(&store, vec![entry(&store, region_two, "b2")]); + store.wait_for_admitted_appends(2).await.unwrap(); + assert!(!first.is_finished()); + assert_eq!(0, latest(&store, region_one)); + store.seal_open_batch().await.unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 2)), (region_two, id(1, 1))]), + first.await.unwrap().unwrap().last_entry_ids + ); + assert_eq!( + HashMap::from([(region_two, id(1, 2))]), + second.await.unwrap().unwrap().last_entry_ids + ); + + // The next object starts every region at position one again. + let third = spawn_append_batch( + &store, + vec![ + entry(&store, region_two, "b3"), + entry(&store, region_one, "a3"), + ], + ); + store.wait_for_admitted_appends(3).await.unwrap(); + store.seal_open_batch().await.unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 1)), (region_two, id(2, 1))]), + third.await.unwrap().unwrap().last_entry_ids + ); + assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + assert_eq!( + expected_entries( + region_one, + &[(id(1, 1), "a1"), (id(1, 2), "a2"), (id(2, 1), "a3")] + ), + read_entries(&store, region_one, 1).await + ); + assert_eq!( + expected_entries(region_two, &[(id(1, 2), "b2"), (id(2, 1), "b3")]), + read_entries(&store, region_two, id(1, 2)).await + ); + assert_eq!(id(2, 1), latest(&store, region_one)); + assert_eq!(id(2, 1), latest(&store, region_two)); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_seals_when_a_region_exhausts_its_positions() { + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + let entries_of = |count: u64| { + (0..count) + .map(|_| entry(&store, region_id, "")) + .collect::>() + }; + + // The open batch holds every position of the region. + let full = spawn_append_batch(&store, entries_of(POSITION_LIMIT - 1)); + store.wait_for_admitted_appends(1).await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + + // The next entry of the region seals the batch and opens the next + // object; another region would still have fit. + let next = spawn_append_batch(&store, vec![entry(&store, region_id, "next")]); + store.wait_for_admitted_appends(2).await.unwrap(); + let response = timeout(WAIT, full).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, POSITION_LIMIT - 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + + // An append that alone runs past the range fits no object: it seals + // the open batch like any append the batch cannot take, then it is + // refused without poisoning the store. + let error = store + .append_batch(entries_of(POSITION_LIMIT)) + .await + .unwrap_err(); + assert!( + matches!(error, Error::WalEntryPositionExhausted { .. }), + "unexpected error: {error:?}" + ); + let response = timeout(WAIT, next).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + let after = spawn_append_batch(&store, vec![entry(&store, region_id, "after")]); + store.wait_for_admitted_appends(3).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, after).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(3, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1, 2, 3], object_seqs(store.io.as_ref()).await); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "next"), (id(3, 1), "after")]), + read_entries(&store, region_id, id(2, 1)).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_flushes_at_the_minimum_interval() { + // An append is acknowledged once its object is durable, so a completed + // append proves the tick after it sealed the batch. The appends are + // spaced far apart, so every one of them lands in its own tick and + // the ticks in between find an empty batch. + let interval = MIN_FLUSH_INTERVAL; + let store = open(memory_store(), &config(interval, u64::MAX)).await; + let data = ["a1", "a2", "a3"]; + for (index, data) in data.iter().enumerate() { + tokio::time::sleep(interval * 5).await; + timeout(WAIT, append(&store, region(1), data)) + .await + .unwrap() + .unwrap(); + let expected_seqs = (0..=index as u64 + 1).collect::>(); + assert_eq!(expected_seqs, object_seqs(store.io.as_ref()).await); + } + + // Ticks with an empty open batch do not create objects. + tokio::time::sleep(interval * 5).await; + let seqs = object_seqs(store.io.as_ref()).await; + assert_eq!(vec![0, 1, 2, 3], seqs); + // Object 0 is the start object of the store. + for object_seq in seqs.into_iter().skip(1) { + let bytes = store.io.get(object_seq).await.unwrap(); + let decoded = decode_object(&bytes).unwrap(); + assert_eq!(1, decoded.records.len(), "object {object_seq} is empty"); + } + assert_eq!( + expected_entries( + region(1), + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(3, 1), "a3")] + ), + read_entries(&store, region(1), 1).await + ); + } + + #[tokio::test] + async fn test_store_rejects_entries_of_another_region() { + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + + let entry = Entry::Naive(NaiveEntry { + provider: provider(region_id), + region_id: region(2), + entry_id: 0, + data: Vec::new(), + }); + let error = store.append_batch(vec![entry]).await.unwrap_err(); + assert!( + matches!(error, Error::MismatchedWalRegion { .. }), + "unexpected error: {error:?}" + ); + let incomplete = Entry::MultiplePart(MultiplePartEntry { + provider: provider(region_id), + region_id, + entry_id: 0, + headers: vec![MultiplePartHeader::First], + parts: vec![b"a1".to_vec()], + }); + let error = store.append_batch(vec![incomplete]).await.unwrap_err(); + assert!( + matches!(error, Error::IncompleteWalEntry { region_id: actual, .. } if actual == region_id), + "unexpected error: {error:?}" + ); + assert_eq!( + common_error::status_code::StatusCode::InvalidArguments, + error.status_code() + ); + assert!( + store + .append_batch(Vec::new()) + .await + .unwrap() + .last_entry_ids + .is_empty() + ); + // Nothing was admitted. + store.seal_open_batch().await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_pipelines_creates_and_acknowledges_in_sequence_order() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 2).await; + + // At most the limit of creates run at a time; the rest wait for a slot. + let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + assert_eq!( + (1..=MAX_IN_FLIGHT_CREATES as u64).collect::>(), + releases.keys().copied().collect::>() + ); + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + assert!(appends.iter().all(|append| !append.is_finished())); + + // Objects 3 and 2 become durable before object 1 and free a slot + // each, but nothing is acknowledged ahead of object 1. + for (released, expected_next) in [(3, 5), (2, 6)] { + releases.remove(&released).unwrap().send(true).unwrap(); + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(expected_next, object_seq); + releases.insert(object_seq, release); + } + round_trip_actor(&store).await; + assert_eq!(vec![0, 2, 3], object_seqs(io.as_ref()).await); + assert!(appends.iter().all(|append| !append.is_finished())); + assert_eq!(0, latest(&store, region_id)); + + // Object 1 releases the acknowledgements of objects 1 to 3. + releases.remove(&1).unwrap().send(true).unwrap(); + let mut appends = appends.into_iter(); + for object_seq in 1..4 { + let response = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!(id(3, 1), latest(&store, region_id)); + + // Object 6 before object 5: the append of object 6 waits for it. + releases.remove(&4).unwrap().send(true).unwrap(); + releases.remove(&6).unwrap().send(true).unwrap(); + let fourth = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(4, 1))]), + fourth.last_entry_ids + ); + let appends = appends.collect::>(); + round_trip_actor(&store).await; + assert!(appends.iter().all(|append| !append.is_finished())); + releases.remove(&5).unwrap().send(true).unwrap(); + for (append, object_seq) in appends.into_iter().zip(5..) { + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!( + MAX_IN_FLIGHT_CREATES, + io.max_in_flight.load(Ordering::SeqCst) + ); + assert_eq!(vec![0, 1, 2, 3, 4, 5, 6], object_seqs(io.as_ref()).await); + // Every batch extends the batch sealed before it, whichever was + // durable first. + for object_seq in 1..=6 { + let header = decode_header(&io.get(object_seq).await.unwrap()).unwrap(); + assert_eq!( + Some(object_seq - 1), + header.prev.map(|link| link.object_seq) + ); + } + assert_eq!( + expected_entries( + region_id, + &[ + (id(1, 1), "a1"), + (id(2, 1), "a2"), + (id(3, 1), "a3"), + (id(4, 1), "a4"), + (id(5, 1), "a5"), + (id(6, 1), "a6") + ] + ), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_transient_failure_fails_every_batch_that_is_not_indexed() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Every slot is taken; the last batch waits for one. + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; + let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + + // Object 1 fails while objects 2 to 4 are in flight and object 5 waits + // for a slot: every batch fails at once. + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } + assert_eq!(0, latest(&store, region_id)); + + // The retried entries take the sequences after the last sealed one. + // The creates of the failed batches still hold their slots, so only + // one retry starts until they complete; what they store is off the + // chain. + let retries = spawn_appends(&store, region_id, 2).await; + let (object_seq, first) = next_create(&mut parked).await; + assert_eq!(6, object_seq); + round_trip_actor(&store).await; + assert!(parked.try_recv().is_err()); + for release in releases.into_values() { + release.send(true).unwrap(); + } + let (object_seq, second) = next_create(&mut parked).await; + assert_eq!(7, object_seq); + first.send(true).unwrap(); + second.send(true).unwrap(); + for (retry, object_seq) in retries.into_iter().zip(6..) { + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!(vec![0, 2, 3, 4, 6, 7], object_seqs(io.as_ref()).await); + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "a1"), (id(7, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_transient_failure_before_a_created_object_keeps_it_off_the_chain() { + let object_store = memory_store(); + let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_creates(&mut parked, 2).await; + + // Object 2 is created while object 1 failed: both batches fail, and + // the store keeps serving. + releases.remove(&2).unwrap().send(true).unwrap(); + timeout(WAIT, async { + while object_seqs(io.as_ref()).await != [0, 2] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + } + assert_eq!(0, latest(&store, region_id)); + + // The retry extends the start object, not the failed batches. + let retry = spawn_appends(&store, region_id, 1).await; + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(3, object_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, retry.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(3, 1))]), + response.last_entry_ids + ); + let bytes = io.get(3).await.unwrap(); + let header = decode_header(&bytes).unwrap(); + assert_eq!( + Some(ChainLink { + object_seq: 0, + epoch: 1, + }), + header.prev + ); + store.stop().await.unwrap(); + + // After a restart only the acknowledged entry replays: object 2 is + // off the chain although it exists. + let store = open(object_store, &eager()).await; + assert_eq!(id(3, 1), latest(&store, region_id)); + assert_eq!( + expected_entries(region_id, &[(id(3, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_permanent_failure_before_a_durable_object_poisons() { + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let foreign = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + // Another writer of the same epoch took sequence 1. + put_foreign(&object_store, 1, 1).await; + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_creates(&mut parked, 2).await; + + // Object 2 is created; object 1 conflicts with the foreign object. + releases.remove(&2).unwrap().send(true).unwrap(); + timeout(WAIT, async { + while object_seqs(&foreign).await != [0, 1, 2] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + round_trip_actor(&store).await; + assert!(appends.iter().all(|append| !append.is_finished())); + releases.remove(&1).unwrap().send(true).unwrap(); + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + let error = store.latest_entry_id(&provider(region_id)).unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_conflict_with_an_earlier_epoch_fails_the_batch_without_poisoning() { + let object_store = memory_store(); + let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + // A late object of an earlier epoch lands under sequence 1. + put_foreign(&object_store, 1, 0).await; + let append = spawn_appends(&store, region_id, 1).await; + next_create(&mut parked).await.1.send(true).unwrap(); + let error = timeout(WAIT, append.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!( + matches!( + unwrap_shared(&error), + Error::StaleWalObject { + existing_epoch: 0, + epoch: 1, + .. + } + ), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + + // The store keeps serving: the retry takes the next sequence. + let retry = spawn_appends(&store, region_id, 1).await; + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(2, object_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, retry.into_iter().next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + store.stop().await.unwrap(); + + // The stale object is off the chain after a restart. + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_obsolete_raises_the_sequence_floor() { + let store = open(memory_store(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + let mut admitted = 0; + let mut spawn_append = |region_id, data: &str| { + admitted += 1; + ( + spawn_append_batch(&store, vec![entry(&store, region_id, data)]), + admitted, + ) + }; + let (first, count) = spawn_append(region_one, "a1"); + store.wait_for_admitted_appends(count).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, first).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 1))]), + response.last_entry_ids + ); + let (second, count) = spawn_append(region_one, "a2"); + store.wait_for_admitted_appends(count).await.unwrap(); + + // A durable watermark names an object below the next sequence and + // changes nothing: the open batch goes on under sequence 2. + store + .obsolete(&provider(region_one), region_one, id(1, 1)) + .await + .unwrap(); + let (third, count) = spawn_append(region_one, "a3"); + store.wait_for_admitted_appends(count).await.unwrap(); + + // A watermark that names a later object, as one inherited from + // another prefix does, cannot move the sequence while the open batch + // handed out ids under it: neither the sequence nor the watermark + // moves. + let error = store + .obsolete(&provider(region_two), region_two, id(3, 7)) + .await + .unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectSequenceUnsettled { object_seq: 2, .. } + ), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert!( + store + .obsolete_entry_ids + .lock() + .unwrap() + .get(®ion_two) + .is_none() + ); + + // Once the batch is durable the sequence moves above that object, so + // the region's next id is greater than the watermark. + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, second).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 1))]), + response.last_entry_ids + ); + let response = timeout(WAIT, third).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_one, id(2, 2))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + store + .obsolete(&provider(region_two), region_two, id(3, 7)) + .await + .unwrap(); + assert_eq!( + Some(&id(3, 7)), + store.obsolete_entry_ids.lock().unwrap().get(®ion_two) + ); + let (fourth, count) = spawn_append(region_two, "b1"); + store.wait_for_admitted_appends(count).await.unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, fourth).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_two, id(4, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1, 2, 4], object_seqs(store.io.as_ref()).await); + // The watermark hides nothing of this prefix; the region has no + // entry at or below it. + assert_eq!( + expected_entries(region_two, &[(id(4, 1), "b1")]), + read_entries(&store, region_two, 1).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_does_not_reuse_the_sequence_of_a_failed_create() { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &manual()).await; + let region_one = region(1); + let region_two = region(2); + + // Two appends share object 1, which is written but reported as + // failed: it exists and is not indexed. + let first = spawn_append_batch(&store, vec![entry(&store, region_one, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let second = spawn_append_batch(&store, vec![entry(&store, region_two, "b1")]); + store.wait_for_admitted_appends(2).await.unwrap(); + io.fail_after_next_put.store(true, Ordering::Relaxed); + store.seal_open_batch().await.unwrap_err(); + for append in [first, second] { + timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + } + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(0, latest(&store, region_one)); + + // Retried apart, the entries land in later objects rather than + // conflicting with object 1. + for (region_id, data, object_seq) in [(region_one, "a1", 2), (region_two, "b1", 3)] { + let retry = spawn_append_batch(&store, vec![entry(&store, region_id, data)]); + store + .wait_for_admitted_appends(object_seq as usize + 1) + .await + .unwrap(); + store.seal_open_batch().await.unwrap(); + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + store.stop().await.unwrap(); + + // Object 2 extends the start object, so recovery leaves object 1 off + // the chain: only the retries replay. + let store = open_over(io, &manual()).await; + assert_eq!( + expected_entries(region_one, &[(id(2, 1), "a1")]), + read_entries(&store, region_one, 0).await + ); + assert_eq!( + expected_entries(region_two, &[(id(3, 1), "b1")]), + read_entries(&store, region_two, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_floor_applies_while_a_create_is_in_flight() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_one = region(1); + let region_two = region(2); + let pending = spawn_append_batch(&store, vec![entry(&store, region_one, "a1")]); + let (_, release) = next_create(&mut parked).await; + + // Object 1 is in flight and the open batch is empty: the floor moves + // the next sequence, which a failure of object 1 does not move back. + store + .obsolete(&provider(region_two), region_two, id(5, 1)) + .await + .unwrap(); + release.send(false).unwrap(); + timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err(); + + let write = spawn_append_batch(&store, vec![entry(&store, region_two, "b1")]); + next_create(&mut parked).await.1.send(true).unwrap(); + let response = timeout(WAIT, write).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_two, id(6, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 6], object_seqs(io.as_ref()).await); + } + + #[tokio::test] + async fn test_store_obsolete_queued_behind_stop_records_the_watermark() { + let (store, mut command_rx) = store_without_actor(); + let region_id = region(1); + let provider_one = provider(region_id); + + // The actor exits with the command still queued and never answers + // it: nothing is assigned an id any more, so the watermark alone holds. + let (result, ()) = tokio::join!(store.obsolete(&provider_one, region_id, 1), async { + let command = timeout(WAIT, command_rx.recv()).await.unwrap(); + assert!(matches!( + command, + Some(Command::Obsolete { entry_id: 1, .. }) + )); + }); + result.unwrap(); + assert_eq!( + Some(&1), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + + // The same holds for a call that finds the actor gone. + drop(command_rx); + store + .obsolete(&provider(region(2)), region(2), id(1, 1)) + .await + .unwrap(); + assert_eq!( + Some(&id(1, 1)), + store.obsolete_entry_ids.lock().unwrap().get(®ion(2)) + ); + } + + #[tokio::test] + async fn test_store_hooks_hold_and_fail_creates() { + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + store.hold_creates(); + let held = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + round_trip_actor(&store).await; + assert!(!held.is_finished()); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + store.release_creates(); + timeout(WAIT, held).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + + store.fail_creates(); + let error = append(&store, region_id, "a2").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + + store.begin_stop(); + assert_stopped(&append(&store, region_id, "a3").await.unwrap_err()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_create_held_when_the_store_is_dropped_never_runs() { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &eager()).await; + store.hold_creates(); + let append = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + append.abort(); + let _ = append.await; + // A command sender outlives the store, so the actor keeps running and + // the parked create completes on the closed hold channel instead of + // being dropped with the actor. + let command_tx = store.command_tx.clone(); + let terminal_error = store.terminal_error.clone(); + drop(store); + timeout(WAIT, async { + while terminal(&terminal_error).is_none() && object_seqs(io.as_ref()).await == [0] { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + let error = terminal(&terminal_error).unwrap(); + assert!( + matches!(error.as_ref(), Error::ObjectStoreWalStopped { .. }), + "unexpected error: {error:?}" + ); + drop(command_tx); + } + + #[tokio::test] + async fn test_store_rollback_and_poison_drop_the_open_batch() { + let object_store = memory_store(); + let (store, io, mut parked) = open_parking_creates(object_store.clone(), &manual()).await; + let region_id = region(1); + let spawn_seal = || { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + + // Object 1 is in flight and a2 waits in the open batch under + // sequence 2 when the create fails: both fail. + let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let seal = spawn_seal(); + let (_, release) = next_create(&mut parked).await; + let open = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); + store.wait_for_admitted_appends(2).await.unwrap(); + release.send(false).unwrap(); + for append in [sealed, open] { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + } + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + + // The open batch had not taken its sequence, so the next entry takes + // the first id of sequence 2 again, alone. + let retry = spawn_append_batch(&store, vec![entry(&store, region_id, "b1")]); + store.wait_for_admitted_appends(3).await.unwrap(); + let seal = spawn_seal(); + next_create(&mut parked).await.1.send(true).unwrap(); + timeout(WAIT, seal).await.unwrap().unwrap().unwrap(); + let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "b1")]), + read_entries(&store, region_id, 0).await + ); + + // Object 3 conflicts while c2 waits in the open batch: both fail. + let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "c1")]); + store.wait_for_admitted_appends(4).await.unwrap(); + let seal = spawn_seal(); + let (_, release) = next_create(&mut parked).await; + let open = spawn_append_batch(&store, vec![entry(&store, region_id, "c2")]); + store.wait_for_admitted_appends(5).await.unwrap(); + put_foreign(&object_store, 3, 1).await; + release.send(true).unwrap(); + for append in [sealed, open] { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + assert_eq!(vec![0, 2, 3], object_seqs(io.as_ref()).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_poisons_once_the_last_sequence_is_taken() { + let object_store = memory_store(); + let last = OBJECT_SEQ_LIMIT - 1; + // The start object of the store takes the sequence before the last. + put_object(&object_store, last - 2, region(1), &[id(last - 2, 1)]).await; + let store = open(object_store, &eager()).await; + let next = entry(&store, region(1), "next"); + + // The batch at the last sequence is created and acknowledged, and the + // store is poisoned as soon as it takes that sequence. + let response = append(&store, region(1), "last").await.unwrap(); + assert_eq!( + HashMap::from([(region(1), id(last, 1))]), + response.last_entry_ids + ); + for error in [ + store.latest_entry_id(&provider(region(1))).unwrap_err(), + store.append_batch(vec![next]).await.unwrap_err(), + ] { + assert!( + matches!( + unwrap_shared(&error), + Error::WalObjectSequenceExhausted { last_object_seq, .. } if *last_object_seq == last + ), + "unexpected error: {error:?}" + ); + } + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_holds_appends_back_while_the_sealed_batches_are_full() { + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + let admitted = || *store.admitted_appends.borrow(); + + // Every create is parked: once no two more sealed batches fit, + // further appends wait in the channel and are not admitted. + let mut appends = spawn_appends(&store, region_id, MAX_SEALED_BATCHES - 1).await; + for index in 0..3 { + let entries = vec![entry(&store, region_id, &format!("w{index}"))]; + appends.push(spawn_append_batch(&store, entries)); + } + let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + round_trip_actor(&store).await; + assert_eq!(MAX_SEALED_BATCHES - 1, admitted()); + + // A durable object frees a place, and one more append is admitted. + releases.remove(&1).unwrap().send(true).unwrap(); + timeout(WAIT, store.wait_for_admitted_appends(MAX_SEALED_BATCHES)) + .await + .unwrap() + .unwrap(); + round_trip_actor(&store).await; + assert_eq!(MAX_SEALED_BATCHES, admitted()); + + // Stop is handled while appends are held back: the batches whose create + // has not started and the appends still waiting learn of the stop. + let stop = begin_spawned_stop(&store).await; + for release in releases.into_values() { + release.send(true).unwrap(); + } + next_create(&mut parked).await.1.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + let mut acknowledged = 0; + for append in appends { + match timeout(WAIT, append).await.unwrap().unwrap() { + Ok(_) => acknowledged += 1, + Err(error) => assert_stopped(&error), + } + } + assert_eq!(MAX_IN_FLIGHT_CREATES + 1, acknowledged); + } + + /// Starts `stop` and waits until it has set the stopped flag. + async fn begin_spawned_stop( + store: &Arc, + ) -> tokio::task::JoinHandle> { + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + while !store.stopped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + stop + } + + #[tokio::test] + async fn test_store_stop_during_failed_flush_reports_stopped() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let pending = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); + let (_, release) = next_create(&mut parked).await; + let stop = begin_spawned_stop(&store).await; + assert!(!stop.is_finished()); + + // The create fails after stop began: its waiter learns of the stop. + release.send(false).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_stopped(&timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_stopped(&append(&store, region(1), "a2").await.unwrap_err()); + } + + #[tokio::test] + async fn test_store_stop_with_several_creates_in_flight() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Four creates are in flight, the fifth batch waits for a slot. + let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; + let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + assert!(parked.try_recv().is_err()); + + let stop = begin_spawned_stop(&store).await; + assert!(!stop.is_finished()); + // Nothing is assigned an id after stop began, so a watermark needs no + // floor even while creates are in flight. + store + .obsolete(&provider(region(2)), region(2), id(9, 1)) + .await + .unwrap(); + assert_eq!( + Some(&id(9, 1)), + store.obsolete_entry_ids.lock().unwrap().get(®ion(2)) + ); + + // The creates in flight run to completion in any order and are + // acknowledged; the batch that never started learns of the stop. + for object_seq in [3, 1, 4, 2] { + releases.remove(&object_seq).unwrap().send(true).unwrap(); + } + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + let mut appends = appends.into_iter(); + for object_seq in 1..=MAX_IN_FLIGHT_CREATES as u64 { + let response = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq, 1))]), + response.last_entry_ids + ); + } + let error = timeout(WAIT, appends.next().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert_stopped(&error); + // No create was started after stop began. + assert!(parked.try_recv().is_err()); + assert_eq!(vec![0, 1, 2, 3, 4], object_seqs(io.as_ref()).await); + assert_eq!(id(4, 1), latest(&store, region_id)); + timeout(WAIT, store.command_tx.closed()).await.unwrap(); + } + + #[tokio::test] + async fn test_store_stop_during_conflicting_flush_reports_stopped_and_poisons() { + let object_store = memory_store(); + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let region_id = region(1); + let pending = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + let (_, release) = next_create(&mut parked).await; + let stop = begin_spawned_stop(&store).await; + // Another writer of the same epoch takes the sequence before the + // create runs. + put_foreign(&object_store, 1, 1).await; + + release.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_stopped(&timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err()); + for error in [ + store.latest_entry_id(&provider(region_id)).unwrap_err(), + store + .read(&provider(region_id), 1, None) + .await + .err() + .unwrap(), + ] { + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + } + } + + fn injected_failure(operation: &'static str, path: String) -> Result { + Err( + object_store::Error::new(object_store::ErrorKind::Unexpected, "injected failure") + .set_temporary(), + ) + .context(WalObjectStoreSnafu { operation, path }) + } + fn chain_header(object_seq: u64, epoch: u64, prev: Option<(u64, u64)>) -> Header { Header { object_seq, @@ -2248,15 +4232,6 @@ mod tests { assert_eq!(vec![0], object_seqs(&io).await); } - async fn object_seqs(io: &ObjectStoreIo) -> Vec { - io.list() - .await - .unwrap() - .into_iter() - .map(|object| object.object_seq) - .collect() - } - #[tokio::test] async fn test_store_open_fails_when_its_start_object_lands_with_an_unknown_outcome() { let object_store = memory_store(); @@ -2368,9 +4343,16 @@ mod tests { } } + /// Object access whose whole-object reads, or conditional creates, park + /// until the test releases them, counting how many are in flight. A + /// release of false fails the operation with a transient error before it + /// reaches the object store. struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, + park_creates: bool, + /// Whether creates park yet; set once the store opened. + creates_parked: AtomicBool, in_flight: AtomicUsize, max_in_flight: AtomicUsize, } @@ -2381,43 +4363,68 @@ mod tests { ) -> ( Arc, mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + Self::new(object_store, false) + } + + fn parking_creates( + object_store: ObjectStore, + ) -> ( + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + ) { + Self::new(object_store, true) + } + + fn new( + object_store: ObjectStore, + park_creates: bool, + ) -> ( + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, ) { let (parked, parked_rx) = mpsc::unbounded_channel(); ( Arc::new(Self { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), parked, + park_creates, + creates_parked: AtomicBool::new(false), in_flight: AtomicUsize::new(0), max_in_flight: AtomicUsize::new(0), }), parked_rx, ) } + + /// Parks until the test releases `object_seq` and returns whether the + /// operation proceeds. + async fn park(&self, object_seq: u64) -> bool { + let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); + let (release, released) = oneshot::channel(); + self.parked.send((object_seq, release)).unwrap(); + let proceed = released.await.unwrap(); + self.in_flight.fetch_sub(1, Ordering::SeqCst); + proceed + } } #[async_trait::async_trait] impl WalObjectIo for ParkedIo { async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + if self.park_creates + && self.creates_parked.load(Ordering::SeqCst) + && !self.park(object_seq).await + { + return injected_failure("write", self.object_path(object_seq)); + } self.inner.put_if_absent(object_seq, content).await } async fn get(&self, object_seq: u64) -> Result { - let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; - self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst); - let (release, released) = oneshot::channel(); - self.parked.send((object_seq, release)).unwrap(); - let success = released.await.unwrap(); - self.in_flight.fetch_sub(1, Ordering::SeqCst); - if !success { - return Err(object_store::Error::new( - object_store::ErrorKind::Unexpected, - "injected failure", - ) - .set_temporary()) - .context(WalObjectStoreSnafu { - operation: "read", - path: self.object_path(object_seq), - }); + if !self.park_creates && !self.park(object_seq).await { + return injected_failure("read", self.object_path(object_seq)); } self.inner.get(object_seq).await } @@ -2437,10 +4444,13 @@ mod tests { type RangeReads = Arc>>; - /// Object access that records every range read as (sequence, offset, length). + /// Object access that records every range read as (sequence, offset, + /// length) and, on request, reports the next conditional create as failed + /// after it wrote the object. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, + fail_after_next_put: AtomicBool, } impl RecordingIo { @@ -2449,6 +4459,7 @@ mod tests { let io = Self { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), reads: reads.clone(), + fail_after_next_put: AtomicBool::new(false), }; (Arc::new(io), reads) } @@ -2457,7 +4468,11 @@ mod tests { #[async_trait::async_trait] impl WalObjectIo for RecordingIo { async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { - self.inner.put_if_absent(object_seq, content).await + let result = self.inner.put_if_absent(object_seq, content).await?; + if self.fail_after_next_put.swap(false, Ordering::Relaxed) { + return injected_failure("write", self.object_path(object_seq)); + } + Ok(result) } async fn get(&self, object_seq: u64) -> Result { From 5ffd01a70a0ae8aa35017e8512cc05f315304b16 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Sun, 27 Sep 2026 02:15:47 +0000 Subject: [PATCH 041/133] fix: preserve primary key order when syncing columns (#9189) * fix: preserve primary key order when syncing columns Signed-off-by: WenyXu * fix: reject invalid SyncColumns metadata Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../resolve_column_metadata.rs | 8 +- src/common/meta/src/reconciliation/utils.rs | 187 +++++++++++++++++- src/mito2/src/engine/alter_test.rs | 129 +++++++++++- 3 files changed, 318 insertions(+), 6 deletions(-) diff --git a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs index 27950f0fd5d..8defc6bd0f2 100644 --- a/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs +++ b/src/common/meta/src/reconciliation/reconcile_table/resolve_column_metadata.rs @@ -25,8 +25,8 @@ use crate::reconciliation::reconcile_table::reconcile_regions::ReconcileRegions; use crate::reconciliation::reconcile_table::update_table_info::UpdateTableInfo; use crate::reconciliation::reconcile_table::{ReconcileTableContext, State, TableMetadataState}; use crate::reconciliation::utils::{ - ResolveColumnMetadataResult, build_column_metadata_from_table_info, - check_column_metadatas_consistent, resolve_column_metadatas_with_latest, + ResolveColumnMetadataResult, build_reconciliation_column_metadata, + check_column_metadatas_consistent, reorder_tag_columns, resolve_column_metadatas_with_latest, resolve_column_metadatas_with_metasrv, }; @@ -92,6 +92,8 @@ impl State for ResolveColumnMetadata { ctx.persistent_ctx.table_info_value = Some(table_info_value); if let Some(column_metadatas) = check_column_metadatas_consistent(&self.region_metadata) { + let column_metadatas = + reorder_tag_columns(&column_metadatas, &self.region_metadata[0].primary_key)?; // Safety: fetched in the above. let table_info_value = ctx.persistent_ctx.table_info_value.clone().unwrap(); info!( @@ -125,7 +127,7 @@ impl State for ResolveColumnMetadata { .table_info .name_to_ids() .context(MissingColumnIdsSnafu)?; - let column_metadata = build_column_metadata_from_table_info( + let column_metadata = build_reconciliation_column_metadata( table_info_value.table_info.meta.schema.column_schemas(), &table_info_value.table_info.meta.primary_key_indices, &name_to_ids, diff --git a/src/common/meta/src/reconciliation/utils.rs b/src/common/meta/src/reconciliation/utils.rs index 6ddc084596b..326516ecd80 100644 --- a/src/common/meta/src/reconciliation/utils.rs +++ b/src/common/meta/src/reconciliation/utils.rs @@ -85,6 +85,70 @@ pub(crate) fn check_column_metadatas_consistent( Some(region_metadatas[0].column_metadatas.clone()) } +/// Returns columns with tag values in primary-key order while preserving the +/// positions of non-tag columns. +pub(crate) fn reorder_tag_columns( + column_metadatas: &[ColumnMetadata], + primary_key: &[u32], +) -> Result> { + let tag_count = column_metadatas + .iter() + .filter(|column| column.semantic_type == SemanticType::Tag) + .count(); + ensure!( + primary_key.len() == tag_count, + UnexpectedSnafu { + err_msg: format!( + "Number of primary key columns {} does not match tag columns {}", + primary_key.len(), + tag_count, + ), + } + ); + + let mut primary_key_ids = HashSet::with_capacity(primary_key.len()); + let mut tags = Vec::with_capacity(primary_key.len()); + for column_id in primary_key { + let column = column_metadatas + .iter() + .find(|column| column.column_id == *column_id) + .with_context(|| UnexpectedSnafu { + err_msg: format!( + "Primary key column {} not found in column metadata", + column_id + ), + })?; + ensure!( + column.semantic_type == SemanticType::Tag, + UnexpectedSnafu { + err_msg: format!("Primary key column {} is not a tag", column_id), + } + ); + ensure!( + primary_key_ids.insert(column_id), + UnexpectedSnafu { + err_msg: format!("Primary key column {} is duplicated", column_id), + } + ); + tags.push(column); + } + + let mut tags = tags.into_iter(); + + column_metadatas + .iter() + .map(|column| { + if column.semantic_type == SemanticType::Tag { + tags.next().cloned().with_context(|| UnexpectedSnafu { + err_msg: "Primary key has fewer columns than tags".to_string(), + }) + } else { + Ok(column.clone()) + } + }) + .collect() +} + /// Resolves column metadata inconsistencies among the given region metadatas /// by using the column metadata from the metasrv as the source of truth. /// @@ -157,8 +221,13 @@ pub(crate) fn resolve_column_metadatas_with_latest( } } - // TODO(weny): verify the new column metadatas are acceptable for regions. - Ok((latest_region_metadata.column_metadatas.clone(), region_ids)) + Ok(( + reorder_tag_columns( + &latest_region_metadata.column_metadatas, + &latest_region_metadata.primary_key, + )?, + region_ids, + )) } /// Constructs a vector of [`ColumnMetadata`] from the provided table information. @@ -206,6 +275,22 @@ pub(crate) fn build_column_metadata_from_table_info( .collect::>>() } +/// Builds SyncColumns metadata with tags in the TableInfo primary-key order. +/// Non-tag columns retain their schema positions. +pub(crate) fn build_reconciliation_column_metadata( + column_schemas: &[ColumnSchema], + primary_key_indexes: &[usize], + name_to_ids: &HashMap, +) -> Result> { + let column_metadatas = + build_column_metadata_from_table_info(column_schemas, primary_key_indexes, name_to_ids)?; + let primary_key = primary_key_indexes + .iter() + .map(|index| column_metadatas[*index].column_id) + .collect::>(); + reorder_tag_columns(&column_metadatas, &primary_key) +} + /// Checks whether the schema invariants hold between the existing and new column metadata. /// /// Invariants: @@ -1180,6 +1265,104 @@ mod tests { assert_matches!(err, Error::MissingColumnInColumnMetadata { .. }); } + #[test] + fn test_reconciliation_columns_follow_primary_key_order() { + let columns = vec![ + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_b", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 1, + }, + ColumnMetadata { + column_schema: ColumnSchema::new("field", ConcreteDataType::int32_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 2, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_a", + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 3, + }, + ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + semantic_type: SemanticType::Timestamp, + column_id: 4, + }, + ]; + let schemas = columns + .iter() + .map(|column| column.column_schema.clone()) + .collect::>(); + let ids = columns + .iter() + .map(|column| (column.column_schema.name.clone(), column.column_id)) + .collect(); + let expected = vec![3, 2, 1, 4]; + assert_eq!( + build_reconciliation_column_metadata(&schemas, &[2, 0], &ids) + .unwrap() + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + + let mut metadata = build_region_metadata(RegionId::new(1024, 0), &columns); + metadata.primary_key = vec![3, 1]; + metadata.schema_version = 2; + + assert_eq!( + check_column_metadatas_consistent(&[metadata.clone()]).unwrap(), + columns + ); + + assert_eq!( + resolve_column_metadatas_with_latest(&[metadata]) + .unwrap() + .0 + .iter() + .map(|column| column.column_id) + .collect::>(), + expected + ); + } + + #[test] + fn test_reorder_tag_columns_rejects_invalid_primary_key() { + let columns = new_test_column_metadatas(); + + let err = reorder_tag_columns(&columns, &[999]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("not found in column metadata")); + + let err = reorder_tag_columns(&columns, &[]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("does not match tag columns")); + + let err = reorder_tag_columns(&columns, &[2]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is not a tag")); + + let mut duplicate_tags = columns.clone(); + duplicate_tags[2].semantic_type = SemanticType::Tag; + let err = reorder_tag_columns(&duplicate_tags, &[0, 0]).unwrap_err(); + assert_matches!(err, Error::Unexpected { .. }); + assert!(err.to_string().contains("is duplicated")); + } + #[test] fn test_check_column_metadatas_consistent() { let column_metadatas = new_test_column_metadatas(); diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index fd74c4b498a..cf039b52278 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -37,7 +37,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::provider::Provider; -use store_api::metadata::ColumnMetadata; +use store_api::metadata::{ColumnMetadata, RegionMetadata}; use store_api::metric_engine_consts::{PRIMARY_KEY_ENCODING, TABLE_COLUMN_METADATA_EXTENSION_KEY}; use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ @@ -193,6 +193,133 @@ fn assert_column_metadatas(column_name: &[(&str, ColumnId)], column_metadatas: & } } +#[tokio::test] +async fn test_sync_columns_preserves_primary_key_order() { + test_sync_columns_preserves_primary_key_order_with_format(false).await; + test_sync_columns_preserves_primary_key_order_with_format(true).await; +} + +async fn assert_sync_columns_readable( + engine: &MitoEngine, + region_id: RegionId, + metadata: &RegionMetadata, + expected: &str, +) { + let current = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(current.primary_key, metadata.primary_key); + let projection = metadata + .column_metadatas + .iter() + .map(|column| current.column_index_by_id(column.column_id).unwrap()) + .collect(); + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + projection: Some(projection), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap(), + expected + ); +} + +async fn test_sync_columns_preserves_primary_key_order_with_format(flat_format: bool) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let mut row_schema = rows_schema(&request); + row_schema.push(tag_column_schema("tag_1", ColumnDataType::String)); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + engine + .handle_request(region_id, RegionRequest::Alter(add_tag1())) + .await + .unwrap(); + let metadata = engine.get_region(region_id).unwrap().metadata(); + assert_eq!(metadata.primary_key, vec![0, 3]); + + put_rows( + &engine, + region_id, + Rows { + schema: row_schema, + rows: build_rows_for_tags("a", "b", 0, 2, 10), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let before = RecordBatches::try_collect( + engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + .scan() + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + let mut columns = metadata.column_metadatas.clone(); + columns.push(ColumnMetadata { + column_schema: ColumnSchema::new("field_1", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 4, + }); + let tags = metadata.primary_key_columns().cloned().collect::>(); + let mut tags = tags.into_iter(); + for column in &mut columns { + if column.semantic_type == SemanticType::Tag { + *column = tags.next().unwrap(); + } + } + engine + .handle_request( + region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::SyncColumns { + column_metadatas: columns, + }, + }), + ) + .await + .unwrap(); + + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + assert_sync_columns_readable(&engine, region_id, &metadata, &before).await; +} + #[tokio::test] async fn test_alter_region() { test_alter_region_with_format(false).await; From 2b596fd52df31dcad779794877a296de7746f736 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 27 Sep 2026 02:36:11 +0000 Subject: [PATCH 042/133] ci: deploy MinIO chart with Silo images (#9362) Signed-off-by: Dennis Zhuang --- .github/actions/setup-minio/action.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/actions/setup-minio/action.yml b/.github/actions/setup-minio/action.yml index 89fbe28baa7..523203be79f 100644 --- a/.github/actions/setup-minio/action.yml +++ b/.github/actions/setup-minio/action.yml @@ -19,6 +19,8 @@ runs: --set rootUser=rootuser,rootPassword=rootpass123 \ --set buckets[0].name=default \ --set service.port=80,service.targetPort=9000 \ + --set image.repository=pgsty/silo,image.tag=RELEASE.2026-09-16T00-00-00Z \ + --set mcImage.repository=pgsty/mc,mcImage.tag=RELEASE.2026-09-16T00-00-00Z \ minio/minio \ --create-namespace \ -n minio From 411c0f173e0220ba50a2340864d9c73e56f0f8e9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 02:50:35 +0000 Subject: [PATCH 043/133] fix(mito2): stop parallel flat scan tasks once the receiver is dropped (#9368) * fix(mito2): stop parallel flat scan tasks once the receiver is dropped spawn_flat_scan_task ignored send errors, so after a query was cancelled the detached task kept reading its source to EOF. Break out of the loop when the receiver is gone. Signed-off-by: Dennis Zhuang * chore(mito2): log when a parallel scan task stops on a dropped receiver Signed-off-by: Dennis Zhuang * chore(mito2): include the send error in the parallel scan stop log Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/read/scan_region.rs | 65 ++++++++++++++++++++++++++++++- 1 file changed, 64 insertions(+), 1 deletion(-) diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 8edca2a0a9c..dc68835a11f 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -1663,7 +1663,15 @@ impl ScanInput { }; match maybe_batch { Some(Ok(batch)) => { - let _ = sender.send(Ok(batch)).await; + // The receiver is gone when the query is cancelled or finishes early, + // so stop reading the source. + if let Err(e) = sender.send(Ok(batch)).await { + debug!( + "Stop parallel scan task, receiver dropped, region_id: {}, error: {}", + region_id, e + ); + break; + } } Some(Err(e)) => { let _ = sender.send(Err(e)).await; @@ -3210,4 +3218,59 @@ mod tests { .build(); assert!(exact_sequence_range(&request, &version).is_err()); } + + #[tokio::test] + async fn test_parallel_flat_source_stops_after_receiver_dropped() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datatypes::arrow::array::Int64Array; + + const SOURCE_BATCHES: usize = 10_000; + + let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); + let mapper = FlatProjectionMapper::new(&metadata, [0, 2, 3]).unwrap(); + let env = SchedulerEnv::new().await; + let input = ScanInput::builder(env.access_layer.clone(), mapper).build(); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "v", + ArrowDataType::Int64, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap(); + let pulled = Arc::new(AtomicUsize::new(0)); + let mut released = Vec::new(); + let sources = (0..2) + .map(|_| { + let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + released.push(release_rx); + let batch = batch.clone(); + let pulled = pulled.clone(); + futures::stream::iter(0..SOURCE_BATCHES) + .map(move |_| { + // Dropped together with the source stream. + let _ = &release_tx; + pulled.fetch_add(1, Ordering::Relaxed); + Ok(batch.clone()) + }) + .boxed() + }) + .collect(); + + let streams = input + .create_parallel_flat_sources(sources, Arc::new(Semaphore::new(2)), 1) + .unwrap(); + drop(streams); + + for release_rx in released { + tokio::time::timeout(std::time::Duration::from_secs(10), release_rx) + .await + .unwrap() + .unwrap_err(); + } + // Each task may have read a few batches ahead into its channel before the + // receiver went away, but must not drain the rest of its source. + assert!(pulled.load(Ordering::Relaxed) < 100); + } } From bf67246645fd7384d1f7685ac6046d796aed1607 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 28 Sep 2026 03:10:21 +0000 Subject: [PATCH 044/133] perf: bound concurrent Metric export writers (#9296) * fix: preserve objects after conditional Metric export collisions Signed-off-by: jeremyhi * perf: share Metric export writer and payload budgets Signed-off-by: jeremyhi * test: cover bounded Metric export conversion and mixed roundtrips Signed-off-by: jeremyhi * fix: preserve export output after ambiguous close Signed-off-by: jeremyhi * test: align HTTP export timeout with managed cancellation Signed-off-by: jeremyhi * test: trim redundant Metric export roundtrips Signed-off-by: jeremyhi * fix: bound pending Metric export writer handles Signed-off-by: jeremyhi * fix: protect export files on failed writes Signed-off-by: jeremyhi * fix: preserve ordinary export cleanup and large rows Signed-off-by: jeremyhi * fix: report Metric export collisions as invalid arguments Signed-off-by: jeremyhi * fix: clean up failed conditional export closes Signed-off-by: jeremyhi * fix: support view arrays in managed ordinary export Signed-off-by: jeremyhi * test: expect successful cleanup after failed close Signed-off-by: jeremyhi * fix: clean up unsynced overwrite after close failure Signed-off-by: jeremyhi * fix: scope export JSON expansion estimates to JSON columns Signed-off-by: jeremyhi * fix: rechunk ordinary export batches within the write budget Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/common/datasource/src/parquet_writer.rs | 243 ++++++++- src/frontend/src/instance.rs | 33 +- src/object-store/src/secure_fs.rs | 254 ++++++++- src/object-store/tests/object_store_test.rs | 72 +++ src/operator/src/statement/copy_table_to.rs | 449 +++++++++++++++- src/operator/src/statement/database_copy.rs | 9 +- src/operator/src/statement/export_database.rs | 93 ++-- .../src/statement/export_logical_tables.rs | 318 ++++++++--- .../statement/export_logical_tables/tests.rs | 501 +++++++++++++++--- .../export_logical_tables/writers.rs | 310 +++++++++++ .../tests/export_logical_tables.rs | 10 +- 11 files changed, 2045 insertions(+), 247 deletions(-) create mode 100644 src/operator/src/statement/export_logical_tables/writers.rs diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index ae1fcec04a9..ed35c314339 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -29,6 +29,13 @@ use tokio_util::sync::CancellationToken; use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::error::{self, Result}; +/// Destination creation policy; conditional failures never authorize path deletion. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParquetCreationPolicy { + Overwrite, + IfNotExists, +} + /// Limits for one Parquet file. Flush thresholds are not hard memory caps. #[derive(Clone, Copy, Debug)] pub struct ParquetWriterLimits { @@ -49,6 +56,8 @@ pub struct ParquetFileWriter { store: ObjectStore, path: String, limits: Option, + creation: ParquetCreationPolicy, + close_started: bool, } impl ParquetFileWriter { @@ -60,6 +69,26 @@ impl ParquetFileWriter { path: &str, concurrency: usize, limits: Option, + ) -> Result { + Self::open_with_creation( + schema, + store, + path, + concurrency, + limits, + ParquetCreationPolicy::Overwrite, + ) + .await + } + + /// Open with a conditional policy only when the caller verified backend support. + pub async fn open_with_creation( + schema: SchemaRef, + store: ObjectStore, + path: &str, + concurrency: usize, + limits: Option, + creation: ParquetCreationPolicy, ) -> Result { let mut props = WriterProperties::builder() .set_compression(Compression::ZSTD(ZstdLevel::default())) @@ -90,6 +119,7 @@ impl ParquetFileWriter { .writer_with(path) .concurrent(concurrency) .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .if_not_exists(creation == ParquetCreationPolicy::IfNotExists) .await .context(error::WriteObjectSnafu { path })?; Ok(Self { @@ -98,6 +128,8 @@ impl ParquetFileWriter { store, path: path.to_owned(), limits, + creation, + close_started: false, }) } @@ -156,9 +188,12 @@ impl ParquetFileWriter { } async fn write_bytes(&mut self, bytes: Vec) -> Result<()> { - if !bytes.is_empty() { + let bytes = Bytes::from(bytes); + let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + // Slices retain the complete encoded allocation until its last submission. + for offset in (0..bytes.len()).step_by(chunk) { self.sink - .write(bytes) + .write(bytes.slice(offset..(offset + chunk).min(bytes.len()))) .await .context(error::WriteObjectSnafu { path: &self.path })?; } @@ -183,6 +218,7 @@ impl ParquetFileWriter { .context(error::JoinHandleSnafu)??; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; + self.close_started = true; self.sink .close() .await @@ -191,13 +227,16 @@ impl ParquetFileWriter { Ok(()) } - /// Abort an exclusively owned file after all in-flight operations have completed. - /// If the backend cannot abort, deletion assumes this attempt owns the path. + /// Abort after all in-flight operations complete. Preserve ambiguous commits; + /// conditional callers delegate cleanup exclusively to the backend. pub async fn abort(mut self) -> Result<()> { let result = self.sink.abort().await; - if result - .as_ref() - .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) + if self.creation == ParquetCreationPolicy::Overwrite + && result.as_ref().is_err_and(|error| { + error.kind() == object_store::ErrorKind::Unsupported + && (!self.close_started + || object_store::secure_fs::is_unsynced_overwrite_abort(error)) + }) { let store = self.store.clone(); let path = self.path.clone(); @@ -288,6 +327,196 @@ mod tests { .unwrap() } + #[tokio::test] + async fn abort_preserves_collisions_and_ambiguous_commits() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close( + &mut self, + ) -> object_store::Result { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + + for (creation, existing) in [ + (ParquetCreationPolicy::Overwrite, false), + (ParquetCreationPolicy::IfNotExists, false), + (ParquetCreationPolicy::IfNotExists, true), + ] { + let directory = common_test_util::temp_dir::create_temp_dir("conditional_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let path = "conditional.parquet"; + if existing { + store.write(path, "original").await.unwrap(); + } + let factory: MockWriterFactory = + Arc::new(|_, _, writer| Box::new(AmbiguousCommit(writer))); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = ParquetFileWriter::open_with_creation( + batch().schema(), + store.clone(), + path, + 1, + None, + creation, + ) + .await + .unwrap(); + writer.write(batch(), None).await.unwrap(); + assert!(writer.finish(None).await.is_err()); + assert!(writer.abort().await.is_err()); + if existing { + assert_eq!( + store.read(path).await.unwrap().to_bytes(), + Bytes::from_static(b"original") + ); + } else { + assert_eq!( + read(&store, path) + .await + .metadata() + .file_metadata() + .num_rows(), + 4 + ); + } + } + } + + #[tokio::test] + async fn overwrite_abort_deletes_after_unsynced_close() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; + + struct FailedClose(oio::Writer); + impl oio::Write for FailedClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close( + &mut self, + ) -> object_store::Result { + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "close failed before sync", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + let directory = common_test_util::temp_dir::create_temp_dir("unsynced_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let path = "partial.parquet"; + let factory: MockWriterFactory = Arc::new(|_, _, writer| Box::new(FailedClose(writer))); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = ParquetFileWriter::open(batch().schema(), store.clone(), path, 1, None) + .await + .unwrap(); + writer.write(batch(), None).await.unwrap(); + assert!(writer.finish(None).await.is_err()); + assert!(store.exists(path).await.unwrap()); + writer.abort().await.unwrap(); + assert!(!store.exists(path).await.unwrap()); + } + + #[tokio::test] + async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() { + use std::sync::Mutex; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + struct ObservedWriter(oio::Writer, Arc>>); + impl oio::Write for ObservedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.1.lock().unwrap().push(bytes.len()); + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let sizes = Arc::new(Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone())) + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut state = 17u64; + let values = (0..600_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as i64 + }) + .collect::>(); + let array = Arc::new(Int64Array::from(values)) as arrow::array::ArrayRef; + let batch = RecordBatch::try_from_iter([("a", array.clone()), ("b", array)]).unwrap(); + let mut writer = + ParquetFileWriter::open(batch.schema(), store.clone(), "large.parquet", 1, None) + .await + .unwrap(); + // No storage-layer chunking: observe the application's actual submissions. + writer.sink = store.writer("large.parquet").await.unwrap(); + writer.write(batch.clone(), None).await.unwrap(); + writer.finish(None).await.unwrap(); + let sizes = sizes.lock().unwrap().clone(); + let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; + assert!(sizes.iter().sum::() > limit); + assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); + let actual = read(&store, "large.parquet") + .await + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(), + batch + ); + } + #[tokio::test] async fn row_and_byte_limits_split_batches() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 92bc74ca10b..6e531214127 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2763,12 +2763,12 @@ mod tests { ) } - struct ReleasedExportSource { + struct CancellableExportSource { schema: GtSchemaRef, channels: std::sync::Mutex, oneshot::Receiver<()>)>>, } - impl DataSource for ReleasedExportSource { + impl DataSource for CancellableExportSource { fn get_stream( &self, _request: ScanRequest, @@ -2788,15 +2788,15 @@ mod tests { } #[tokio::test] - async fn test_metric_export_http_timeout_drains_ordinary_writer() { + async fn test_metric_export_http_timeout_cancels_ordinary_source() { let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); let (started_tx, started_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel(); + let (mut release_tx, release_rx) = oneshot::channel(); let info = test_table_info(1024, "source").unwrap(); let source = Arc::new(Table::new( Arc::new(info.clone()), FilterPushDownType::Unsupported, - Arc::new(ReleasedExportSource { + Arc::new(CancellableExportSource { schema: info.meta.schema.clone(), channels: std::sync::Mutex::new(Some((started_tx, release_rx))), }), @@ -2834,9 +2834,9 @@ mod tests { let server_task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); + let uri = reqwest::Url::from_directory_path(destination.path()).unwrap(); let sql = format!( - "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", - destination.path().display() + "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')" ); let request = tokio::spawn(async move { reqwest::Client::builder() @@ -2855,21 +2855,10 @@ mod tests { .unwrap(); let response = request.await.unwrap(); assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); - // A dropped ordinary stream would close this receiver before release. - release_tx.send(()).unwrap(); - let file = destination.path().join("source.parquet"); - tokio::time::timeout(Duration::from_secs(5), async { - loop { - if std::fs::read(&file) - .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .unwrap(); + tokio::time::timeout(Duration::from_secs(5), release_tx.closed()) + .await + .unwrap(); + assert!(!destination.path().join("source.parquet").exists()); server_task.abort(); } diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 18afa3ac820..feaaa42fdf2 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -274,6 +274,7 @@ impl Service for SecureFsBackend { path, args, file: None, + synced: false, }) } @@ -387,6 +388,29 @@ struct SecureFsWriter { path: PathBuf, args: OpWrite, file: Option, + synced: bool, +} + +#[derive(Debug)] +struct UnsyncedOverwrite { + flush_error: Option, +} + +impl fmt::Display for UnsyncedOverwrite { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "overwrite file was not synced") + } +} + +impl std::error::Error for UnsyncedOverwrite { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.flush_error.as_ref().map(|error| error as _) + } +} + +/// Whether an abort error proves an opened overwrite has not been synced. +pub fn is_unsynced_overwrite_abort(error: &Error) -> bool { + std::error::Error::source(error).is_some_and(|source| source.is::()) } impl SecureFsWriter { @@ -438,10 +462,18 @@ impl oio::Write for SecureFsWriter { } async fn close(&mut self) -> Result { - let file = self.ensure_file().await?; - file.flush().await.map_err(new_std_io_error)?; - file.sync_all().await.map_err(new_std_io_error)?; - let metadata = file.metadata().await.map_err(new_std_io_error)?; + { + let file = self.ensure_file().await?; + file.flush().await.map_err(new_std_io_error)?; + file.sync_all().await.map_err(new_std_io_error)?; + } + self.synced = true; + let metadata = self + .ensure_file() + .await? + .metadata() + .await + .map_err(new_std_io_error)?; let mut builder = MetadataBuilder::file(metadata.len()); builder.last_modified(Timestamp::try_from( metadata.modified().map_err(new_std_io_error)?, @@ -450,10 +482,43 @@ impl oio::Write for SecureFsWriter { } async fn abort(&mut self) -> Result<()> { - Err(Error::new( + // Tokio writes may finish in the blocking pool after write_all returns. + let flush = match self.file.as_mut() { + Some(file) => file.flush().await.map_err(new_std_io_error), + None => Ok(()), + }; + if self.args.if_not_exists() { + // A failed exclusive create owns no file. Once data is synced, preserve + // potentially committed output for the caller's deliberate retry. + if let Some(file) = self.file.take() { + drop(file); + if !self.synced { + let root = self.root.clone(); + let path = self.path.clone(); + let cleanup = + common_runtime::spawn_blocking_global(move || root.dir.remove_file(path)) + .await + .map_err(new_task_join_error) + .and_then(|result| result.map_err(new_std_io_error)); + return flush.and(cleanup); + } + } + return flush; + } + if self.file.is_none() { + return flush; + } + let error = Error::new( ErrorKind::Unsupported, "filesystem writes cannot be aborted without atomic writes", - )) + ); + if !self.synced { + return Err(error.set_source(UnsyncedOverwrite { + flush_error: flush.err(), + })); + } + flush?; + Err(error) } } @@ -800,16 +865,181 @@ mod tests { } #[tokio::test] - async fn test_writer_abort_is_unsupported_without_atomic_write() { - let temp_dir = create_temp_dir("secure_fs_writer_abort"); + async fn test_conditional_abort_only_removes_owned_partial_file() { + let temp_dir = create_temp_dir("secure_fs_conditional_abort"); let operator = SecureFsRoot::open(temp_dir.path()) .unwrap() .build_operator(); - let mut writer = operator.writer("partial").await.unwrap(); - writer.write(Bytes::from_static(b"partial")).await.unwrap(); + for started in [false, true] { + let mut writer = operator + .writer_with("partial") + .if_not_exists(true) + .await + .unwrap(); + if started { + writer.write(Bytes::from_static(b"partial")).await.unwrap(); + } + writer.abort().await.unwrap(); + assert!(!operator.exists("partial").await.unwrap()); + } + } - let error = writer.abort().await.unwrap_err(); + #[tokio::test] + async fn test_overwrite_abort_preserves_unopened_destination() { + use std::path::PathBuf; - assert_eq!(ErrorKind::Unsupported, error.kind()); + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + let temp_dir = create_temp_dir("secure_fs_unopened_abort"); + std::fs::write(temp_dir.path().join("existing"), b"original").unwrap(); + let mut writer = super::SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("existing"), + args: OpWrite::default(), + file: None, + synced: false, + }; + + writer.abort().await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("existing")).unwrap(), + b"original" + ); + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn test_abort_after_background_write_error() { + use std::path::PathBuf; + + use opendal::options::WriteOptions; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + use tokio::io::AsyncWriteExt; + + for if_not_exists in [true, false] { + let temp_dir = create_temp_dir("secure_fs_flush_error_abort"); + let root = SecureFsRoot::open(temp_dir.path()).unwrap(); + let (args, _) = OpWrite::from_options( + &root.build_operator().info().capability(), + WriteOptions { + if_not_exists, + ..Default::default() + }, + ) + .unwrap(); + let mut writer = super::SecureFsWriter { + root, + path: PathBuf::from("partial"), + args, + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let mut failing_file = tokio::fs::OpenOptions::new() + .write(true) + .open("/dev/full") + .await + .unwrap(); + failing_file.write_all(b"partial").await.unwrap(); + writer.file = Some(failing_file); + + let error = writer.abort().await.unwrap_err(); + if if_not_exists { + assert!(!temp_dir.path().join("partial").exists()); + } else { + assert_eq!(error.kind(), opendal::ErrorKind::Unsupported); + assert!(std::error::Error::source(&error).is_some()); + } + } + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn test_conditional_abort_after_close_flush_error() { + use std::path::PathBuf; + + use opendal::options::WriteOptions; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + use tokio::io::AsyncWriteExt; + + let temp_dir = create_temp_dir("secure_fs_close_flush_error"); + let root = SecureFsRoot::open(temp_dir.path()).unwrap(); + let (args, _) = OpWrite::from_options( + &root.build_operator().info().capability(), + WriteOptions { + if_not_exists: true, + ..Default::default() + }, + ) + .unwrap(); + let mut writer = super::SecureFsWriter { + root, + path: PathBuf::from("partial"), + args, + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let mut failing_file = tokio::fs::OpenOptions::new() + .write(true) + .open("/dev/full") + .await + .unwrap(); + failing_file.write_all(b"partial").await.unwrap(); + writer.file = Some(failing_file); + + assert!(writer.close().await.is_err()); + writer.abort().await.unwrap(); + assert!(!temp_dir.path().join("partial").exists()); + } + + #[test] + fn test_writer_abort_drains_before_reporting_unsupported() { + use std::path::PathBuf; + + use opendal::Buffer; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + use super::SecureFsWriter; + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap() + .block_on(async { + let temp_dir = create_temp_dir("secure_fs_writer_abort"); + let mut writer = SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("partial"), + args: OpWrite::default(), + file: None, + synced: false, + }; + writer.ensure_file().await.unwrap(); + let (started, ready) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel(); + let blocking = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + blocked.recv().unwrap(); + }); + ready.await.unwrap(); + writer.write(Buffer::from("partial")).await.unwrap(); + let abort = writer.abort(); + tokio::pin!(abort); + let pending = futures::poll!(&mut abort).is_pending(); + release.send(()).unwrap(); + assert!(pending); + assert_eq!(ErrorKind::Unsupported, abort.await.unwrap_err().kind()); + blocking.await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("partial")).unwrap(), + b"partial" + ); + }); } } diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index 735d2c56f15..dc2482d1722 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -290,6 +290,9 @@ async fn test_s3_backend() -> Result<()> { .secret_access_key(&env::var("GT_S3_ACCESS_KEY")?) .region(&env::var("GT_S3_REGION")?) .bucket(&bucket); + if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") { + builder = builder.endpoint(&endpoint); + } // Honors an S3-compatible endpoint (MinIO in CI) so the test does not // fall back to resolving the bucket against real AWS. @@ -306,6 +309,7 @@ async fn test_s3_backend() -> Result<()> { test_object_crud(&store).await?; test_object_list(&store).await?; test_object_list_start_after(&store).await?; + test_conditional_creation(&store).await?; assert_opendal_metrics(); guard.remove_all().await?; } @@ -313,6 +317,74 @@ async fn test_s3_backend() -> Result<()> { Ok(()) } +async fn test_conditional_creation(store: &ObjectStore) -> Result<()> { + const PART: usize = 8 * 1024 * 1024; + for size in [32, 2 * PART + 17] { + let payload = Bytes::from(vec![42; size]); + for preexisting in [false, true] { + let path = format!("conditional-{size}-{preexisting}"); + if preexisting { + store.write(&path, "original").await?; + } + let attempt = |value: Bytes| { + let path = &path; + async move { + let mut writer = store + .writer_with(path) + .if_not_exists(true) + .concurrent(1) + .chunk(PART) + .await + .unwrap(); + let result = async { + for offset in (0..value.len()).step_by(PART) { + writer + .write(value.slice(offset..(offset + PART).min(value.len()))) + .await?; + } + writer.close().await + } + .await; + if result.is_err() { + writer.abort().await.unwrap(); + } + result + } + }; + let (a, b) = tokio::join!(attempt(payload.clone()), attempt(payload.clone())); + assert_eq!( + usize::from(a.is_ok()) + usize::from(b.is_ok()), + usize::from(!preexisting) + ); + for error in [a.err(), b.err()].into_iter().flatten() { + assert!( + matches!( + error.kind(), + object_store::ErrorKind::ConditionNotMatch + | object_store::ErrorKind::AlreadyExists + ), + "{error:?}" + ); + } + let expected = if preexisting { + Bytes::from_static(b"original") + } else { + payload.clone() + }; + assert_eq!(store.read(&path).await?.to_bytes(), expected); + store.delete(&path).await?; + } + } + Ok(()) +} + +#[tokio::test] +async fn test_secure_fs_conditional_creation() -> Result<()> { + let dir = TempDir::new()?; + let store = object_store::secure_fs::SecureFsRoot::open(dir.path())?.build_operator(); + test_conditional_creation(&store).await +} + #[tokio::test] async fn test_oss_backend() -> Result<()> { common_telemetry::init_default_ut_logging(); diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 21b82f51a3b..d3d44f7655f 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -15,6 +15,8 @@ use std::collections::HashMap; use std::sync::Arc; +use arrow::array::{Array, ArrayRef, AsArray, UInt64Array, make_array}; +use arrow::datatypes::DataType; use client::OutputData; use common_base::readable_size::ReadableSize; use common_datasource::file_format::Format; @@ -22,6 +24,7 @@ use common_datasource::file_format::csv::stream_to_csv; use common_datasource::file_format::json::stream_to_json; use common_datasource::file_format::parquet::stream_to_parquet; use common_datasource::object_store::build_backend_for_write_with_path; +use common_datasource::parquet_writer::ParquetFileWriter; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; use common_recordbatch::{ @@ -32,16 +35,22 @@ use common_telemetry::{debug, tracing}; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference as DfTableReference; use datafusion_expr::LogicalPlanBuilder; +use futures::StreamExt; use object_store::ObjectStore; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use table::TableRef; use table::requests::CopyTableRequest; use table::table::adapter::DfTableProviderAdapter; use table::table_reference::TableReference; +use tokio_util::sync::CancellationToken; use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result}; use crate::statement::StatementExecutor; +use crate::statement::export_logical_tables::writers::ExportWriteBudget; +use crate::statement::export_logical_tables::{ + expand_export_batch, map_writer_error, rows_within_budget, +}; // The buffer size should be greater than 5MB (minimum multipart upload size). /// Buffer size to flush data to object stores. @@ -117,6 +126,17 @@ impl StatementExecutor { table: TableRef, req: CopyTableRequest, query_ctx: QueryContextRef, + ) -> Result { + self.copy_captured_table_to_managed(table, req, query_ctx, None) + .await + } + + pub(crate) async fn copy_captured_table_to_managed( + &self, + table: TableRef, + req: CopyTableRequest, + query_ctx: QueryContextRef, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, ) -> Result { let info = table.table_info(); let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name); @@ -165,7 +185,7 @@ impl StatementExecutor { } = &req; debug!("Copy table: {table_id} to location: {location}"); - self.copy_to_file(&format, output, location, connection) + self.copy_to_file_managed(&format, output, location, connection, managed) .await } @@ -176,9 +196,25 @@ impl StatementExecutor { location: &str, connection: &HashMap, ) -> Result { - let output = output - .map_dictionary_to_values() - .context(error::BuildRecordBatchSnafu)?; + self.copy_to_file_managed(format, output, location, connection, None) + .await + } + + async fn copy_to_file_managed( + &self, + format: &Format, + output: Output, + location: &str, + connection: &HashMap, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, + ) -> Result { + let output = if managed.is_none() { + output + .map_dictionary_to_values() + .context(error::BuildRecordBatchSnafu)? + } else { + output + }; let stream = match output.data { OutputData::Stream(stream) => stream, OutputData::RecordBatches(record_batches) => record_batches.as_stream(), @@ -192,7 +228,408 @@ impl StatementExecutor { let filename = backend.object_path.context(error::UnexpectedSnafu { violated: format!("Expected filename, path: {location}"), })?; - self.stream_to_file(stream, format, backend.object_store, &filename) + if let Some((budget, token)) = managed { + stream_to_managed_parquet(stream, backend.object_store, &filename, budget, token).await + } else { + self.stream_to_file(stream, format, backend.object_store, &filename) + .await + } + } +} + +pub(crate) async fn stream_to_managed_parquet( + mut stream: SendableRecordBatchStream, + store: ObjectStore, + path: &str, + budget: &ExportWriteBudget, + token: &CancellationToken, +) -> Result { + use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema}; + let original = stream.schema(); + let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone()); + let json_columns = expanded_schema + .column_schemas() + .iter() + .enumerate() + .filter_map(|(index, column)| { + (column.data_type.is_json() && !column.data_type.is_json2()).then_some(index) + }) + .collect::>(); + let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone()); + let output_schema = if json { + mapped_schema + } else { + expanded_schema.clone() + }; + let mut writer = + ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None) .await + .context(error::WriteStreamToFileSnafu { path })?; + let mut started = false; + let result = async { + let mut rows = 0; + loop { + let batch = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + batch = stream.next() => batch, + }; + let Some(batch) = batch else { break }; + let batch = batch + .context(error::BuildRecordBatchSnafu)? + .into_df_record_batch(); + let mut offset = 0; + while offset < batch.num_rows() { + // The scan batch remains query-owned. Bound the downstream slice and + // detach its buffers below if it would retain more than its reservation. + let (conversion, retained) = + ExportWriteBudget::conversion_budget(0, batch.num_columns(), usize::MAX)?; + let input = batch.clone(); + let json_columns = json_columns.clone(); + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, offset, input.num_rows(), conversion, &json_columns) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, token).await?; + let batch = batch.clone(); + let (original, expanded_schema, output_schema) = ( + original.clone(), + expanded_schema.clone(), + output_schema.clone(), + ); + let (batch, len) = common_runtime::spawn_blocking_global(move || { + let mut batch = RecordBatch::from_df_record_batch( + original.clone(), + batch.slice(offset, len), + ); + if expand { + batch = RecordBatch::from_df_record_batch( + expanded_schema.clone(), + expand_export_batch( + &batch.into_df_record_batch(), + expanded_schema.arrow_schema().clone(), + )?, + ); + } + if json { + batch = map_json_type_to_string(batch, &expanded_schema, &output_schema) + .context(error::BuildRecordBatchSnafu)?; + } + let mut batch = batch.into_df_record_batch(); + if batch.get_array_memory_size() > reservation { + let indices = UInt64Array::from_iter_values(0..batch.num_rows() as u64); + batch = arrow::compute::take_record_batch(&batch, &indices) + .context(error::ComputeArrowSnafu)?; + let arrays = batch + .columns() + .iter() + .map(compact_view_buffers) + .collect::>>()?; + batch = arrow::record_batch::RecordBatch::try_new(batch.schema(), arrays) + .context(error::ComputeArrowSnafu)?; + } + ensure!( + batch.get_array_memory_size() <= reservation, + error::LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok::<_, error::Error>((batch, len)) + }) + .await + .context(error::JoinTaskSnafu)??; + started = true; + let write = writer.write(batch, Some(token)).await; + drop(permit); + write.map_err(|error| map_writer_error(error, path))?; + rows += len; + offset += len; + } + } + started = true; + writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, path))?; + Ok(rows) + } + .await; + if result.is_err() { + token.cancel(); + if started && let Err(error) = writer.abort().await { + common_telemetry::warn!(error; "Failed to abort ordinary export file"); + } + } + result +} + +// Arrow take copies ordinary buffers but shares view data buffers, including +// views nested inside lists or structs. Reclaim those unselected values too. +fn compact_view_buffers(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8View => Ok(Arc::new(array.as_string_view().gc())), + DataType::BinaryView => Ok(Arc::new(array.as_binary_view().gc())), + _ => { + let data = array.to_data(); + if data.child_data().is_empty() { + return Ok(array.clone()); + } + let children = data + .child_data() + .iter() + .map(|child| { + compact_view_buffers(&make_array(child.clone())).map(|array| array.to_data()) + }) + .collect::>>()?; + Ok(make_array( + data.into_builder() + .child_data(children) + .build() + .context(error::ComputeArrowSnafu)?, + )) + } + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{ + ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int32Array, StringArray, + StringViewArray, + }; + use arrow::datatypes::Int32Type; + use common_recordbatch::{RecordBatch, RecordBatches}; + use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + + use super::*; + + #[tokio::test] + async fn managed_copy_preserves_existing_file_before_sink_open() { + let temp_dir = common_test_util::temp_dir::create_temp_dir("managed_copy_existing_file"); + let store = object_store::secure_fs::SecureFsRoot::open(temp_dir.path()) + .unwrap() + .build_operator(); + let path = "existing.parquet"; + store.write(path, "original").await.unwrap(); + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = arrow::record_batch::RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let batches = async_stream::stream! { + yield Ok(batch); + yield Err(datafusion::error::DataFusionError::Execution("source failed".into())); + }; + let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema.arrow_schema().clone(), + batches, + ); + let stream = + common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream)) + .unwrap(); + + let result = stream_to_managed_parquet( + Box::pin(stream), + store.clone(), + path, + &ExportWriteBudget::new(1), + &CancellationToken::new(), + ) + .await; + + assert!(result.is_err()); + assert_eq!( + store.read(path).await.unwrap().to_bytes().as_ref(), + b"original" + ); + } + + #[tokio::test] + async fn managed_copy_rechunks_large_backing_buffers() { + let value = "x".repeat(9 * 1024); + for view in [false, true] { + let data_type = if view { + ConcreteDataType::utf8_view_datatype() + } else { + ConcreteDataType::string_datatype() + }; + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", data_type, false, + )])); + let values = vec![value.as_str(); 8192]; + let array: ArrayRef = if view { + Arc::new(StringViewArray::from(values)) + } else { + Arc::new(StringArray::from(values)) + }; + let batch = arrow::record_batch::RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![array], + ) + .unwrap(); + assert!(batch.get_array_memory_size() > 64 * 1024 * 1024); + // Also exercise a small slice that still references the large allocation. + for input in [batch.clone(), batch.slice(200, 3)] { + let rows = input.num_rows(); + let batches = RecordBatches::try_new( + schema.clone(), + vec![RecordBatch::from_df_record_batch(schema.clone(), input)], + ) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(1); + assert_eq!( + stream_to_managed_parquet( + batches.as_stream(), + store.clone(), + "large.parquet", + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(), + rows + ); + let reader = ParquetRecordBatchReaderBuilder::try_new( + store.read("large.parquet").await.unwrap().to_bytes(), + ) + .unwrap() + .build() + .unwrap(); + let mut actual_rows = 0; + for batch in reader { + let batch = batch.unwrap(); + let strings = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + assert!( + strings + .as_string::() + .iter() + .all(|item| item == Some(value.as_str())) + ); + actual_rows += batch.num_rows(); + } + assert_eq!(actual_rows, rows); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } + } + } + + #[tokio::test] + async fn managed_copy_preserves_dictionary_json_views_and_empty_schema() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "host", + ConcreteDataType::dictionary_datatype( + ConcreteDataType::int32_datatype(), + ConcreteDataType::string_datatype(), + ), + true, + ), + ColumnSchema::new("json", ConcreteDataType::json_datatype(), true), + ColumnSchema::new("text_view", ConcreteDataType::utf8_view_datatype(), true), + ColumnSchema::new( + "binary_view", + ConcreteDataType::binary_view_datatype(), + true, + ), + ])); + let dictionary = DictionaryArray::::new( + Int32Array::from(vec![Some(0), None, Some(1), Some(1)]), + Arc::new(StringArray::from(vec!["tag", &"x".repeat(3 * 1024 * 1024)])), + ); + let json = datatypes::types::parse_string_to_jsonb(&format!( + r#"{{"value":"escaped\ntext{}","n":123456789}}"#, + "y".repeat(160_000) + )) + .unwrap(); + let arrays = vec![ + Arc::new(dictionary) as ArrayRef, + Arc::new(BinaryArray::from(vec![ + Some(json.as_slice()), + None, + None, + Some(json.as_slice()), + ])), + Arc::new(StringViewArray::from(vec![ + Some("short"), + None, + Some("long view"), + Some("long view"), + ])), + Arc::new(BinaryViewArray::from(vec![ + Some(&b"short"[..]), + None, + Some(&b"long view"[..]), + Some(&b"long view"[..]), + ])), + ]; + let batch = + arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays) + .unwrap(); + let batch = RecordBatch::from_df_record_batch(schema.clone(), batch); + for empty in [false, true] { + let batches = RecordBatches::try_new( + schema.clone(), + if empty { vec![] } else { vec![batch.clone()] }, + ) + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let managed = stream_to_managed_parquet( + batches.as_stream(), + store.clone(), + "managed.parquet", + &budget, + &token, + ) + .await + .unwrap(); + let output = Output::new_with_record_batches(batches) + .map_dictionary_to_values() + .unwrap(); + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected batches"); + }; + let stream = SendableRecordBatchMapper::new( + batches.as_stream(), + map_json_type_to_string, + map_json_type_to_string_schema, + ); + let ordinary = stream_to_parquet( + Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))), + store.clone(), + "ordinary.parquet", + WRITE_CONCURRENCY, + ) + .await + .unwrap(); + assert_eq!(managed, ordinary); + let mut outputs = Vec::new(); + for name in ["managed.parquet", "ordinary.parquet"] { + let reader = ParquetRecordBatchReaderBuilder::try_new( + store.read(name).await.unwrap().to_bytes(), + ) + .unwrap(); + let schema = reader.schema().clone(); + let values = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + outputs.push((schema, values)); + } + assert_eq!(outputs[0], outputs[1]); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } } } diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs index 117104b9d8f..9878560afaf 100644 --- a/src/operator/src/statement/database_copy.rs +++ b/src/operator/src/statement/database_copy.rs @@ -27,6 +27,7 @@ use table::TableRef; use table::metadata::TableType; use table::requests::CopyDatabaseRequest; use table::table_reference::TableReference; +use tokio::sync::Semaphore; use url::Url; use crate::error::{self, Result}; @@ -56,7 +57,7 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap().ok()) .unwrap_or_else(get_total_cpu_cores) - .max(1) + .clamp(1, Semaphore::MAX_PERMITS) } /// Rejects import-only layouts before either database export path creates output. @@ -321,5 +322,11 @@ mod tests { let options = HashMap::from([("parallelism".to_string(), "0".to_string())]); assert_eq!(parse_parallelism_from_option_map(&options), 1); + + let options = HashMap::from([("parallelism".to_string(), usize::MAX.to_string())]); + assert_eq!( + parse_parallelism_from_option_map(&options), + Semaphore::MAX_PERMITS + ); } } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index 7c5de554487..dbe18ba763a 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -37,6 +37,7 @@ use crate::statement::database_copy::{ DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory, validate_database_export_layout, }; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error}; use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; /// A validated request-scoped selection, not a metadata snapshot or an ACL token. @@ -202,46 +203,52 @@ impl StatementExecutor { } output_files.sort(); let req = &plan.request; - let rows = run_database_export_jobs( - plan.jobs, - parse_parallelism_from_option_map(&req.with), - cancellation, - |job, token| { - let ctx = ctx.clone(); - async move { - match job { - DatabaseExportJob::Metric(unit) => self - .export_logical_tables( - &unit, - &req.location, - &req.connection, - req.time_range.as_ref(), - LogicalTableExportLimits::default(), - &token, - ctx, - ) - .await - .map(|summary| summary.rows), - DatabaseExportJob::Ordinary { table, output } => { - let info = table.table_info(); - let copy = CopyTableRequest { - catalog_name: info.catalog_name.clone(), - schema_name: info.schema_name.clone(), - table_name: info.name.clone(), - location: output.location, - with: req.with.clone(), - connection: req.connection.clone(), - pattern: None, - direction: CopyDirection::Export, - timestamp_range: req.time_range, - limit: None, - }; - self.copy_captured_table_to(table, copy, ctx).await - } + let parallelism = parse_parallelism_from_option_map(&req.with); + let budget = ExportWriteBudget::new(parallelism); + let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| { + let ctx = ctx.clone(); + let budget = budget.clone(); + async move { + match job { + DatabaseExportJob::Metric(unit) => self + .export_logical_tables_managed( + &unit, + &req.location, + &req.connection, + req.time_range.as_ref(), + LogicalTableExportLimits::default(), + &token, + ctx, + budget, + ) + .await + .map(|summary| summary.rows), + DatabaseExportJob::Ordinary { table, output } => { + let info = table.table_info(); + let copy = CopyTableRequest { + catalog_name: info.catalog_name.clone(), + schema_name: info.schema_name.clone(), + table_name: info.name.clone(), + location: output.location, + with: req.with.clone(), + connection: req.connection.clone(), + pattern: None, + direction: CopyDirection::Export, + timestamp_range: req.time_range, + limit: None, + }; + let _permit = budget.writer(&token).await?; + self.copy_captured_table_to_managed( + table, + copy, + ctx, + Some((&budget, &token)), + ) + .await } } - }, - ) + } + }) .await?; Ok(DatabaseExportSummary { rows, output_files }) } @@ -261,6 +268,7 @@ async fn run_database_export_jobs>>( loop { while first_error.is_none() && !cancellation.is_cancelled() + && !token.is_cancelled() && active.len() < parallelism.max(1) { let Some(job) = jobs.next() else { break }; @@ -284,11 +292,12 @@ async fn run_database_export_jobs>>( }; match result { Some(Ok(count)) => rows += count, - Some(Err(err)) if first_error.is_none() => { - first_error = Some(err); + Some(Err(err)) => { + if !cancellation.is_cancelled() || first_error.is_none() { + retain_error(&mut first_error, err); + } token.cancel(); } - Some(Err(err)) => common_telemetry::warn!(err; "Failed to drain database export job"), None => break, } } @@ -339,7 +348,7 @@ mod tests { } .fail(); } - // Ordinary COPY continues its I/O even when Metric jobs cancel. + // Already-started ordinary I/O drains after cancellation. token.cancelled().await; started.send(10).unwrap(); finish_io.acquire().await.unwrap().forget(); diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 7547a3f54de..4a2e186d477 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -19,16 +19,21 @@ //! exclusively by the export attempt. Completion and cleanup of closed files //! belong to the caller's chunk protocol, not individual file completion. +pub(crate) mod writers; + use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; -use arrow::array::{Array, AsArray, UInt32Array}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, AsArray, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{cast, take}; use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; use common_datasource::object_store::build_backend_for_write; -use common_datasource::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; +use common_datasource::parquet_writer::{ + ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits, +}; use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; use common_query::OutputData; use common_recordbatch::SendableRecordBatchStream; @@ -52,6 +57,7 @@ use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; use crate::statement::StatementExecutor; use crate::statement::database_copy::DatabaseExportFile; +use crate::statement::export_logical_tables::writers::{ExportWriteBudget, Payload, TableWriters}; /// Export preprocessing and per-file limits. Query memory and spill remain /// governed by the query engine. @@ -101,7 +107,7 @@ pub struct LogicalTableExport { logical_tables: BTreeMap, } -struct LogicalTableProjection { +pub(crate) struct LogicalTableProjection { output: DatabaseExportFile, schema: SchemaRef, projection: Vec, @@ -306,6 +312,31 @@ impl StatementExecutor { limits: LogicalTableExportLimits, cancellation: &CancellationToken, query_ctx: QueryContextRef, + ) -> Result { + self.export_logical_tables_managed( + unit, + directory, + connection, + time_range, + limits, + &cancellation.child_token(), + query_ctx, + ExportWriteBudget::new(1), + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn export_logical_tables_managed( + &self, + unit: &LogicalTableExport, + directory: &str, + connection: &HashMap, + time_range: Option<&TimestampRange>, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + query_ctx: QueryContextRef, + budget: Arc, ) -> Result { limits.validate()?; let (store, stream) = tokio::select! { @@ -325,10 +356,11 @@ impl StatementExecutor { Ok((store, stream)) } => result?, }; - export_stream(unit, stream, &store, limits, cancellation).await + export_stream_managed(unit, stream, &store, limits, cancellation, budget).await } } +#[cfg(test)] async fn export_stream( unit: &LogicalTableExport, stream: SendableRecordBatchStream, @@ -336,19 +368,42 @@ async fn export_stream( limits: LogicalTableExportLimits, cancellation: &CancellationToken, ) -> Result { - let mut active = None; - let result = write_tables(unit, stream, store, limits, cancellation, &mut active).await; - if result.is_err() - && let Some(writer) = active - && let Err(cleanup_error) = writer - .writer - .abort() - .await - .map_err(|error| map_writer_error(error, &writer.path)) - { - common_telemetry::warn!(cleanup_error; "Failed to clean up incomplete Metric export file"); - } - result + export_stream_managed( + unit, + stream, + store, + limits, + &cancellation.child_token(), + ExportWriteBudget::new(1), + ) + .await +} + +async fn export_stream_managed( + unit: &LogicalTableExport, + stream: SendableRecordBatchStream, + store: &ObjectStore, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + budget: Arc, +) -> Result { + let mut writers = TableWriters::new(budget.clone()); + let result = write_tables( + unit, + stream, + store, + limits, + cancellation, + &mut writers, + &budget, + ) + .await; + let (summary, result) = match result { + Ok(summary) => (summary, Ok(())), + Err(error) => (LogicalTableExportSummary::default(), Err(error)), + }; + writers.drain(result, cancellation).await?; + Ok(summary) } fn check_cancelled(cancellation: &CancellationToken) -> Result<()> { @@ -365,7 +420,8 @@ async fn write_tables( store: &ObjectStore, limits: LogicalTableExportLimits, cancellation: &CancellationToken, - active: &mut Option, + writers: &mut TableWriters, + budget: &ExportWriteBudget, ) -> Result { let id_index = stream @@ -423,22 +479,19 @@ async fn write_tables( while end < batch.num_rows() && ids.value(end) == id { end += 1; } - if active.as_ref().is_some_and(|writer| writer.table_id != id) { - finish_active(active, cancellation).await?; + if writers.table_id().is_some_and(|last| last != id) { + writers.close_input(); } check_cancelled(cancellation)?; if let Some(file) = unit.logical_tables.get(&id) { - if active.is_none() { - *active = Some(ActiveWriter::open(id, file, store, limits).await?); + if writers.table_id().is_none() { + writers.open(id, file, store, limits, cancellation).await?; written.insert(id); summary.files += 1; } let projected = batch .project(&file.projection) .context(error::ProjectSchemaSnafu)?; - let writer = active.as_mut().context(error::UnexpectedSnafu { - violated: "missing logical writer", - })?; let mut offset = start; while offset < end { let (expanded, consumed) = expand_bounded_slice( @@ -447,15 +500,12 @@ async fn write_tables( offset, end, limits.conversion_bytes, + budget, + cancellation, ) .await?; check_cancelled(cancellation)?; - // Do not drop an in-flight file operation before cleanup. - writer - .writer - .write(expanded, Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; + writers.send(expanded, cancellation).await?; check_cancelled(cancellation)?; offset += consumed; summary.rows += consumed; @@ -466,12 +516,12 @@ async fn write_tables( start = end; } } - finish_active(active, cancellation).await?; + writers.close_input(); for (&id, file) in &unit.logical_tables { if !written.contains(&id) { check_cancelled(cancellation)?; - *active = Some(ActiveWriter::open(id, file, store, limits).await?); - finish_active(active, cancellation).await?; + writers.open(id, file, store, limits, cancellation).await?; + writers.close_input(); summary.files += 1; } } @@ -480,85 +530,158 @@ async fn write_tables( } struct ActiveWriter { - table_id: u32, path: String, writer: ParquetFileWriter, } impl ActiveWriter { async fn open( - id: u32, table: &LogicalTableProjection, store: &ObjectStore, limits: LogicalTableExportLimits, ) -> Result { let path = table.output.path.clone(); - ensure!( - !store - .exists(&path) - .await - .context(error::ReadObjectSnafu { path: &path })?, - InvalidLogicalTableExportSnafu { - reason: format!("output already exists: {path}") - } - ); - let writer = ParquetFileWriter::open( + let conditional = store.info().capability().write_with_if_not_exists; + if !conditional { + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + InvalidLogicalTableExportSnafu { + reason: format!("output already exists: {path}") + } + ); + } + let writer = ParquetFileWriter::open_with_creation( table.schema.clone(), store.clone(), &path, 1, Some(limits.writer), + if conditional { + ParquetCreationPolicy::IfNotExists + } else { + ParquetCreationPolicy::Overwrite + }, ) .await .map_err(|error| map_writer_error(error, &path))?; - Ok(Self { - table_id: id, - path, - writer, - }) + Ok(Self { path, writer }) } } -async fn finish_active( - active: &mut Option, - cancellation: &CancellationToken, -) -> Result<()> { - if let Some(writer) = active.as_mut() { - writer - .writer - .finish(Some(cancellation)) - .await - .map_err(|error| map_writer_error(error, &writer.path))?; - check_cancelled(cancellation)?; - *active = None; - } - Ok(()) -} - async fn expand_bounded_slice( batch: RecordBatch, schema: SchemaRef, start: usize, end: usize, - budget: usize, -) -> Result<(RecordBatch, usize)> { + requested: usize, + budget: &ExportWriteBudget, + cancellation: &CancellationToken, +) -> Result<(Payload, usize)> { + let (conversion, retained) = ExportWriteBudget::conversion_budget( + batch.get_array_memory_size(), + batch.num_columns(), + requested, + )?; + let input = batch.clone(); + let (len, estimated) = common_runtime::spawn_blocking_global(move || { + rows_within_budget(&input, start, end, conversion, &[]) + }) + .await + .context(error::JoinTaskSnafu)??; + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { - let len = rows_within_budget(&batch, start, end, budget)?; let slice = batch.slice(start, len); - let arrays = slice - .columns() - .iter() - .zip(schema.fields()) - .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) - .collect::>>()?; - let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; - Ok((expanded, len)) + let expanded = expand_export_batch(&slice, schema)?; + ensure!( + expanded.get_array_memory_size() <= reservation, + LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok(( + Payload { + batch: expanded, + permit, + }, + len, + )) }) .await .context(error::JoinTaskSnafu)? } -fn map_writer_error(source: common_datasource::error::Error, path: &str) -> error::Error { +/// Limits nested casts to selected children: Arrow's List cast otherwise expands +/// the entire values array, even when the parent has been sliced. +pub(crate) fn expand_export_batch(batch: &RecordBatch, schema: SchemaRef) -> Result { + let arrays = batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu) +} + +fn expand_export_array(array: &ArrayRef, target: &DataType) -> Result { + if array.data_type() == target { + return Ok(array.clone()); + } + match (array.data_type(), target) { + (DataType::Dictionary(_, _), _) => { + downcast_dictionary_array! { + array => { + // Select dictionary entries before recursively converting nested values. + let selected = take(array.values().as_ref(), array.keys(), None) + .context(error::ComputeArrowSnafu)?; + expand_export_array(&selected, target) + }, + _ => error::UnexpectedSnafu { violated: "invalid dictionary array" }.fail(), + } + } + (DataType::List(_), DataType::List(field)) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0]; + let end = offsets[list.len()]; + let values = list.values().slice(start as usize, (end - start) as usize); + let values = expand_export_array(&values, field.data_type())?; + let offsets = OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - start) + .collect::>() + .into(), + ); + Ok(Arc::new( + ListArray::try_new(field.clone(), offsets, values, list.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + (DataType::Struct(_), DataType::Struct(fields)) => { + let array = array.as_struct(); + let columns = array + .columns() + .iter() + .zip(fields) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + Ok(Arc::new( + StructArray::try_new(fields.clone(), columns, array.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + _ => cast(array, target).context(error::ComputeArrowSnafu), + } +} + +pub(crate) fn map_writer_error( + source: common_datasource::error::Error, + path: &str, +) -> error::Error { match source { common_datasource::error::Error::ParquetWriteCancelled {} => { error::LogicalTableExportCancelledSnafu.build() @@ -572,6 +695,12 @@ fn map_writer_error(source: common_datasource::error::Error, path: &str) -> erro common_datasource::error::Error::ParquetWriterResource { reason } => { LogicalTableExportResourceSnafu { reason }.build() } + common_datasource::error::Error::WriteObject { error, .. } + if error.kind() == object_store::ErrorKind::ConditionNotMatch => + { + let reason = format!("output already exists: {path}"); + InvalidLogicalTableExportSnafu { reason }.build() + } source => error::WriteStreamToFileSnafu { path }.into_error(source), } } @@ -579,7 +708,7 @@ fn map_writer_error(source: common_datasource::error::Error, path: &str) -> erro // Count only selected logical values before dictionary expansion, including nested // histogram lists/structs. Offset and validity overhead is charged per value. fn estimate_value_size(array: &dyn Array, row: usize) -> Result { - if array.is_null(row) { + if array.is_null(row) && !matches!(array.data_type(), DataType::Struct(_) | DataType::List(_)) { return Ok(32); } let bytes = match array.data_type() { @@ -587,8 +716,10 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { DataType::Null => 0, DataType::Utf8 => array.as_string::().value(row).len(), DataType::LargeUtf8 => array.as_string::().value(row).len(), + DataType::Utf8View => array.as_string_view().value(row).len(), DataType::Binary => array.as_binary::().value(row).len(), DataType::LargeBinary => array.as_binary::().value(row).len(), + DataType::BinaryView => array.as_binary_view().value(row).len(), DataType::Struct(_) => { array .as_struct() @@ -629,18 +760,31 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { Ok(bytes.saturating_add(16)) } -fn rows_within_budget( +pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, end: usize, budget: usize, -) -> Result { + json_columns: &[usize], +) -> Result<(usize, usize)> { let mut bytes = 0usize; let mut row = start; while row < end { - let row_bytes = batch.columns().iter().try_fold(0usize, |sum, array| { - Ok::<_, error::Error>(sum.saturating_add(estimate_value_size(array.as_ref(), row)?)) - })?; + let row_bytes = + batch + .columns() + .iter() + .enumerate() + .try_fold(0usize, |sum, (index, array)| { + // JSON escaping and number formatting only expand JSON values. + let expansion = if json_columns.contains(&index) && !array.is_null(row) { + 8 + } else { + 1 + }; + let size = estimate_value_size(array.as_ref(), row)?; + Ok::<_, error::Error>(sum.saturating_add(size.saturating_mul(expansion))) + })?; if row_bytes > budget.saturating_sub(bytes) { break; } @@ -653,7 +797,7 @@ fn rows_within_budget( reason: "one expanded logical row exceeds conversion byte budget" } ); - Ok(row - start) + Ok((row - start, bytes)) } #[cfg(test)] diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 4576c1e395d..a03fdf1a80a 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -18,6 +18,7 @@ use arrow::array::{ use arrow::buffer::OffsetBuffer; use arrow::datatypes::{Field, Schema, UInt32Type}; use bytes::Bytes; +use common_error::status_code::StatusCode; use common_recordbatch::{RecordBatch as GreptimeRecordBatch, RecordBatches}; use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use table::test_util::EmptyTable; @@ -137,13 +138,14 @@ async fn routes_across_batches_and_writes_empty_files() { vec![Some("a"), Some("b"), Some("unselected")], ), ]; - let result = write_tables( + let budget = ExportWriteBudget::new(4); + let result = export_stream_managed( &unit, stream(batches), &store, export_limits(), &CancellationToken::new(), - &mut None, + budget.clone(), ) .await .unwrap(); @@ -174,6 +176,7 @@ async fn routes_across_batches_and_writes_empty_files() { let (schema, requests) = read(&store, "requests.parquet").await; assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields()); assert_eq!(requests[0].column(1).null_count(), 1); + assert_eq!(budget.available(), (4, 64 * 1024 * 1024)); } #[tokio::test] @@ -211,21 +214,16 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { ), ] { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let mut active = None; - let err = write_tables( + let err = export_stream( &unit, stream(batches), &store, limits, &CancellationToken::new(), - &mut active, ) .await .unwrap_err(); assert!(err.to_string().contains(message), "{err}"); - if let Some(writer) = active { - writer.writer.abort().await.unwrap(); - } } } @@ -233,17 +231,17 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() { async fn existing_outputs_are_not_overwritten() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); store.write("cpu.v1.parquet", "keep").await.unwrap(); - let err = write_tables( + let err = export_stream( &unit(), stream(vec![batch(vec![Some(1025)], vec![None])]), &store, LogicalTableExportLimits::default(), &CancellationToken::new(), - &mut None, ) .await .unwrap_err(); - assert!(err.to_string().contains("already exists")); + let status = common_error::ext::ErrorExt::status_code(&err); + assert_eq!(status, StatusCode::InvalidArguments); assert_eq!( store.read("cpu.v1.parquet").await.unwrap().to_bytes(), Bytes::from_static(b"keep") @@ -272,10 +270,10 @@ fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { ("histogram", Arc::new(histogram)), ]) .unwrap(); - assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap(), 1); + assert_eq!(rows_within_budget(&batch, 0, 3, 4300, &[]).unwrap().0, 1); // The dictionary and container overhead fit; the nested list elements do not. - assert!(rows_within_budget(&batch, 0, 3, 4180).is_err()); - assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap(), 3); + assert!(rows_within_budget(&batch, 0, 3, 4180, &[]).is_err()); + assert_eq!(rows_within_budget(&batch, 0, 3, 15000, &[]).unwrap().0, 3); } #[test] @@ -440,64 +438,195 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter { } #[tokio::test] -async fn cancellation_waits_for_file_creation_before_cleanup() { - use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; - let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); - let access = - common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); - let store = build_backend_for_write( - &format!("{}/", directory.path().display()), - &HashMap::new(), - &access, - ) - .await - .unwrap(); - let started = Arc::new(tokio::sync::Notify::new()); - let release = Arc::new(tokio::sync::Notify::new()); - let factory: MockWriterFactory = Arc::new({ - let started = started.clone(); - let release = release.clone(); - move |_, _, inner| { - Box::new(PausedFileWriter { - inner: Some(inner), - started: started.clone(), - release: release.clone(), - }) +async fn cancellation_drains_storage_and_preserves_committed_files() { + for large in [false, true] { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); + let access = + common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap(); + let store = build_backend_for_write( + &format!("{}/", directory.path().display()), + &HashMap::new(), + &access, + ) + .await + .unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let started = started.clone(); + let release = release.clone(); + move |_, _, inner| { + Box::new(PausedFileWriter { + inner: Some(inner), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let cancellation = CancellationToken::new(); + let (unit, input) = if large { + let field = Field::new("value", DataType::Utf8, true); + let unit = LogicalTableExport::try_new( + table( + 1024, + "phy", + vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], + true, + ), + &[table(1025, "cpu.v1", vec![field], false)], + ) + .unwrap(); + let mut state = 17u64; + let strings = (0..512) + .map(|_| { + (0..32768) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + char::from(b' ' + (state % 95) as u8) + }) + .collect::() + }) + .collect::>(); + let input = RecordBatch::try_from_iter([ + ( + TABLE_ID, + Arc::new(UInt32Array::from(vec![1025; 512])) as ArrayRef, + ), + ("value", Arc::new(StringArray::from(strings))), + ]) + .unwrap(); + (unit, input) + } else { + (unit(), batch(vec![Some(1025)], vec![Some("a")])) + }; + let budget = ExportWriteBudget::new(1); + let mut limits = LogicalTableExportLimits::default(); + limits.writer.flush_threshold_bytes = 1024 * 1024; + let export = export_stream_managed( + &unit, + stream(vec![input]), + &store, + limits, + &cancellation, + budget.clone(), + ); + tokio::pin!(export); + tokio::select! { + result = &mut export => panic!("export completed before the file open: {result:?}"), + _ = started.notified() => {}, } - }); - let store = store.layer( - MockLayerBuilder::default() - .writer_factory(factory) - .build() - .unwrap(), - ); - let cancellation = CancellationToken::new(); - let unit = unit(); - let export = export_stream( - &unit, - stream(vec![batch(vec![Some(1025)], vec![Some("a")])]), - &store, - LogicalTableExportLimits::default(), - &cancellation, - ); - tokio::pin!(export); - tokio::select! { - result = &mut export => panic!("export completed before the file open: {result:?}"), - _ = started.notified() => {}, + if large { + assert!(budget.available().1 < 64 * 1024 * 1024); + } + let held = budget.available(); + cancellation.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + assert_eq!(budget.available().0, held.0); + if large { + // Cancelled admission releases its pending reservation; the worker + // still owns payload capacity while the storage operation is paused. + assert!(budget.available().1 < 64 * 1024 * 1024); + } + release.notify_one(); + let result = export.await; + assert!(matches!( + result, + Err(error::Error::LogicalTableExportCancelled { .. }) + )); + assert_eq!(store.exists("cpu.v1.parquet").await.unwrap(), !large); + if !large { + let (_, batches) = read(&store, "cpu.v1.parquet").await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); } - cancellation.cancel(); - assert!(futures::poll!(&mut export).is_pending()); - release.notify_one(); - let result = export.await; - assert!(matches!( - result, - Err(error::Error::LogicalTableExportCancelled { .. }) - )); - assert!(!store.exists("cpu.v1.parquet").await.unwrap()); } struct FailedAbortWriter(object_store::layers::mock::oio::Writer); +#[tokio::test] +async fn metric_fallback_and_ordinary_preserve_ambiguous_commits() { + use object_store::layers::CapabilityOverrideLayer; + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + let factory: MockWriterFactory = Arc::new(|_, args, inner| { + assert!(!args.if_not_exists()); + Box::new(AmbiguousCommit(inner)) + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer(CapabilityOverrideLayer::new(|mut capability| { + capability.write_with_if_not_exists = false; + capability + })) + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + assert!(!store.info().capability().write_with_if_not_exists); + let rows = || stream(vec![batch(vec![Some(1025)], vec![Some("a")])]); + assert!( + export_stream( + &unit(), + rows(), + &store, + export_limits(), + &CancellationToken::new() + ) + .await + .is_err() + ); + let budget = ExportWriteBudget::new(1); + assert!( + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &CancellationToken::new() + ) + .await + .is_err() + ); + for path in ["cpu.v1.parquet", "ordinary.parquet"] { + let (_, batches) = read(&store, path).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); +} + impl object_store::layers::mock::oio::Write for FailedAbortWriter { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { self.0.write(bytes).await @@ -574,3 +703,243 @@ async fn validates_membership_by_table_route() { } } } + +#[tokio::test] +async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() { + let input = batch(vec![Some(1025); 8192], vec![Some("shared"); 8192]); + let tiny = input.slice(0, 1); + let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])); + let projected = tiny.project(&[2]).unwrap(); + let full = input.project(&[2]).unwrap().get_array_memory_size(); + assert_eq!(projected.get_array_memory_size(), full); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let blocker = budget.reserve(64 * 1024 * 1024, &token).await.unwrap(); + let convert = expand_bounded_slice(projected, schema, 0, 1, 1024, &budget, &token); + tokio::pin!(convert); + assert!(futures::poll!(&mut convert).is_pending()); + drop(blocker); + let (payload, rows) = convert.await.unwrap(); + assert_eq!(rows, 1); + assert!( + 64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size() + ); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err()); +} + +#[tokio::test] +async fn groups_and_ordinary_files_share_writer_admission_and_drain() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use tokio::sync::Semaphore; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct CountedWriter { + inner: oio::Writer, + active: Arc, + started: Arc, + release: Arc, + } + impl Drop for CountedWriter { + fn drop(&mut self) { + self.active.fetch_sub(1, Ordering::SeqCst); + } + } + impl oio::Write for CountedWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.started.add_permits(1); + self.release.acquire().await.unwrap().forget(); + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + for parallelism in [1, 4] { + for cancel in [false, true] { + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let started = Arc::new(Semaphore::new(0)); + let release = Arc::new(Semaphore::new(0)); + let factory: MockWriterFactory = Arc::new({ + let (active, peak, started, release) = ( + active.clone(), + peak.clone(), + started.clone(), + release.clone(), + ); + move |_, args, inner| { + assert_eq!(args.concurrent(), 1); + let now = active.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(now, Ordering::SeqCst); + Box::new(CountedWriter { + inner, + active: active.clone(), + started: started.clone(), + release: release.clone(), + }) + } + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut a = unit(); + let mut b = unit(); + for (group, unit) in [("a", &mut a), ("b", &mut b)] { + for table in unit.logical_tables.values_mut() { + table.output.path = format!("{group}/{}", table.output.path); + } + } + let rows = || { + stream(vec![batch( + vec![Some(1025), Some(1027)], + vec![Some("a"), Some("b")], + )]) + }; + let ordinary = async { + let _permit = budget.writer(&token).await?; + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &token, + ) + .await + }; + let work = async { + tokio::join!( + export_stream_managed( + &a, + rows(), + &store, + export_limits(), + &token, + budget.clone() + ), + export_stream_managed( + &b, + rows(), + &store, + export_limits(), + &token, + budget.clone() + ), + ordinary, + ) + }; + tokio::pin!(work); + tokio::select! { + _ = started.acquire_many(parallelism as u32) => {}, + result = &mut work => panic!("completed while close paused: {result:?}"), + } + assert_eq!(budget.available().0, 0); + assert_eq!(peak.load(Ordering::SeqCst), parallelism); + if cancel { + token.cancel(); + } + assert!(futures::poll!(&mut work).is_pending()); + release.add_permits(10); + let (a, b, ordinary) = work.await; + assert_eq!(a.is_err(), cancel); + assert_eq!(b.is_err(), cancel); + assert_eq!(ordinary.is_err(), cancel); + assert_eq!(active.load(Ordering::SeqCst), 0); + assert!(peak.load(Ordering::SeqCst) <= parallelism); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } + } +} + +#[tokio::test] +async fn nested_dictionary_conversion_limits_child_ranges_and_charges_null_parents() { + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 1024]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let list = ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0i32, 1023, 1024].into()), + values, + None, + ); + let input = RecordBatch::try_from_iter([("list", Arc::new(list) as ArrayRef)]).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + )])); + let budget = ExportWriteBudget::new(1); + let (payload, rows) = expand_bounded_slice( + input, + schema, + 1, + 2, + 8192, + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(rows, 1); + assert!(payload.batch.get_array_memory_size() < 16384); + let list = payload.batch.column(0).as_list::(); + assert_eq!(list.value_offsets(), &[0, 1]); + assert_eq!(list.values().as_string::().value(0), "x".repeat(4096)); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let structure = StructArray::new( + vec![Arc::new(Field::new( + "child", + values.data_type().clone(), + true, + ))] + .into(), + vec![values], + Some(arrow::buffer::NullBuffer::from(vec![false])), + ); + let input = RecordBatch::try_from_iter([("struct", Arc::new(structure) as ArrayRef)]).unwrap(); + assert!(rows_within_budget(&input, 0, 1, 128, &[]).is_err()); +} + +#[tokio::test] +async fn completed_table_workers_are_reaped_during_admission() { + for parallelism in [1, 4] { + let mut unit = unit(); + let file = unit.logical_tables.get_mut(&1025).unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + for id in 0..64 { + file.output.path = format!("empty-{id}.parquet"); + writers + .open(id, file, &store, export_limits(), &token) + .await + .unwrap(); + assert!(writers.pending_tasks() <= parallelism + 1); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(writers.pending_tasks(), 0); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs new file mode 100644 index 00000000000..38a63936a5b --- /dev/null +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -0,0 +1,310 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow::record_batch::RecordBatch; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt}; +use object_store::ObjectStore; +use snafu::{OptionExt, ResultExt, ensure}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::statement::export_logical_tables::{ + ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled, + map_writer_error, +}; + +const PAYLOAD_BYTES: usize = 64 * 1024 * 1024; + +/// Request-wide downstream payload and writer admission, separate from query memory. +pub(crate) struct ExportWriteBudget { + writers: Arc, + bytes: Arc, + max_writers: usize, +} + +impl ExportWriteBudget { + pub(crate) fn new(parallelism: usize) -> Arc { + let max_writers = parallelism.max(1); + Arc::new(Self { + writers: Arc::new(Semaphore::new(max_writers)), + bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)), + max_writers, + }) + } + + #[cfg(test)] + pub(crate) fn available(&self) -> (usize, usize) { + ( + self.writers.available_permits(), + self.bytes.available_permits(), + ) + } + + pub(crate) async fn writer(&self, token: &CancellationToken) -> Result { + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.writers.clone().acquire_owned() => permit.map_err(|_| error::UnexpectedSnafu { violated: "writer budget closed" }.build()), + } + } + + pub(crate) async fn reserve( + &self, + size: usize, + token: &CancellationToken, + ) -> Result { + ensure!( + size <= PAYLOAD_BYTES, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + permit = self.bytes.clone().acquire_many_owned(size as u32) => permit.map_err(|_| error::UnexpectedSnafu { violated: "payload budget closed" }.build()), + } + } + + /// Accounts for retained input allocations, conversion capacity growth and + /// array metadata. Pass zero backing when the caller detaches oversized slices. + pub(crate) fn conversion_budget( + backing: usize, + columns: usize, + requested: usize, + ) -> Result<(usize, usize)> { + let overhead = columns.saturating_mul(1024); + let available = PAYLOAD_BYTES + .saturating_sub(backing) + .saturating_sub(overhead) + / 4; + let conversion = requested.min(available); + ensure!( + conversion > 0, + error::LogicalTableExportResourceSnafu { + reason: "batch backing buffers exceed request payload budget" + } + ); + Ok((conversion, backing.saturating_add(overhead))) + } +} + +pub(crate) struct Payload { + pub(crate) batch: RecordBatch, + pub(crate) permit: OwnedSemaphorePermit, +} + +pub(crate) struct TableWriters { + current: Option<(u32, mpsc::Sender)>, + tasks: FuturesUnordered>>, + budget: Arc, +} + +impl TableWriters { + pub(crate) fn new(budget: Arc) -> Self { + Self { + current: None, + tasks: FuturesUnordered::new(), + budget, + } + } + + #[cfg(test)] + pub(crate) fn pending_tasks(&self) -> usize { + self.tasks.len() + } + + pub(crate) fn table_id(&self) -> Option { + self.current.as_ref().map(|(id, _)| *id) + } + + pub(crate) fn close_input(&mut self) { + self.current = None; + } + + pub(crate) async fn open( + &mut self, + id: u32, + table: &LogicalTableProjection, + store: &ObjectStore, + limits: LogicalTableExportLimits, + token: &CancellationToken, + ) -> Result<()> { + // EOF must precede acquiring the next slot, including when P is one. + self.close_input(); + let permit = self.budget.writer(token).await?; + self.reap_for_admission(token).await?; + let writer = ActiveWriter::open(table, store, limits).await?; + let (sender, receiver) = mpsc::channel(2); + self.current = Some((id, sender)); + let token = token.clone(); + self.tasks.push(common_runtime::spawn_global(async move { + let _permit = permit; + let guard = token.clone().drop_guard(); + let result = run_writer(writer, receiver, &token).await; + if result.is_ok() { + guard.disarm(); + } + result + })); + Ok(()) + } + + async fn reap_for_admission(&mut self, token: &CancellationToken) -> Result<()> { + while let Some(Some(result)) = self.tasks.next().now_or_never() { + result.context(error::JoinTaskSnafu)??; + } + // A worker can release its permit before its JoinHandle becomes ready. + while self.tasks.len() > self.budget.max_writers { + let result = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + result = self.tasks.next() => result.context(error::UnexpectedSnafu { violated: "writer task queue unexpectedly empty" })?, + }; + result.context(error::JoinTaskSnafu)??; + } + Ok(()) + } + + pub(crate) async fn send(&self, payload: Payload, token: &CancellationToken) -> Result<()> { + let (_, sender) = self.current.as_ref().context(error::UnexpectedSnafu { + violated: "missing logical writer", + })?; + tokio::select! { + biased; + _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(), + result = sender.send(payload) => result.map_err(|_| error::LogicalTableExportCancelledSnafu.build()), + } + } + + pub(crate) async fn drain( + &mut self, + result: Result<()>, + token: &CancellationToken, + ) -> Result<()> { + self.close_input(); + let mut first_error = result.err(); + if first_error.is_some() { + token.cancel(); + } + while let Some(result) = self.tasks.next().await { + if let Err(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) { + retain_error(&mut first_error, err); + token.cancel(); + } + } + match first_error { + Some(err) => Err(err), + None => Ok(()), + } + } +} + +pub(crate) fn retain_error(first: &mut Option, error: error::Error) { + if first.as_ref().is_none_or(|err| { + matches!( + err, + error::Error::LogicalTableExportCancelled { .. } + | error::Error::DatabaseExportCancelled { .. } + ) + }) { + *first = Some(error); + } +} + +async fn run_writer( + mut writer: ActiveWriter, + mut receiver: mpsc::Receiver, + token: &CancellationToken, +) -> Result<()> { + let result = async { + loop { + let payload = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + payload = receiver.recv() => payload, + }; + let Some(Payload { batch, permit }) = payload else { + break; + }; + let result = writer.writer.write(batch, Some(token)).await; + drop(permit); + result.map_err(|error| map_writer_error(error, &writer.path))?; + } + writer + .writer + .finish(Some(token)) + .await + .map_err(|error| map_writer_error(error, &writer.path))?; + check_cancelled(token) + } + .await; + if result.is_err() { + token.cancel(); + receiver.close(); + drop(receiver); + if let Err(error) = writer.writer.abort().await { + common_telemetry::warn!(error; "Failed to abort Metric export file"); + } + } + result +} + +#[cfg(test)] +mod tests { + use tokio::sync::oneshot; + + use super::*; + + #[tokio::test] + async fn released_permit_does_not_hide_pending_join_handles() { + for parallelism in [1, 4] { + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + let mut resume = Vec::new(); + for _ in 0..=parallelism { + let permit = budget.writer(&token).await.unwrap(); + let (released_tx, released_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + writers.tasks.push(common_runtime::spawn_global(async move { + drop(permit); + released_tx.send(()).unwrap(); + resume_rx.await.unwrap(); + Ok(()) + })); + released_rx.await.unwrap(); + resume.push(resume_tx); + } + assert_eq!(budget.available().0, parallelism); + + let mut reap = Box::pin(writers.reap_for_admission(&token)); + assert!(reap.as_mut().now_or_never().is_none()); + resume.pop().unwrap().send(()).unwrap(); + reap.await.unwrap(); + assert_eq!(writers.pending_tasks(), parallelism); + + for tx in resume { + tx.send(()).unwrap(); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(budget.available().0, parallelism); + } + } +} diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 41fe8382c82..1f9161d970f 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -228,7 +228,7 @@ fn database_export_request(directory: &std::path::Path) -> table::requests::Copy } } -async fn database_export_roundtrip(instance: &Arc) { +async fn database_export_roundtrip(instance: &Arc, parallelism: usize) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); let (first_logical_table_names, _, renamed_physical_table) = create_metric_export_source_tables(instance, "db_a", "dense").await; @@ -253,7 +253,9 @@ async fn database_export_roundtrip(instance: &Arc) { ]; let mut names = selected.clone(); names.extend([renamed_physical_table, "dashboard".into()]); - let req = database_export_request(&destination.path().join("data")); + let mut req = database_export_request(&destination.path().join("data")); + req.with + .insert("parallelism".into(), parallelism.to_string()); let executor = instance.statement_executor(); let captured = executor .capture_database_export_tables(&req, None, &QueryContext::arc()) @@ -448,7 +450,7 @@ async fn database_export_standalone_roundtrip() { let standalone = GreptimeDbStandaloneBuilder::new("database_export") .build() .await; - database_export_roundtrip(standalone.fe_instance()).await; + database_export_roundtrip(standalone.fe_instance(), 1).await; } #[tokio::test(flavor = "multi_thread")] @@ -464,7 +466,7 @@ async fn database_export_distributed_roundtrip() { ) .build(false) .await; - database_export_roundtrip(cluster.fe_instance()).await; + database_export_roundtrip(cluster.fe_instance(), 4).await; } #[tokio::test(flavor = "multi_thread")] From 4ffd93d8dc0e11dec3602a39df61338686fe4728 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 28 Sep 2026 03:35:36 +0000 Subject: [PATCH 045/133] fix(query): preserve global limits with DataFusion optimizer fixes (#9071) * fix(query): preserve global fetch through physical-optimizer repartitioning Aggregate soft-limit pushdown inserts a global CoalescePartitionsExec(fetch) below an AggregateExec whose input is HashPartitioned/KeyPartitioned. The DataFusion physical optimizer's EnforceDistribution/EnforceSorting passes then dropped or rewrote that fetched coalesce, losing the global limit. Root-caused and fixed in the DataFusion fork (GreptimeTeam/datafusion PR #35); this repo pins to that fix commit and adds regression coverage. - pin datafusion to GreptimeTeam/datafusion 48510c7f5 (fetch preservation fix) - global_limit.rs: adapt to new DF API (input_distribution_requirements, child_distribution, replace_children+Recompute); accept HashPartitioned and KeyPartitioned as partitioning to restore; add regression unit tests - sqlness: extend standalone common ssts and refresh distributed ssts_limit result for the preserved fetch plan Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor: drop dead HashPartitioned compat arms and test wrapper - remove HashPartitioned compatibility matches (pinned DF only emits KeyPartitioned); drop their #[expect(deprecated)] attributes - remove single-use agg_with_limit helper whose seed limit is always overwritten by the soft-limit transform - restore Cargo.lock dependency edges unrelated to the DataFusion pin bump (cargo update --precise had downgraded 26 unrelated packages) Addresses oracle ablation review. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore(deps): pin DataFusion to merged fetch preservation fix Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 70 +++--- Cargo.toml | 30 +-- src/query/src/optimizer/global_limit.rs | 229 +++++++++++++++++- .../information_schema/ssts_limit.result | 43 ++-- .../common/information_schema/ssts.result | 14 ++ .../common/information_schema/ssts.sql | 8 + 6 files changed, 321 insertions(+), 73 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2cae7d03322..597ad270305 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3672,7 +3672,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3724,7 +3724,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3748,7 +3748,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3771,7 +3771,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3797,7 +3797,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "futures", "log", @@ -3807,7 +3807,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-compression", @@ -3843,7 +3843,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ipc 59.2.0", @@ -3867,7 +3867,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3890,7 +3890,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -3913,7 +3913,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3945,12 +3945,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" [[package]] name = "datafusion-execution" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -3975,7 +3975,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -3999,7 +3999,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4010,7 +4010,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-buffer 59.2.0", @@ -4041,7 +4041,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4061,7 +4061,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4072,7 +4072,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-ord 59.2.0", @@ -4096,7 +4096,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "async-trait", @@ -4111,7 +4111,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4127,7 +4127,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4136,7 +4136,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-doc", "quote", @@ -4146,7 +4146,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "chrono", @@ -4194,7 +4194,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -4217,7 +4217,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4231,7 +4231,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "chrono", @@ -4248,7 +4248,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4267,7 +4267,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "arrow-data 59.2.0", @@ -4304,7 +4304,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-catalog", @@ -4330,7 +4330,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4340,7 +4340,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -4350,7 +4350,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "datafusion-common", @@ -4365,7 +4365,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow-schema 59.2.0", "async-trait", @@ -4379,7 +4379,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "arrow 59.2.0", "bigdecimal 0.4.8", @@ -4398,7 +4398,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=2aa87d52cdce7006af492330064738f33ed294c1#2aa87d52cdce7006af492330064738f33ed294c1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4bd567a884dcbfd89d0f0a7859ef8367e38fff97#4bd567a884dcbfd89d0f0a7859ef8367e38fff97" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 5dc37afceba..e057db65826 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -363,21 +363,21 @@ rev = "f743c3a5fe8c1f57363eac5866c356b7d91f9b39" [patch.crates-io] substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "91ec978b0649417ad3da8390e7a515baec723b1b" } -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "2aa87d52cdce7006af492330064738f33ed294c1" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" } datafusion-pg-catalog = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } arrow-pg = { git = "https://github.com/datafusion-contrib/datafusion-postgres.git", rev = "ffb14a52b3d7c7489559812e1ec38096f225ce93" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "9e9019bb1c7040ed956f654e39378dd42ab17884" } diff --git a/src/query/src/optimizer/global_limit.rs b/src/query/src/optimizer/global_limit.rs index 5abb9953b7b..b29907f01ab 100644 --- a/src/query/src/optimizer/global_limit.rs +++ b/src/query/src/optimizer/global_limit.rs @@ -254,22 +254,215 @@ fn inherited_partitioning_to_restore( #[cfg(test)] mod tests { - use datafusion::arrow::array::Int32Array; + use datafusion::arrow::array::{Array, Int32Array}; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; + use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::{col, lit}; + use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; + use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, + }; use datafusion::physical_plan::filter::FilterExecBuilder; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; - use datafusion::physical_plan::limit::GlobalLimitExec; + use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::test::TestMemoryExec; use datafusion_common::{JoinType, NullEquality}; use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr}; use super::*; + async fn optimize_and_collect_twice( + mut plan: Arc, + config: &ConfigOptions, + ) -> Vec> { + let mut results = Vec::with_capacity(2); + for _ in 0..2 { + for rule in PhysicalOptimizer::new().rules { + plan = rule.optimize(plan, config).unwrap(); + } + let batches = datafusion::physical_plan::collect( + Arc::clone(&plan), + Arc::new(TaskContext::default()), + ) + .await + .unwrap(); + results.push( + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(), + ); + } + results + } + + #[tokio::test] + async fn physical_optimizer_keeps_global_distinct_limit_across_two_passes() { + for soft_limit in [false, true] { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + config.optimizer.enable_distinct_aggregation_soft_limit = soft_limit; + + for mode in [ + AggregateMode::FinalPartitioned, + AggregateMode::SinglePartitioned, + ] { + // The query limit is the only initial limit. + let input = input_with_all_hash_partitions(); + let aggregate = match mode { + AggregateMode::FinalPartitioned => { + agg(agg(input, AggregateMode::Partial), mode) + } + AggregateMode::SinglePartitioned => agg(hash_repartition(input), mode), + _ => unreachable!(), + }; + let plan = + Arc::new(GlobalLimitExec::new(aggregate, 0, Some(1))) as Arc; + + let results = optimize_and_collect_twice(plan, &config).await; + assert_eq!( + results.iter().map(Vec::len).collect::>(), + vec![1, 1], + "soft limit enabled: {soft_limit}, mode: {mode:?}", + ); + } + } + } + + #[tokio::test] + async fn physical_optimizer_keeps_count_over_distinct_limit_across_two_passes() { + use datafusion::datasource::MemTable; + use datafusion::execution::context::{SessionConfig, SessionContext as DFSessionContext}; + use datafusion_common::ScalarValue; + + for soft_limit in [false, true] { + let mut config = SessionConfig::new().with_target_partitions(3); + config + .options_mut() + .optimizer + .enable_distinct_aggregation_soft_limit = soft_limit; + let ctx = DFSessionContext::new_with_config(config.clone()); + let schema = schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + let table = MemTable::try_new(schema, vec![vec![batch]; 3]).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + let mut plan = ctx + .sql("SELECT COUNT(*) FROM (SELECT DISTINCT a FROM t LIMIT 1) AS limited") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + for pass in 0..=2 { + if pass > 0 { + for rule in PhysicalOptimizer::new().rules { + plan = rule.optimize(plan, config.options()).unwrap(); + } + } + let batches = datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(|row| { + ScalarValue::try_from_array(batch.column(0).as_ref(), row).unwrap() + }) + }) + .collect::>(); + assert_eq!( + values, + vec![ScalarValue::Int64(Some(1))], + "soft limit enabled: {soft_limit}, additional optimizer passes: {pass}", + ); + } + } + } + + #[tokio::test] + async fn physical_optimizer_keeps_local_limit_per_partition_across_two_passes() { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + let input = input_with_all_hash_partitions(); + let local_limit = Arc::new(LocalLimitExec::new( + agg(hash_repartition(input), AggregateMode::SinglePartitioned), + 1, + )) as Arc; + + let results = optimize_and_collect_twice(local_limit, &config).await; + assert_eq!(results.iter().map(Vec::len).collect::>(), vec![3, 3]); + } + + #[tokio::test] + async fn physical_optimizer_keeps_ordered_topk_across_two_passes() { + let schema = schema(); + let losing = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![100]))]) + .unwrap(); + let winning = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + let input = Arc::new( + TestMemoryExec::try_new(&[vec![losing], vec![winning]], schema.clone(), None).unwrap(), + ); + let ordering = ordering(schema.as_ref(), false); + let topk = Arc::new( + SortExec::new( + ordering, + agg_with_limit_options( + input, + AggregateMode::FinalPartitioned, + Some(LimitOptions::new_with_order(1, false)), + ), + ) + .with_fetch(Some(1)), + ) as Arc; + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + assert_eq!( + optimize_and_collect_twice(topk, &config).await, + vec![vec![1], vec![1]], + ); + } + + #[tokio::test] + async fn physical_optimizer_keeps_ordered_offset_limit_across_two_passes() { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 3; + let offset = Arc::new(GlobalLimitExec::new( + Arc::new(SortExec::new( + ordering(schema().as_ref(), false), + unordered_input(), + )), + 3, + Some(1), + )) as Arc; + + assert_eq!( + optimize_and_collect_twice(offset, &config).await, + vec![vec![2], vec![2]], + ); + } + #[test] fn adds_global_limit_for_multi_partition_filter_fetch() { let filter = filter_fetch(unordered_input(), 1); @@ -563,6 +756,38 @@ mod tests { Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap()) } + fn input_with_all_hash_partitions() -> Arc { + let schema = schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]]; + Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap()) + } + + fn agg(input: Arc, mode: AggregateMode) -> Arc { + agg_with_limit_options(input, mode, None) + } + + fn agg_with_limit_options( + input: Arc, + mode: AggregateMode, + limit_options: Option, + ) -> Arc { + let schema = input.schema(); + let group_by = PhysicalGroupBy::new_single(vec![( + col("a", schema.as_ref()).unwrap(), + "a".to_string(), + )]); + Arc::new( + AggregateExec::try_new(mode, group_by, vec![], vec![], input, schema) + .unwrap() + .with_limit_options(limit_options), + ) + } + fn ordered_input() -> (Arc, LexOrdering) { let schema = schema(); let ordering = ordering(schema.as_ref(), false); diff --git a/tests/cases/distributed/information_schema/ssts_limit.result b/tests/cases/distributed/information_schema/ssts_limit.result index 248593410fc..1b21a84c7af 100644 --- a/tests/cases/distributed/information_schema/ssts_limit.result +++ b/tests/cases/distributed/information_schema/ssts_limit.result @@ -72,27 +72,28 @@ FROM ( LIMIT 1 ); -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | -| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | -| | Limit: skip=0, fetch=1 | -| | Projection: information_schema.ssts_manifest.region_id | -| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | -| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | -| | ]] | -| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | -| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec | -| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | -| | CoalescePartitionsExec: fetch=1 | -| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | -| | RepartitionExec: REDACTED -| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [2], filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | -| | | -+---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: count(Int64(1)) AS count(*) AS filtered_limited_rows | +| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] | +| | Limit: skip=0, fetch=1 | +| | Projection: information_schema.ssts_manifest.region_id | +| | Filter: information_schema.ssts_manifest.table_id > UInt32(0) | +| | TableScan: information_schema.ssts_manifest, partial_filters=[information_schema.ssts_manifest.table_id > UInt32(0)] | +| | ]] | +| physical_plan | ProjectionExec: expr=[count(Int64(1))@0 as filtered_limited_rows] | +| | AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] | +| | CoalescePartitionsExec | +| | AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] | +| | RepartitionExec: REDACTED +| | CoalescePartitionsExec: fetch=1 | +| | FilterExec: table_id@0 > 0, projection=[], fetch=1 | +| | RepartitionExec: REDACTED +| | DistributedInspectExec: kind=SstManifest, scan=ScanRequest { projection: [2], filters: [table_id > UInt32(0)] }, schema=Schema { fields: [Field { name: "table_id", data_type: UInt32 }], metadata: {"greptime:version": "0"} } | +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ DROP TABLE ssts_limit_case; diff --git a/tests/cases/standalone/common/information_schema/ssts.result b/tests/cases/standalone/common/information_schema/ssts.result index 9777e3b12a9..289b4599387 100644 --- a/tests/cases/standalone/common/information_schema/ssts.result +++ b/tests/cases/standalone/common/information_schema/ssts.result @@ -229,6 +229,20 @@ SELECT * FROM information_schema.ssts_storage order by file_path; | data/greptime/public//_/index/.puffin || || +---------------------------------------------------------------------------------------------+-----------+-------------------------+---------+ +-- Regression: repeated physical optimization must preserve the global DISTINCT limit. +SELECT COUNT(*) AS distinct_limited_rows +FROM ( + SELECT DISTINCT region_id + FROM information_schema.ssts_manifest + LIMIT 1 +); + ++-----------------------+ +| distinct_limited_rows | ++-----------------------+ +| 1 | ++-----------------------+ + DROP TABLE sst_case; Affected Rows: 0 diff --git a/tests/cases/standalone/common/information_schema/ssts.sql b/tests/cases/standalone/common/information_schema/ssts.sql index 43512de6808..7b40b0056bb 100644 --- a/tests/cases/standalone/common/information_schema/ssts.sql +++ b/tests/cases/standalone/common/information_schema/ssts.sql @@ -72,4 +72,12 @@ SELECT * FROM information_schema.ssts_index_meta ORDER BY meta_json; -- SQLNESS REPLACE (/public/\d+/\d+_\d+) /public//_ SELECT * FROM information_schema.ssts_storage order by file_path; +-- Regression: repeated physical optimization must preserve the global DISTINCT limit. +SELECT COUNT(*) AS distinct_limited_rows +FROM ( + SELECT DISTINCT region_id + FROM information_schema.ssts_manifest + LIMIT 1 +); + DROP TABLE sst_case; From d5a6a9ac0513ccedde012c410a2c0c5abec624dc Mon Sep 17 00:00:00 2001 From: DeviousCardi <115358213+DeviousCardi@users.noreply.github.com> Date: Mon, 28 Sep 2026 04:33:46 +0000 Subject: [PATCH 046/133] fix(auth): follow symlink chains in watch_file_user_provider (#9365) * fix(auth): follow symlink chains in watch_file_user_provider The shared file watcher only watched the immediate parent directory of the configured path, so swapping a directory symlink elsewhere in the chain (e.g. `ln -sfn secrets.2 secrets`, as sops-nix does) never produced an event and the new users file was not loaded until restart. The watcher now resolves each path component by component and also watches the directory holding every symlink it meets, plus the directory of the final file. After each relevant event the chain is re-resolved and the watches are re-armed, so later in-place edits of the new target are seen. Only the final file's directory is required at startup; a link directory that cannot be watched (e.g. mode 0711) is logged and retried later. Events are now filtered to the chain's links and final files, so that watching e.g. /tmp does not reload on unrelated files. Names are compared case-insensitively on macOS and Windows. TLS cert reload uses the same helper. Close #9310 Signed-off-by: Aarav Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bvab9gpXYjKyLN8Mo6JWWL * fix(common-config): warn when the file watcher watches the root dir On macOS /tmp, /var and /etc are symlinks under /, so following a symlink chain there adds / to the watched directories, and the FSEvents backend then sees every event on the volume. Log a warning when that happens. Also shorten the comment on skipped link directories. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bvab9gpXYjKyLN8Mo6JWWL Signed-off-by: Aarav --------- Signed-off-by: Aarav Co-authored-by: Claude Opus 5.5 --- .../user_provider/watch_file_user_provider.rs | 42 ++ src/common/config/src/file_watcher.rs | 500 +++++++++++++++--- 2 files changed, 473 insertions(+), 69 deletions(-) diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs index 3208948b11d..bba97f6df1f 100644 --- a/src/auth/src/user_provider/watch_file_user_provider.rs +++ b/src/auth/src/user_provider/watch_file_user_provider.rs @@ -199,4 +199,46 @@ pub mod test { test_authenticate(&provider, "root", "654321", false, Some(timeout)).await; test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await; } + + /// Regression test for #9310: the users file is reached through a chain + /// of symlinks and a directory symlink in the chain gets swapped + /// atomically (like `ln -sfn`, as done by sops-nix). + #[cfg(unix)] + #[tokio::test] + async fn test_file_provider_symlink_chain() { + use std::os::unix::fs::symlink; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_provider_symlink_chain"); + let base = dir.path(); + let timeout = Duration::from_secs(10); + + // secrets.1/users <- secrets (dir symlink) <- conf/users (file symlink) + std::fs::create_dir(base.join("secrets.1")).unwrap(); + std::fs::write(base.join("secrets.1/users"), "root=123456\n").unwrap(); + symlink(base.join("secrets.1"), base.join("secrets")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink(base.join("secrets/users"), base.join("conf/users")).unwrap(); + + let file_path = base.join("conf/users"); + let provider = WatchFileUserProvider::new(file_path.to_str().unwrap()).unwrap(); + test_authenticate(&provider, "root", "123456", true, None).await; + + // Atomically swap the directory symlink to a new generation. + std::fs::create_dir(base.join("secrets.2")).unwrap(); + std::fs::write(base.join("secrets.2/users"), "root=654321\n").unwrap(); + symlink(base.join("secrets.2"), base.join("secrets.tmp")).unwrap(); + std::fs::rename(base.join("secrets.tmp"), base.join("secrets")).unwrap(); + std::fs::remove_dir_all(base.join("secrets.1")).unwrap(); + + test_authenticate(&provider, "root", "654321", true, Some(timeout)).await; + test_authenticate(&provider, "root", "123456", false, None).await; + + // In-place edits of the new target are picked up too, which requires + // the watch to follow the chain to the new directory. + std::fs::write(base.join("secrets.2/users"), "root=abcdef\n").unwrap(); + test_authenticate(&provider, "root", "abcdef", true, Some(timeout)).await; + test_authenticate(&provider, "root", "654321", false, None).await; + } } diff --git a/src/common/config/src/file_watcher.rs b/src/common/config/src/file_watcher.rs index d0766920191..da2c84a8f60 100644 --- a/src/common/config/src/file_watcher.rs +++ b/src/common/config/src/file_watcher.rs @@ -21,9 +21,16 @@ //! file itself. This ensures that file deletions and recreations are properly //! tracked, which is common with editors that use atomic saves or when //! configuration files are replaced. +//! +//! If the path goes through symlinks (the file itself or any ancestor +//! directory), the directory holding each symlink is watched as well, so that +//! swapping a symlink anywhere in the chain (e.g. `ln -sfn secrets.2 secrets`, +//! as done by secret managers like sops-nix or Kubernetes volume mounts) is +//! detected. The chain is re-resolved after every relevant event and the set +//! of watched directories is updated accordingly. -use std::collections::HashSet; -use std::path::{Path, PathBuf}; +use std::collections::{HashSet, VecDeque}; +use std::path::{Component, Path, PathBuf}; use std::sync::mpsc::channel; use common_telemetry::{error, info, warn}; @@ -54,10 +61,12 @@ impl FileWatcherConfig { /// /// The watcher monitors the parent directory of each file to handle file /// deletion and recreation properly. Events are filtered to only trigger -/// callbacks for the specific files being watched. +/// callbacks for the specific files being watched (or the symlinks leading +/// to them). pub struct FileWatcherBuilder { config: FileWatcherConfig, - /// Canonicalized paths of files to watch. + /// Paths of files to watch, as configured. Symlinks in them are resolved + /// when spawning and again after every relevant event. file_paths: Vec, } @@ -79,7 +88,7 @@ impl FileWatcherBuilder { /// Add a file path to watch. /// /// Returns an error if the path is a directory. - /// The path is canonicalized for reliable comparison with events. + /// Symlinks in the path are followed, see [`FileWatcherBuilder::spawn`]. pub fn watch_path>(mut self, path: P) -> Result { let path = path.as_ref(); snafu::ensure!( @@ -110,7 +119,10 @@ impl FileWatcherBuilder { /// /// The callback is invoked when relevant file events are detected for /// the watched files. The watcher monitors the parent directories to - /// handle file deletion and recreation properly. + /// handle file deletion and recreation properly. When a path goes through + /// symlinks, the directory holding each symlink is watched too, and the + /// watches are re-armed after every relevant event, so the callback also + /// fires when a symlink anywhere in the chain is replaced. /// /// The spawned watcher thread runs for the lifetime of the process. pub fn spawn(self, callback: F) -> Result<()> @@ -121,42 +133,61 @@ impl FileWatcherBuilder { let mut watcher = notify::recommended_watcher(tx).context(FileWatchSnafu { path: "" })?; - // Collect unique parent directories to watch - let mut watched_dirs: HashSet = HashSet::new(); - for file_path in &self.file_paths { - if let Some(parent) = file_path.parent() - && watched_dirs.insert(parent.to_path_buf()) - { - watcher - .watch(parent, RecursiveMode::NonRecursive) - .context(FileWatchSnafu { - path: parent.display().to_string(), - })?; + let mut targets = resolve_watch_targets(&self.file_paths); + let mut skipped = Vec::new(); + for dir in &targets.dirs { + match watcher.watch(dir, RecursiveMode::NonRecursive) { + Ok(()) => warn_if_root(dir), + // Only the directories holding the files are required. A + // directory holding a symlink may be unreadable (e.g. mode + // 0711); it is skipped and retried by `rearm_watches`. + Err(err) if !targets.file_dirs.contains(dir) => { + warn!("Failed to watch {:?}, skipping: {}", dir, err); + skipped.push(dir.clone()); + } + Err(err) => { + return Err(err).context(FileWatchSnafu { + path: dir.display().to_string(), + }); + } } } + for dir in skipped { + targets.dirs.remove(&dir); + } let config = self.config; + let file_paths = self.file_paths; info!( - "Spawning file watcher for paths: {:?} (watching parent directories)", - self.file_paths + "Spawning file watcher for paths: {:?} (watching directories: {:?})", + file_paths .iter() .map(|p| p.display().to_string()) - .collect::>() + .collect::>(), + targets.dirs ); std::thread::spawn(move || { // Keep watcher alive in the thread - let _watcher = watcher; + let mut watcher = watcher; while let Ok(res) = rx.recv() { match res { Ok(event) => { - if !is_relevant_event(&event.kind, &config) { + if !is_relevant_event(&event.kind, &config) + || !event + .paths + .iter() + .any(|p| targets.entries.contains(&entry_key(p))) + { continue; } info!(?event.kind, ?event.paths, "Detected folder change"); + // A symlink in the chain may have been replaced, so + // re-resolve and re-arm the watches before reloading. + targets = rearm_watches(&mut watcher, &file_paths, targets); callback(); } Err(err) => { @@ -178,6 +209,166 @@ impl Default for FileWatcherBuilder { } } +/// Maximum number of symlinks followed while resolving a single path. This +/// mirrors Linux's `MAXSYMLINKS` and guards against symlink loops. +const MAX_SYMLINK_HOPS: usize = 40; + +/// Directories to watch and the directory entries whose events are relevant. +#[derive(Debug, Default)] +struct WatchTargets { + /// Directories holding a symlink of a chain, or the final file. + dirs: HashSet, + /// The subset of `dirs` holding the final files. + file_dirs: HashSet, + /// The symlinks of every chain and the final files, as they appear + /// inside `dirs`, normalized with [`entry_key`]. + entries: HashSet, +} + +/// Key used to compare an event path with [`WatchTargets::entries`]. +/// +/// The entries are built from the configured spelling, which may differ in +/// case from what is on disk (and reported in events) on the usually +/// case-insensitive filesystems of macOS and Windows, so compare lowercased +/// there. At worst this lets a few extra events through on a case-sensitive +/// volume. +fn entry_key(path: &Path) -> PathBuf { + #[cfg(any(target_os = "macos", windows))] + { + PathBuf::from(path.to_string_lossy().to_lowercase()) + } + #[cfg(not(any(target_os = "macos", windows)))] + { + path.to_path_buf() + } +} + +/// Resolves every path component by component, following symlinks, and +/// collects what has to be watched to notice any change of what the path +/// points to. +/// +/// For each symlink met on the way (the file itself or any ancestor +/// directory), the directory holding it is watched, so replacing the symlink +/// is noticed. The directory holding the final file is watched too, so +/// in-place edits are noticed. The collected directories are free of +/// symlinks; the collected entries are the symlinks themselves and the final +/// files, located inside those directories. +/// +/// If a component is missing (e.g. a symlink being swapped non-atomically), +/// resolution of that path stops and its containing directory is watched, so +/// the watch is re-armed once the component shows up again. +fn resolve_watch_targets(paths: &[PathBuf]) -> WatchTargets { + let mut targets = WatchTargets::default(); + for path in paths { + resolve_one(path, &mut targets); + } + targets +} + +fn resolve_one(path: &Path, targets: &mut WatchTargets) { + let Ok(path) = std::path::absolute(path) else { + return; + }; + + let mut pending: VecDeque = path + .components() + .map(|c| PathBuf::from(c.as_os_str())) + .collect(); + // Always free of symlinks, since each component is resolved before + // being appended. + let mut resolved = PathBuf::new(); + let mut hops = 0; + + while let Some(component) = pending.pop_front() { + match component.components().next() { + Some(Component::Prefix(_) | Component::RootDir) => resolved.push(&component), + Some(Component::CurDir) | None => {} + Some(Component::ParentDir) => { + resolved.pop(); + } + Some(Component::Normal(name)) => { + let candidate = resolved.join(name); + let is_last = pending.is_empty(); + match std::fs::symlink_metadata(&candidate) { + Ok(meta) if meta.file_type().is_symlink() => { + targets.dirs.insert(resolved.clone()); + targets.entries.insert(entry_key(&candidate)); + hops += 1; + let Ok(link) = std::fs::read_link(&candidate) else { + return; + }; + if hops > MAX_SYMLINK_HOPS { + warn!("Too many levels of symlinks while resolving {:?}", path); + return; + } + // Relative targets resolve against `resolved`; absolute + // ones reset it through their root component. + for c in link.components().rev() { + pending.push_front(PathBuf::from(c.as_os_str())); + } + } + Ok(_) if !is_last => resolved = candidate, + Ok(_) | Err(_) => { + if is_last { + targets.file_dirs.insert(resolved.clone()); + } + targets.dirs.insert(resolved); + targets.entries.insert(entry_key(&candidate)); + return; + } + } + } + } + } +} + +/// Re-resolves `file_paths` and updates the watched directories, returning +/// the new targets. Failures are logged and retried on the next event that +/// matches a chain entry. +fn rearm_watches( + watcher: &mut impl Watcher, + file_paths: &[PathBuf], + old: WatchTargets, +) -> WatchTargets { + let mut new = resolve_watch_targets(file_paths); + + for dir in old.dirs.difference(&new.dirs) { + // The directory may already be gone; nothing to do then. + let _ = watcher.unwatch(dir); + } + let mut failed = Vec::new(); + for dir in new.dirs.difference(&old.dirs) { + match watcher.watch(dir, RecursiveMode::NonRecursive) { + Ok(()) => { + info!("File watcher now watching {:?}", dir); + warn_if_root(dir); + } + Err(err) => { + warn!("Failed to watch {:?}: {}", dir, err); + failed.push(dir.clone()); + } + } + } + for dir in failed { + new.dirs.remove(&dir); + } + new +} + +/// Warns when the root directory is watched, which happens when a path goes +/// through a symlink directly under it (e.g. `/tmp`, `/var` or `/etc` on +/// macOS). On macOS the FSEvents backend then receives every event on the +/// volume and filters them in user space. +fn warn_if_root(dir: &Path) { + if dir.parent().is_none() { + warn!( + "File watcher is watching the root directory {:?} because a watched path goes \ + through a symlink in it; on macOS this receives every event on the volume", + dir + ); + } +} + /// Check if an event kind is relevant based on the configuration. fn is_relevant_event(kind: &EventKind, config: &FileWatcherConfig) -> bool { match kind { @@ -189,47 +380,75 @@ fn is_relevant_event(kind: &EventKind, config: &FileWatcherConfig) -> bool { #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; + use std::sync::{Arc, Mutex}; + use std::time::{Duration, Instant}; use common_test_util::temp_dir::create_temp_dir; use super::*; + const TIMEOUT: Duration = Duration::from_secs(10); + + /// Polls `cond` until it holds or `TIMEOUT` elapses. + fn wait_until(cond: impl Fn() -> bool) -> bool { + let deadline = Instant::now() + TIMEOUT; + while Instant::now() < deadline { + if cond() { + return true; + } + std::thread::sleep(Duration::from_millis(20)); + } + cond() + } + + /// Spawns a watcher on `path` and returns a counter of callback calls. + fn spawn_counter(path: &Path) -> Arc { + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = counter.clone(); + FileWatcherBuilder::new() + .watch_path(path) + .unwrap() + .config(FileWatcherConfig::new()) + .spawn(move || { + counter_clone.fetch_add(1, Ordering::SeqCst); + }) + .unwrap(); + counter + } + + /// Spawns a watcher on `path` whose callback stores the content read + /// through `path`, like the real reload callbacks do. + fn spawn_reader(path: &Path) -> Arc> { + let content = Arc::new(Mutex::new(String::new())); + let content_clone = content.clone(); + let path_owned = path.to_path_buf(); + FileWatcherBuilder::new() + .watch_path(path) + .unwrap() + .spawn(move || { + if let Ok(c) = std::fs::read_to_string(&path_owned) { + *content_clone.lock().unwrap() = c; + } + }) + .unwrap(); + content + } + #[test] fn test_file_watcher_detects_changes() { common_telemetry::init_default_ut_logging(); let dir = create_temp_dir("test_file_watcher"); let file_path = dir.path().join("test_file.txt"); - - // Create initial file std::fs::write(&file_path, "initial content").unwrap(); - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); - - FileWatcherBuilder::new() - .watch_path(&file_path) - .unwrap() - .config(FileWatcherConfig::new()) - .spawn(move || { - counter_clone.fetch_add(1, Ordering::SeqCst); - }) - .unwrap(); - - // Give watcher time to start - std::thread::sleep(Duration::from_millis(100)); - - // Modify the file + // Watches are registered before `spawn` returns, so no need to wait. + let counter = spawn_counter(&file_path); std::fs::write(&file_path, "modified content").unwrap(); - // Wait for the event to be processed - std::thread::sleep(Duration::from_millis(500)); - assert!( - counter.load(Ordering::SeqCst) >= 1, + wait_until(|| counter.load(Ordering::SeqCst) >= 1), "Watcher should have detected at least one change" ); } @@ -240,38 +459,181 @@ mod tests { let dir = create_temp_dir("test_file_watcher_recreate"); let file_path = dir.path().join("test_file.txt"); - - // Create initial file std::fs::write(&file_path, "initial content").unwrap(); - let counter = Arc::new(AtomicUsize::new(0)); - let counter_clone = counter.clone(); + let content = spawn_reader(&file_path); - FileWatcherBuilder::new() - .watch_path(&file_path) - .unwrap() - .config(FileWatcherConfig::new()) - .spawn(move || { - counter_clone.fetch_add(1, Ordering::SeqCst); - }) - .unwrap(); - - // Give watcher time to start - std::thread::sleep(Duration::from_millis(100)); - - // Delete the file std::fs::remove_file(&file_path).unwrap(); - std::thread::sleep(Duration::from_millis(100)); - // Recreate the file - this should still be detected because we watch the directory std::fs::write(&file_path, "recreated content").unwrap(); - // Wait for the event to be processed - std::thread::sleep(Duration::from_millis(500)); - assert!( - counter.load(Ordering::SeqCst) >= 1, + wait_until(|| *content.lock().unwrap() == "recreated content"), "Watcher should have detected file recreation" ); } + + #[test] + fn test_file_watcher_ignores_unrelated_files() { + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_unrelated"); + let file_path = dir.path().join("watched.txt"); + std::fs::write(&file_path, "initial content").unwrap(); + + let counter = spawn_counter(&file_path); + + std::fs::write(dir.path().join("unrelated.txt"), "noise").unwrap(); + // Write the watched file afterwards: once its event is seen, the + // earlier unrelated event has been processed too. + std::fs::write(&file_path, "modified content").unwrap(); + assert!(wait_until(|| counter.load(Ordering::SeqCst) >= 1)); + std::thread::sleep(Duration::from_millis(200)); + let seen = counter.load(Ordering::SeqCst); + + // Only unrelated activity from now on. + std::fs::write(dir.path().join("unrelated.txt"), "more noise").unwrap(); + std::fs::remove_file(dir.path().join("unrelated.txt")).unwrap(); + std::thread::sleep(Duration::from_millis(500)); + assert_eq!( + counter.load(Ordering::SeqCst), + seen, + "Unrelated files in the watched directory must not trigger the callback" + ); + } + + #[cfg(unix)] + #[test] + fn test_resolve_watch_targets_symlink_chain() { + use std::os::unix::fs::symlink; + + let dir = create_temp_dir("test_resolve_watch_targets"); + let base = dir.path().canonicalize().unwrap(); + let keys = |paths: &[PathBuf]| paths.iter().map(|p| entry_key(p)).collect::>(); + + // real/users <- secrets (relative dir symlink) <- conf/users (relative + // file symlink going through `..`). + std::fs::create_dir(base.join("real")).unwrap(); + std::fs::write(base.join("real/users"), "").unwrap(); + symlink("real", base.join("secrets")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink("../secrets/users", base.join("conf/users")).unwrap(); + + let targets = resolve_watch_targets(&[base.join("conf/users")]); + assert_eq!( + targets.dirs, + HashSet::from([base.clone(), base.join("conf"), base.join("real")]) + ); + assert_eq!(targets.file_dirs, HashSet::from([base.join("real")])); + assert_eq!( + targets.entries, + keys(&[ + base.join("conf/users"), + base.join("secrets"), + base.join("real/users"), + ]) + ); + + // A plain file only needs its parent directory. + let targets = resolve_watch_targets(&[base.join("real/users")]); + assert_eq!(targets.dirs, HashSet::from([base.join("real")])); + assert_eq!(targets.entries, keys(&[base.join("real/users")])); + + // A dangling link keeps the directory where the target should appear. + std::fs::remove_file(base.join("secrets")).unwrap(); + let targets = resolve_watch_targets(&[base.join("conf/users")]); + assert_eq!( + targets.dirs, + HashSet::from([base.clone(), base.join("conf")]) + ); + assert!(targets.entries.contains(&entry_key(&base.join("secrets")))); + + // Symlink loops terminate. + symlink("loop", base.join("loop")).unwrap(); + let targets = resolve_watch_targets(&[base.join("loop")]); + assert_eq!(targets.dirs, HashSet::from([base.clone()])); + } + + #[cfg(unix)] + #[test] + fn test_file_watcher_detects_dir_symlink_swap() { + use std::os::unix::fs::symlink; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_symlink_swap"); + let base = dir.path(); + std::fs::create_dir(base.join("gen1")).unwrap(); + std::fs::write(base.join("gen1/file"), "1").unwrap(); + std::fs::create_dir(base.join("gen2")).unwrap(); + std::fs::write(base.join("gen2/file"), "2").unwrap(); + symlink(base.join("gen1"), base.join("current")).unwrap(); + std::fs::create_dir(base.join("conf")).unwrap(); + symlink(base.join("current/file"), base.join("conf/file")).unwrap(); + + let content = spawn_reader(&base.join("conf/file")); + + // Swap `current` atomically, like `ln -sfn gen2 current`. + symlink(base.join("gen2"), base.join("current.tmp")).unwrap(); + std::fs::rename(base.join("current.tmp"), base.join("current")).unwrap(); + assert!( + wait_until(|| *content.lock().unwrap() == "2"), + "Watcher should have detected the swap" + ); + // Let any remaining events of the swap drain, so they can't read the + // edit below by accident. + std::thread::sleep(Duration::from_millis(200)); + + // The watch follows the chain to the new target directory. + std::fs::write(base.join("gen2/file"), "22").unwrap(); + assert!( + wait_until(|| *content.lock().unwrap() == "22"), + "Watcher should have detected the edit of the new target" + ); + } + + /// A directory holding a symlink of the chain may be traversable but not + /// readable. It can't be watched, but that must not prevent watching the + /// file itself. + #[cfg(unix)] + #[test] + fn test_file_watcher_skips_unreadable_link_dir() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + common_telemetry::init_default_ut_logging(); + + let dir = create_temp_dir("test_file_watcher_unreadable_dir"); + let base = dir.path(); + std::fs::create_dir(base.join("real")).unwrap(); + std::fs::write(base.join("real/file"), "1").unwrap(); + std::fs::create_dir(base.join("locked")).unwrap(); + symlink(base.join("real/file"), base.join("locked/file")).unwrap(); + + /// Restores a readable mode on drop, even if the test panics: the + /// temp dir can't be removed while `locked` is unreadable, since + /// removing it requires listing its entries. + struct RestoreMode(PathBuf); + impl Drop for RestoreMode { + fn drop(&mut self) { + let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o755)); + } + } + + // Declared after `dir`, so it drops (restoring the mode) first. + let _restore = RestoreMode(base.join("locked")); + // Traversable (so `locked/file` still resolves) but not readable. + std::fs::set_permissions(base.join("locked"), std::fs::Permissions::from_mode(0o311)) + .unwrap(); + if std::fs::read_dir(base.join("locked")).is_ok() { + // Running as root (or with CAP_DAC_OVERRIDE): nothing to test. + return; + } + + // Would panic if watching `locked` were fatal. + let content = spawn_reader(&base.join("locked/file")); + + // The directory of the final file is still watched. + std::fs::write(base.join("real/file"), "2").unwrap(); + assert!(wait_until(|| *content.lock().unwrap() == "2")); + } } From 678aa81cae1ad7089f97a6397b820f675d201020 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 28 Sep 2026 04:39:20 +0000 Subject: [PATCH 047/133] fix(client): complete transport lane isolation (#9030) * fix(client): complete transport lane isolation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(client): deprecate legacy single-manager constructors Per review: mark the legacy single-manager constructors and helper as deprecated so callers move to the isolated query/control manager pair. Tests intentionally exercising the legacy path are annotated with `#[allow(deprecated)]`. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(client): update deprecated constructor callers for CI Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/client/src/client.rs | 100 +++++++++++++- src/client/src/client_manager.rs | 4 +- src/client/src/region.rs | 1 + src/cmd/src/cli.rs | 7 +- src/cmd/src/frontend.rs | 2 - src/flow/src/batching_mode/frontend_client.rs | 56 +++++--- src/frontend/src/frontend.rs | 6 +- src/frontend/src/server.rs | 13 +- src/meta-srv/src/procedure/utils.rs | 3 + src/meta-srv/src/utils/database.rs | 69 +++++++++- src/query/src/datafusion.rs | 122 ++++-------------- tests-integration/src/cluster.rs | 6 +- tests-integration/src/grpc/flight.rs | 2 + tests-integration/src/grpc/network.rs | 2 + tests-integration/src/test_util.rs | 10 +- tests-integration/tests/grpc.rs | 3 + tests/runner/src/client.rs | 6 +- 17 files changed, 283 insertions(+), 129 deletions(-) diff --git a/src/client/src/client.rs b/src/client/src/client.rs index fc40c98a723..dfc0e44fad5 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -87,6 +87,7 @@ struct Inner { impl Default for Inner { fn default() -> Self { + #[allow(deprecated)] Self::with_manager_and_peers(ChannelManager::new(), Vec::new(), ClientOptions::default()) } } @@ -105,6 +106,7 @@ struct Peers { } impl Inner { + #[deprecated(note = "legacy single-manager path shares its pool between lanes")] fn with_manager_and_peers( channel_manager: ChannelManager, peers: Vec, @@ -113,6 +115,7 @@ impl Inner { Self::with_managers_and_peers(channel_manager.clone(), channel_manager, peers, options) } + // The explicit dual-manager path keeps query and control pools independent. fn with_managers_and_peers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, @@ -222,36 +225,61 @@ fn random_initial_delay(max_delay: Duration) -> Duration { } impl Client { + /// Creates a client whose query and control lanes intentionally share the default manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn new() -> Self { Default::default() } + /// Creates a client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_urls(urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_urls_and_options(urls, ClientOptions::default()) } /// Creates a client with URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_urls_and_options(urls: A, options: ClientOptions) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(ChannelManager::new(), urls, options) } + /// Creates a TLS client whose query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls(urls: A, client_tls: ClientTlsOption) -> Result where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_tls_and_urls_and_options(urls, client_tls, ClientOptions::default()) } /// Creates a client with TLS URLs and custom options. + /// + /// The query and control lanes intentionally share one manager. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_tls_and_urls_and_options( urls: A, client_tls: ClientTlsOption, @@ -265,6 +293,7 @@ impl Client { let tls_config = load_client_tls_config(Some(client_tls)).context(error::CreateTlsChannelSnafu)?; let channel_manager = ChannelManager::with_config(channel_config, tls_config); + #[allow(deprecated)] Ok(Self::with_manager_and_urls_and_options( channel_manager, urls, @@ -272,15 +301,24 @@ impl Client { )) } + /// Creates a client with one manager shared intentionally by query and control lanes. + #[deprecated( + note = "shares one manager between query and control lanes; use `with_query_and_control_managers` with independently constructed managers instead" + )] pub fn with_manager_and_urls(channel_manager: ChannelManager, urls: A) -> Self where U: AsRef, A: AsRef<[U]>, { + #[allow(deprecated)] Self::with_manager_and_urls_and_options(channel_manager, urls, ClientOptions::default()) } - pub(crate) fn with_managers_and_urls( + /// Creates a client with query and control lanes backed by the supplied managers. + /// + /// The lanes are isolated only when the supplied managers are independently constructed; + /// this constructor does not enforce that they are distinct. + pub fn with_query_and_control_managers( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -289,7 +327,7 @@ impl Client { U: AsRef, A: AsRef<[U]>, { - Self::with_managers_and_urls_and_options( + Self::with_query_and_control_managers_and_options( query_channel_manager, control_channel_manager, urls, @@ -297,7 +335,7 @@ impl Client { ) } - fn with_managers_and_urls_and_options( + fn with_query_and_control_managers_and_options( query_channel_manager: ChannelManager, control_channel_manager: ChannelManager, urls: A, @@ -323,6 +361,11 @@ impl Client { } /// Creates a client with a channel manager, URLs, and custom options. + /// + /// The query and control lanes intentionally share this manager and its pool. + #[deprecated( + note = "shares this manager and its pool between query and control lanes; use `with_query_and_control_managers` instead" + )] pub fn with_manager_and_urls_and_options( channel_manager: ChannelManager, urls: A, @@ -333,7 +376,8 @@ impl Client { A: AsRef<[U]>, { let channel_manager_for_query = channel_manager.clone(); - Self::with_managers_and_urls_and_options( + // Legacy constructors intentionally share the manager and therefore its pool. + Self::with_query_and_control_managers_and_options( channel_manager_for_query, channel_manager, urls, @@ -541,6 +585,23 @@ impl Client { Ok(()) } + /// Returns the number of cached channels in the query and control pools for tests. + #[cfg(feature = "testing")] + pub fn channel_pool_sizes(&self) -> (usize, usize) { + let pool_size = |channel_manager: &ChannelManager| { + let mut size = 0; + channel_manager.retain_channel(|_, _| { + size += 1; + true + }); + size + }; + ( + pool_size(&self.inner.query_channel_manager), + pool_size(&self.inner.control_channel_manager), + ) + } + /// Returns peer addresses grouped by active and inactive state for tests. #[cfg(feature = "testing")] pub fn peer_addresses_by_state(&self) -> (Vec, Vec) { @@ -559,6 +620,7 @@ impl Client { } #[cfg(test)] +#[allow(deprecated)] mod tests { use std::collections::HashSet; use std::sync::Arc; @@ -666,6 +728,36 @@ mod tests { ] } + #[tokio::test] + async fn test_explicit_dual_manager_constructor_uses_isolated_reused_channel_pools() { + let query_channel_manager = ChannelManager::new(); + let control_channel_manager = ChannelManager::new(); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + ["127.0.0.1:3001"], + ); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(0, channel_pool_size(&control_channel_manager)); + + client.make_control_flight_client(false, false).unwrap(); + client.make_control_flight_client(false, false).unwrap(); + assert_eq!(1, channel_pool_size(&query_channel_manager)); + assert_eq!(1, channel_pool_size(&control_channel_manager)); + } + + fn channel_pool_size(manager: &ChannelManager) -> usize { + let mut size = 0; + manager.retain_channel(|_, _| { + size += 1; + true + }); + size + } + #[test] fn test_inner() { let inner = Inner::default(); diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index 008ada00d23..4d723361e80 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -77,6 +77,7 @@ impl FlownodeManager for NodeClients { } impl NodeClients { + /// Creates node clients with independent query and control channel pools. pub fn new(config: ChannelConfig) -> Self { Self { query_channel_manager: ChannelManager::with_config(config.clone(), None), @@ -91,7 +92,7 @@ impl NodeClients { pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { - Client::with_managers_and_urls( + Client::with_query_and_control_managers( self.query_channel_manager.clone(), self.control_channel_manager.clone(), vec![datanode.addr.clone()], @@ -153,6 +154,7 @@ mod tests { assert_pool_has_one_address(&node_clients.control_channel_manager); let manager = ChannelManager::new(); + #[allow(deprecated)] let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]); legacy.make_flight_client(false, false).unwrap(); legacy.make_control_flight_client(false, false).unwrap(); diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 4a4d243f36a..f0c683ec155 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -418,6 +418,7 @@ pub fn check_response_header(header: &Option) -> Result<()> { } #[cfg(test)] +#[allow(deprecated)] mod test { use api::v1::Status as PbStatus; use api::v1::region::region_server::{Region, RegionServer}; diff --git a/src/cmd/src/cli.rs b/src/cmd/src/cli.rs index 09ec9646d8b..3739d4cc4fc 100644 --- a/src/cmd/src/cli.rs +++ b/src/cmd/src/cli.rs @@ -108,6 +108,7 @@ mod tests { use clap::Parser; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::ChannelManager; use common_telemetry::logging::LoggingOptions; use rand::Rng; @@ -174,7 +175,11 @@ mod tests { let mut instance = standalone.build(standalone_opts).await?; instance.start().await?; - let client = Client::with_urls([rpc_addr.as_str()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [rpc_addr.as_str()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); database .sql(r#"CREATE DATABASE "cli.export.create_table";"#) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 6cae195f3e6..b000f7e1f02 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -467,8 +467,6 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; - // Source Flight streams and sink unary responses share pooled connections. - channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index aa2f2aa072a..822d8dea45b 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -91,7 +91,8 @@ impl HandlerMutable { pub enum FrontendClient { Distributed { meta_client: Arc, - chnl_mgr: ChannelManager, + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, query: QueryOptions, batch_opts: BatchingModeOptions, }, @@ -136,17 +137,19 @@ impl FrontendClient { batch_opts: BatchingModeOptions, ) -> Result { common_telemetry::info!("Frontend client build without auth"); + let cfg = ChannelConfig::new() + .connect_timeout(batch_opts.grpc_conn_timeout) + .timeout(Some(batch_opts.query_timeout)); + let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) + .context(InvalidClientConfigSnafu)?; + // Keep separate TLS config handles and pools for query and control traffic. + let query_channel_manager = ChannelManager::with_config(cfg.clone(), tls_config.clone()); + let control_channel_manager = ChannelManager::with_config(cfg, tls_config); + Ok(Self::Distributed { meta_client, - chnl_mgr: { - let cfg = ChannelConfig::new() - .connect_timeout(batch_opts.grpc_conn_timeout) - .timeout(Some(batch_opts.query_timeout)); - - let tls_config = load_client_tls_config(batch_opts.frontend_tls.clone()) - .context(InvalidClientConfigSnafu)?; - ChannelManager::with_config(cfg, tls_config) - }, + query_channel_manager, + control_channel_manager, query, batch_opts, }) @@ -227,7 +230,8 @@ impl FrontendClient { frontends: &[Peer], ) -> Result, Error> { let Self::Distributed { - chnl_mgr, + query_channel_manager, + control_channel_manager, batch_opts, .. } = self @@ -240,10 +244,15 @@ impl FrontendClient { .iter() .map(|peer| { let addr = peer.addr.clone(); - let chnl_mgr = chnl_mgr.clone(); + let query_channel_manager = query_channel_manager.clone(); + let control_channel_manager = control_channel_manager.clone(); async move { - let client = Client::with_manager_and_urls(chnl_mgr, vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager, + control_channel_manager, + vec![addr.clone()], + ); let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); match tokio::time::timeout(probe_timeout, database.sql("SELECT 1")).await { @@ -282,7 +291,8 @@ impl FrontendClient { ) -> Result { let Self::Distributed { meta_client: _, - chnl_mgr, + query_channel_manager, + control_channel_manager, query: _, batch_opts, } = self @@ -302,7 +312,11 @@ impl FrontendClient { for peer in frontends { let addr = peer.addr.clone(); - let client = Client::with_manager_and_urls(chnl_mgr.clone(), vec![addr.clone()]); + let client = Client::with_query_and_control_managers( + query_channel_manager.clone(), + control_channel_manager.clone(), + vec![addr.clone()], + ); let database = Database::new(catalog, schema, client); let db = DatabaseWithPeer::new(database, peer); match db.try_select_one().await { @@ -1247,7 +1261,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, @@ -1284,7 +1302,11 @@ mod tests { let database = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls([addr.as_str()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.as_str()], + ), ); let db = DatabaseWithPeer::new( database, diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index b2178e1ff35..71e5888b8fa 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -643,7 +643,11 @@ max_batch_rows = 25 expected: std::result::Result<&str, (StatusCode, &str)>, ) { let addr = frontend.server_handlers().addr(GRPC_SERVER).unwrap(); - let client = Client::with_urls([addr.to_string()]); + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [addr.to_string()], + ); let client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); let response = client.sql("SELECT 1").await; diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index da89d871ac0..a5568b24be5 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -520,6 +520,7 @@ mod tests { use async_trait::async_trait; use auth::{UserProviderRef, static_user_provider_from_option}; use client::{Client, Database}; + use common_grpc::channel_manager::ChannelManager; use meta_client::client::MetaClientBuilder; use servers::grpc::GRPC_SERVER; use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream}; @@ -938,12 +939,20 @@ mod tests { let public_database = Database::new( "greptime", "public", - Client::with_urls([public_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [public_addr.to_string()], + ), ); let internal_database = Database::new( "greptime", "public", - Client::with_urls([internal_addr.to_string()]), + Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [internal_addr.to_string()], + ), ); let internal_result = internal_database.sql("SELECT 1").await; diff --git a/src/meta-srv/src/procedure/utils.rs b/src/meta-srv/src/procedure/utils.rs index c30519da984..49964de24b7 100644 --- a/src/meta-srv/src/procedure/utils.rs +++ b/src/meta-srv/src/procedure/utils.rs @@ -336,6 +336,9 @@ pub mod mock { }), ) .unwrap(); + // The mock transport is a single in-memory duplex connection, so both lanes + // must resolve through this one pre-wired manager. + #[allow(deprecated)] Client::with_manager_and_urls(channel_manager, vec![datanode.addr.clone()]) } } diff --git a/src/meta-srv/src/utils/database.rs b/src/meta-srv/src/utils/database.rs index 976be8a85b6..796d63009d8 100644 --- a/src/meta-srv/src/utils/database.rs +++ b/src/meta-srv/src/utils/database.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use client::error::{ExternalSnafu, Result as ClientResult}; use client::{Client, Database, Output}; use common_error::ext::BoxedError; +use common_grpc::channel_manager::ChannelManager; use common_meta::peer::PeerDiscoveryRef; use common_telemetry::{debug, warn}; use snafu::{ResultExt, ensure}; @@ -43,6 +44,9 @@ impl<'a> DatabaseContext<'a> { } /// A cached frontend database operator used by metasrv. +/// +/// Its cached clients use independent query and control channel-manager pools; +/// the legacy single-manager client constructors intentionally remain shared. pub struct DatabaseOperator { peer_discovery: PeerDiscoveryRef, client: RwLock>, @@ -110,7 +114,11 @@ impl DatabaseOperator { debug!("Available frontend addresses: {:?}", urls); - Ok(Client::with_urls(urls)) + Ok(Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + urls, + )) } async fn maybe_init_client(&self) -> ClientResult { @@ -147,3 +155,62 @@ fn should_reset_client(result: &client::error::Result) -> bool { .map(|err| err.is_connection_error()) .unwrap_or(false) } + +#[cfg(test)] +mod tests { + use api::v1::meta::heartbeat_request::NodeWorkloads; + use common_meta::cluster::{FrontendStatus, NodeInfo, NodeStatus}; + use common_meta::peer::{Peer, PeerDiscovery}; + + use super::*; + + struct TestPeerDiscovery; + + #[async_trait::async_trait] + impl PeerDiscovery for TestPeerDiscovery { + async fn active_frontends(&self) -> common_meta::error::Result> { + Ok(vec![NodeInfo { + peer: Peer::new(1, "127.0.0.1:3001".to_string()), + last_activity_ts: 0, + status: NodeStatus::Frontend(FrontendStatus::default()), + version: String::new(), + git_commit: String::new(), + start_time_ms: 0, + total_cpu_millicores: 0, + total_memory_bytes: 0, + cpu_usage_millicores: 0, + memory_usage_bytes: 0, + hostname: String::new(), + env_vars: Default::default(), + }]) + } + + async fn active_datanodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + + async fn active_flownodes( + &self, + _filter: Option fn(&'a NodeWorkloads) -> bool>, + ) -> common_meta::error::Result> { + unreachable!() + } + } + + #[tokio::test] + async fn test_build_client_uses_isolated_reused_channel_pools() { + let operator = DatabaseOperator::new(Arc::new(TestPeerDiscovery)); + let client = operator.build_client().await.unwrap(); + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_eq!((1, 0), client.channel_pool_sizes()); + + client.find_channel().unwrap(); + client.find_channel().unwrap(); + assert_eq!((1, 1), client.channel_pool_sizes()); + } +} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 7df99f059ca..970f3109804 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -31,7 +31,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -45,7 +45,6 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; -use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -82,29 +81,6 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; -// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared -// HTTP/2 connection are pending, trading bounded memory for request liveness. -async fn forward_record_batches( - mut stream: SendableRecordBatchStream, - batch_tx: tokio::sync::mpsc::UnboundedSender>, -) -> Result<()> { - while let Some(batch) = stream.next().await { - match batch.context(CreateRecordBatchSnafu) { - Ok(batch) => { - if batch_tx.send(Ok(batch)).is_err() { - break; - } - tokio::task::yield_now().await; - } - Err(error) => { - let _ = batch_tx.send(Err(error)); - break; - } - } - } - Ok(()) -} - fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -230,7 +206,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let stream = match data { + let mut stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -241,31 +217,32 @@ impl DatafusionQueryEngine { match dml.op { WriteOp::Insert(_) => { - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_insert_record_batches( - batch_rx, - &table_name, - table.schema(), - query_ctx.clone(), - ); - let ((), (rows, cost)) = try_join(producer, consumer).await?; - affected_rows += rows; - insert_cost += cost; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(&table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } } WriteOp::Delete => { - // Keep DELETE on the same producer/consumer schedule as INSERT so the source - // stream can continue draining while mutation RPCs are pending. - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_delete_record_batches( - batch_rx, - &table_name, - &table, - query_ctx.clone(), - ); - let (rows, ()) = try_join(consumer, producer).await?; - affected_rows += rows; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(&table_name, &table, column_vectors, query_ctx.clone()) + .await?; + } } _ => unreachable!("guarded by the 'ensure!' at the beginning"), } @@ -275,53 +252,6 @@ impl DatafusionQueryEngine { )) } - async fn consume_insert_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table_schema: Arc, - query_ctx: QueryContextRef, - ) -> Result<(usize, usize)> { - let mut affected_rows = 0; - let mut insert_cost = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table_schema.clone()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - // We ignore the insert op. - let output = self - .insert(table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - Ok((affected_rows, insert_cost)) - } - - async fn consume_delete_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table: &TableRef, - query_ctx: QueryContextRef, - ) -> Result { - let mut affected_rows = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - affected_rows += self - .delete(table_name, table, column_vectors, query_ctx.clone()) - .await?; - } - Ok(affected_rows) - } - #[tracing::instrument(skip_all)] async fn delete( &self, diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index bc2ed6f20a4..4f762ba709b 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -682,11 +682,15 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ( addr.to_string(), - Client::with_manager_and_urls(channel_manager, [addr]), + // The mock connector pool lives on this single manager, so both lanes + // intentionally share it (same semantics as the legacy constructor). + Client::with_query_and_control_managers(channel_manager.clone(), channel_manager, [addr]), ) } +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { use api::v1::flow::FlowRequest; use api::v1::region::{ diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index fb27a273516..2cd0c44d504 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod test { use std::collections::HashMap; use std::net::SocketAddr; diff --git a/tests-integration/src/grpc/network.rs b/tests-integration/src/grpc/network.rs index 427317fb05d..97e4c56494a 100644 --- a/tests-integration/src/grpc/network.rs +++ b/tests-integration/src/grpc/network.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Mock connectors are registered on a shared channel manager. #[cfg(test)] +#[allow(deprecated)] mod tests { use std::io; use std::pin::Pin; diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 622d3ba6c16..62c32091d76 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -977,8 +977,14 @@ pub async fn setup_authenticated_grpc_database( .unwrap(); let grpc_server = setup_grpc_server_for_frontend_instance(instance, Some(user_provider)).await; let grpc_addr = grpc_server.bind_addr().unwrap().to_string(); - let mut database = - Database::new_with_dbname("greptime-public", Client::with_urls(vec![grpc_addr])); + let mut database = Database::new_with_dbname( + "greptime-public", + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_addr], + ), + ); database.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic { username: username.to_string(), password: password.to_string(), diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 970f65a2383..c73f80da78e 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Retain coverage of the legacy constructors, including TLS configuration. +#![allow(deprecated)] + use std::sync::Arc; use api::v1::alter_table_expr::Kind; diff --git a/tests/runner/src/client.rs b/tests/runner/src/client.rs index 163e827e53a..a595c877940 100644 --- a/tests/runner/src/client.rs +++ b/tests/runner/src/client.rs @@ -64,7 +64,11 @@ impl MultiProtocolClient { let grpc_client = Database::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, - Client::with_urls(vec![grpc_server_addr]), + Client::with_query_and_control_managers( + Default::default(), + Default::default(), + vec![grpc_server_addr], + ), ); let pg_client = create_postgres_client(pg_server_addr).await; let mysql_client = create_mysql_client(mysql_server_addr).await; From 825402dc70af64f754727f8558f8c8aca8c68b5e Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 28 Sep 2026 05:07:16 +0000 Subject: [PATCH 048/133] ci: add manually triggered tracesbench workflow (#9372) Signed-off-by: WenyXu --- .github/workflows/tracesbench.yml | 323 ++++++++++++++++++++++++++++++ 1 file changed, 323 insertions(+) create mode 100644 .github/workflows/tracesbench.yml diff --git a/.github/workflows/tracesbench.yml b/.github/workflows/tracesbench.yml new file mode 100644 index 00000000000..60081d69ad9 --- /dev/null +++ b/.github/workflows/tracesbench.yml @@ -0,0 +1,323 @@ +name: Tracesbench +on: + workflow_dispatch: + inputs: + dataset: + description: Dataset profile (main is experimental; native queries may exceed memory limits) + type: choice + options: [sanity, preflight, main] + default: sanity + required: true + runtime_image: + description: Aliyun runtime image (tag or digest reference) + type: string + default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260928-4e52d0b0@sha256:39d2d87593788e0c47a54768f6323679b8e03d7dc72a8158a9b22cf29f883baa + required: true + provider: + description: Cloud provider for the benchmark runner + type: choice + options: [Aliyun, AWS] + default: Aliyun + required: true + instance_type: + description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' + type: string + default: auto + required: true + benchmark_root: + description: Absolute disk-backed work directory (outside the o11ybench checkout) + type: string + default: /home/runner/benchmark-data + required: true + system_disk_gib: + description: 'System disk GiB: auto (or 0) = 40 for sanity/preflight, 200 for main; enter 20..2048 to override' + type: string + default: auto + required: true + benchmark_timeout_minutes: + description: Benchmark job timeout in minutes (1..360) + type: number + default: 360 + required: true + janitor_ttl_hours: + description: Instance lifetime before janitor eligibility (1..168 hours; at least timeout + 2h) + type: number + default: 8 + required: true + db_cpus: + description: CPU limit per DB container + type: string + default: '8' + required: true + db_memory: + description: Memory limit per DB container + type: string + default: 16g + required: true + runtime_cpus: + description: CPU limit per runtime container + type: string + default: '2' + required: true + runtime_memory: + description: Memory limit per runtime container (separate from DB memory) + type: string + default: 4g + required: true +permissions: + contents: read +env: + DATASET: ${{ inputs.dataset }} + RUNTIME_CPUS: ${{ inputs.runtime_cpus }} + RUNTIME_MEMORY: ${{ inputs.runtime_memory }} + BENCHMARK_ROOT: ${{ inputs.benchmark_root }} + RUNTIME_IMAGE: ${{ inputs.runtime_image }} + PROVIDER: ${{ inputs.provider }} + REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} + SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} + BENCHMARK_TIMEOUT_MINUTES: ${{ inputs.benchmark_timeout_minutes }} + JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + DB_CPUS: ${{ inputs.db_cpus }} + DB_MEMORY: ${{ inputs.db_memory }} +jobs: + validate: + runs-on: ubuntu-latest + timeout-minutes: 15 + outputs: + instance_type: ${{ steps.runner.outputs.instance_type }} + system_disk_gib: ${{ steps.runtime.outputs.system_disk_gib }} + runtime_image: ${{ steps.runtime.outputs.runtime_image }} + o11ybench_ref: ${{ steps.runtime.outputs.o11ybench_ref }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Resolve CI runner + id: runner + run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + - name: Validate budgets and shared runtime + id: runtime + run: | + set -euo pipefail + [[ "$BENCHMARK_ROOT" == /* && "$BENCHMARK_ROOT" != / && "$BENCHMARK_ROOT" != */ ]] + [[ "$BENCHMARK_ROOT" != /tmp && "$BENCHMARK_ROOT" != /tmp/* && "$BENCHMARK_ROOT" != *,* ]] + [[ "$BENCHMARK_ROOT" != *$'\n'* && "$BENCHMARK_ROOT" != *$'\r'* ]] + [[ "/${BENCHMARK_ROOT#/}/" != */../* && "/${BENCHMARK_ROOT#/}/" != */./* ]] + case "$DATASET" in sanity|preflight|main) ;; *) echo "Unknown dataset"; exit 1;; esac + if [[ "$SYSTEM_DISK_GIB" == auto || "$SYSTEM_DISK_GIB" == 0 ]]; then + if [[ "$DATASET" == main ]]; then SYSTEM_DISK_GIB=200; else SYSTEM_DISK_GIB=40; fi + fi + for value in "$SYSTEM_DISK_GIB" "$BENCHMARK_TIMEOUT_MINUTES" "$JANITOR_TTL_HOURS"; do + [[ "$value" =~ ^[1-9][0-9]{0,3}$ ]] || { echo 'Resource budgets must be positive integers'; exit 1; } + done + ((SYSTEM_DISK_GIB >= 20 && SYSTEM_DISK_GIB <= 2048)) + ((BENCHMARK_TIMEOUT_MINUTES <= 360 && JANITOR_TTL_HOURS <= 168)) + ((JANITOR_TTL_HOURS * 60 >= BENCHMARK_TIMEOUT_MINUTES + 120)) + [[ "$DB_CPUS" =~ ^[1-9][0-9]*$ && "$DB_MEMORY" =~ ^[1-9][0-9]*[mMgG]$ ]] + [[ "$RUNTIME_CPUS" =~ ^[1-9][0-9]*$ && "$RUNTIME_MEMORY" =~ ^[1-9][0-9]*[mMgG]$ ]] + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] + docker pull "$RUNTIME_IMAGE" + metadata=$(docker image inspect "$RUNTIME_IMAGE") + resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") + revision=$(jq -er '.[0].Config.Labels["org.opencontainers.image.revision"]' <<< "$metadata") + [[ "$revision" =~ ^[0-9a-f]{40}$ ]] + # Refuse runtimes without the trace runner and its validated dataset/report contract. + help=$(docker run --rm --cpus "$RUNTIME_CPUS" --memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY" "$resolved" scripts/run_tracesbench_docker.py --help) + grep -F -- '--dataset {sanity,preflight,main}' <<< "$help" + grep -F -- '--runtime-cpus' <<< "$help" + grep -F -- '--runtime-memory' <<< "$help" + docker run --rm --cpus "$RUNTIME_CPUS" --memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY" "$resolved" scripts/check_run_tracesbench_docker.py + printf 'runtime_image=%s\no11ybench_ref=%s\nsystem_disk_gib=%s\n' "$resolved" "$revision" "$SYSTEM_DISK_GIB" >> "$GITHUB_OUTPUT" + provision: + needs: validate + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.ecs.outputs.label || steps.aws.outputs.label }} + instance_id: ${{ steps.ecs.outputs.instance_id || steps.aws.outputs.instance_id }} + runner_name: ${{ steps.ecs.outputs.runner_name || steps.aws.outputs.runner_name }} + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Create ephemeral ECS runner + if: inputs.provider == 'Aliyun' + id: ecs + uses: ./greptimedb/.github/actions/aliyun-ecs-create + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + vswitch-id: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + security-group-id: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + image-id: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + instance-type: ${{ needs.validate.outputs.instance_type }} + enable-docker: 'true' + system-disk-gib: ${{ needs.validate.outputs.system_disk_gib }} + ttl-hours: ${{ inputs.janitor_ttl_hours }} + resource-group-id: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + runner-uid: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + runner-gid: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Create ephemeral EC2 runner + if: inputs.provider == 'AWS' + id: aws + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_IMAGE_ID: ${{ vars.BENCHMARK_EC2_IMAGE_ID || 'auto' }} + AWS_EC2_INSTANCE_TYPE: ${{ needs.validate.outputs.instance_type }} + AWS_EC2_SUBNET_ID: ${{ vars.EC2_RUNNER_SUBNET_ID }} + AWS_EC2_SECURITY_GROUP_ID: ${{ vars.EC2_RUNNER_SECURITY_GROUP_ID }} + AWS_EC2_SYSTEM_DISK_GIB: ${{ needs.validate.outputs.system_disk_gib }} + AWS_EC2_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-provision.py + benchmark: + needs: + - validate + - provision + runs-on: ${{ needs.provision.outputs.label }} + timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} + env: + RESOLVED_RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} + RESOLVED_O11YBENCH_REF: ${{ needs.validate.outputs.o11ybench_ref }} + RUN_ROOT: ${{ inputs.benchmark_root }}/tracesbench-${{ github.run_id }}-${{ github.run_attempt }} + OWNER: o11ybench-traces-${{ github.run_id }}-${{ github.run_attempt }} + METADATA_ROOT: ${{ github.workspace }}/tracesbench-metadata + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Set up benchmark dependencies + uses: ./greptimedb/.github/actions/setup-benchmark + - name: Checkout matching benchmark tools + uses: actions/checkout@v4 + with: + repository: GreptimeTeam/o11ybench + token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ref: ${{ env.RESOLVED_O11YBENCH_REF }} + path: o11ybench + persist-credentials: false + - name: Check runner and runtime identity + id: runtime + run: | + set -euo pipefail + test "$(git -C o11ybench rev-parse HEAD)" = "$RESOLVED_O11YBENCH_REF" + mkdir -p "$BENCHMARK_ROOT" "$METADATA_ROOT" + test "$(stat -f -c %T "$BENCHMARK_ROOT")" != tmpfs || { echo 'Benchmark root must be disk-backed'; exit 1; } + df -h "$BENCHMARK_ROOT" /var/lib/docker + minimum_gib=10 + if [[ "$DATASET" == main ]]; then minimum_gib=100; fi + for path in "$BENCHMARK_ROOT" /var/lib/docker; do + available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') + ((available >= minimum_gib * 1024 * 1024)) || { echo "Need at least $minimum_gib GiB free on $path"; exit 1; } + done + docker pull "$RESOLVED_RUNTIME_IMAGE" + docker image inspect "$RESOLVED_RUNTIME_IMAGE" > "$METADATA_ROOT/runtime-image.json" + test "$(jq -r '.[0].Config.Labels["org.opencontainers.image.revision"]' "$METADATA_ROOT/runtime-image.json")" = "$RESOLVED_O11YBENCH_REF" + # Preserve the pinned catalog images and require all three targets in the report. + - name: Generate selected dataset once + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage generate \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: GreptimeDB load, correctness and benchmark + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: VictoriaTraces load, correctness and benchmark + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target victoriatraces --port 19428 \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Tempo load, correctness and benchmark + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target tempo --port 13200 \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Validate and render report + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage report \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Clean up owned containers + if: always() && steps.runtime.outcome == 'success' + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage cleanup \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Stage benchmark artifacts + if: always() + run: | + mkdir -p "$RUNNER_TEMP/tracesbench-artifacts" + if [[ -d "$METADATA_ROOT" ]]; then + rsync -a --no-owner --no-group "$METADATA_ROOT" "$RUNNER_TEMP/tracesbench-artifacts/" + fi + for entry in runs report report.json report.md; do + if [[ -e "$RUN_ROOT/$entry" ]]; then + rsync -a --no-owner --no-group "$RUN_ROOT/$entry" "$RUNNER_TEMP/tracesbench-artifacts/" + fi + done + - name: Upload HTML and results (no corpus) + if: always() + uses: actions/upload-artifact@v4 + with: + name: tracesbench-${{ inputs.dataset }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ runner.temp }}/tracesbench-artifacts + if-no-files-found: warn + retention-days: 14 + teardown: + needs: + - provision + - benchmark + if: ${{ always() && (needs.provision.outputs.instance_id != '' || (inputs.provider == 'AWS' && needs.provision.result != 'skipped')) }} + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + with: + path: greptimedb + persist-credentials: false + - name: Delete ECS and unregister runner + if: inputs.provider == 'Aliyun' + uses: ./greptimedb/.github/actions/aliyun-ecs-delete + with: + access-key-id: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + access-key-secret: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + region-id: ${{ vars.ALIYUN_ECS_REGION_ID }} + instance-id: ${{ needs.provision.outputs.instance_id }} + runner-name: ${{ needs.provision.outputs.runner_name }} + - uses: aws-actions/configure-aws-credentials@v4 + if: inputs.provider == 'AWS' + with: + aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} + aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + aws-region: ${{ vars.EC2_RUNNER_REGION }} + - uses: astral-sh/setup-uv@v6 + if: inputs.provider == 'AWS' + - name: Delete EC2 and unregister runner + if: inputs.provider == 'AWS' + env: + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + AWS_REGION: ${{ vars.EC2_RUNNER_REGION }} + AWS_EC2_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + AWS_EC2_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + AWS_EC2_RUN_ID: ${{ github.run_id }}-${{ github.run_attempt }} + run: uv run greptimedb/.github/scripts/aws-ec2-runner-teardown.py From f956f4da30a9702097e5825a91eaba60a2178645 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 28 Sep 2026 05:30:33 +0000 Subject: [PATCH 049/133] refactor(promql): split planner into focused modules (#9354) * refactor(promql): move planner tests into separate module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move function-specific planner methods into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move OR planner method into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): move binary island planner into module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): keep binary result labels in planner root Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): limit planner helper visibility to parent module Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): group set operator planning and localize imports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): use crate-rooted planner imports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/planner.rs | 9415 +---------------- .../src/promql/planner/function_plans.rs | 473 + src/query/src/promql/planner/island.rs | 536 + src/query/src/promql/planner/set_operator.rs | 832 ++ src/query/src/promql/planner/test.rs | 7585 +++++++++++++ 5 files changed, 9451 insertions(+), 9390 deletions(-) create mode 100644 src/query/src/promql/planner/function_plans.rs create mode 100644 src/query/src/promql/planner/island.rs create mode 100644 src/query/src/promql/planner/set_operator.rs create mode 100644 src/query/src/promql/planner/test.rs diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 73332c237c8..08e7005f9fb 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -12,9 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod function_plans; + +mod island; + mod matching_filters; -use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; +mod set_operator; + +use std::collections::{BTreeSet, HashSet, VecDeque}; use std::sync::Arc; use std::time::UNIX_EPOCH; @@ -46,14 +52,12 @@ use datafusion::logical_expr::{ BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; -use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; use datafusion_common::{DFSchema, NullEquality, TableReference}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; -use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit}; use datafusion_functions::core::coalesce; @@ -62,8 +66,7 @@ use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond, - RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, + EmptyMetric, InstantManipulate, Millisecond, RangeManipulate, SeriesDivide, SeriesNormalize, build_special_time_expr, }; use promql::functions::{ @@ -72,24 +75,22 @@ use promql::functions::{ NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, - NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, - NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, - NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, - NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, - Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, UniqueMatchGroup, - quantile_udaf, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, + NativeHistogramPresentOverTime, NativeHistogramRate, NativeHistogramResets, + NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, + NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, PredictLinear, + PresentOverTime, PromqlFloatToString, QuantileOverTime, Rate, Resets, Round, StddevOverTime, + StdvarOverTime, SumOverTime, UniqueMatchGroup, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ AggregateExpr, BinModifier, BinaryExpr as PromBinaryExpr, Call, EvalStmt, Expr as PromExpr, - Function, FunctionArgs as PromFunctionArgs, LabelModifier, MatrixSelector, NumberLiteral, - Offset, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality, - VectorSelector, token, + Function, LabelModifier, MatrixSelector, NumberLiteral, Offset, ParenExpr, StringLiteral, + SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, token, }; use regex::{self, Regex}; use snafu::{OptionExt, ResultExt, ensure}; @@ -104,14 +105,12 @@ use crate::parser::{ EXPLAIN_VERBOSE_NODE_NAME, }; use crate::promql::error::{ - CatalogSnafu, ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, - ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, - InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, - MultipleMetricMatchersSnafu, MultipleVectorSnafu, NoMetricMatcherSnafu, PromqlPlanNodeSnafu, - Result, SameLabelSetSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, - UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu, - UnsupportedMatcherOpSnafu, UnsupportedVectorMatchSnafu, ValueNotFoundSnafu, - ZeroRangeSelectorSnafu, + CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, + FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, + InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, Result, SameLabelSetSnafu, TableNameNotFoundSnafu, + TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, + UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, }; use crate::query_engine::QueryEngineState; @@ -147,8 +146,6 @@ const FIELD_COLUMN_MATCHER: &str = "__field__"; const SCHEMA_COLUMN_MATCHER: &str = "__schema__"; const DB_COLUMN_MATCHER: &str = "__database__"; -/// Prefix for generated binary island leaf aliases. -const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; const OR_FLOAT_FIELD_PREFIX: &str = "__promql_or_float_"; const OR_HISTOGRAM_FIELD_PREFIX: &str = "__promql_or_histogram_"; const TIMESTAMP_VALUE_PREFIX: &str = "__promql_timestamp_value_"; @@ -198,103 +195,6 @@ struct PromPlannerContext { range: Option, } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct VectorLeafKey { - metric_name: String, - matchers: Vec<(String, String, String)>, - or_matchers: Vec>, - offset_ms: i128, - at: String, -} - -#[derive(Debug, Clone)] -struct IslandLeaf { - selector: VectorSelector, - display_table: String, -} - -#[derive(Debug, Clone)] -enum IslandExpr { - VectorLeaf(usize), - Scalar(DfExpr), - Unary { - input: Box, - }, - Binary { - op: TokenType, - lhs: Box, - rhs: Box, - }, -} - -impl IslandExpr { - fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { - if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { - return Some(Self::Scalar(expr)); - } - - match expr { - PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), - PromExpr::VectorSelector(selector) => { - let leaf = env.intern_leaf(selector)?; - Some(Self::VectorLeaf(leaf)) - } - PromExpr::Unary(UnaryExpr { expr }) => { - let input = Self::try_new(expr, env)?; - Some(Self::Unary { - input: Box::new(input), - }) - } - PromExpr::Binary(PromBinaryExpr { - lhs, - rhs, - op, - modifier, - }) if matches!( - op.id(), - token::T_ADD - | token::T_SUB - | token::T_MUL - | token::T_DIV - | token::T_MOD - | token::T_POW - | token::T_ATAN2 - ) && modifier.as_ref().is_none_or(|modifier| { - !modifier.return_bool - && modifier.matching.is_none() - && matches!(modifier.card, VectorMatchCardinality::OneToOne) - && modifier.fill_values.lhs.is_none() - && modifier.fill_values.rhs.is_none() - }) => - { - let lhs = Self::try_new(lhs, env)?; - let rhs = Self::try_new(rhs, env)?; - Some(Self::Binary { - op: *op, - lhs: Box::new(lhs), - rhs: Box::new(rhs), - }) - } - _ => None, - } - } -} - -#[derive(Debug, Default)] -struct IslandCollectEnv { - leaf_by_key: HashMap, - leaves: Vec, - vector_occurrences: usize, -} - -#[derive(Debug)] -struct PlannedIslandLeaf { - plan: LogicalPlan, - ctx: PromPlannerContext, - alias: TableReference, - display_table: String, -} - /// Result labels a vector-vector binary operation derives from its matching modifier, projected /// from the operand each label belongs to. #[derive(Debug)] @@ -324,81 +224,6 @@ impl BinaryResultLabels { } } -#[derive(Debug)] -struct IslandFieldExprs { - exprs: Vec, - names: Vec, - scalar: bool, -} - -impl VectorLeafKey { - fn from_selector(selector: &VectorSelector) -> Option { - let mut metric_name = selector.name.clone(); - let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); - let matcher_key = |matcher: &Matcher| { - ( - matcher.name.clone(), - matcher.op.to_string(), - matcher.value.clone(), - ) - }; - - for matcher in &selector.matchers.matchers { - if matcher.name == METRIC_NAME { - if matcher.op != MatchOp::Equal || metric_name.is_some() { - return None; - } - metric_name = Some(matcher.value.clone()); - } else { - matchers.push(matcher_key(matcher)); - } - } - matchers.sort(); - - let mut or_matchers = selector - .matchers - .or_matchers - .iter() - .map(|group| { - let mut group = group.iter().map(matcher_key).collect::>(); - group.sort(); - group - }) - .collect::>(); - or_matchers.sort(); - - Some(Self { - metric_name: metric_name?, - matchers, - or_matchers, - offset_ms: match &selector.offset { - Some(Offset::Pos(duration)) => duration.as_millis() as i128, - Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), - None => 0, - }, - at: format!("{:?}", selector.at), - }) - } -} - -impl IslandCollectEnv { - fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { - self.vector_occurrences += 1; - let key = VectorLeafKey::from_selector(selector)?; - if let Some(id) = self.leaf_by_key.get(&key) { - return Some(*id); - } - - let id = self.leaves.len(); - self.leaves.push(IslandLeaf { - selector: selector.clone(), - display_table: key.metric_name.clone(), - }); - self.leaf_by_key.insert(key, id); - Some(id) - } -} - impl PromPlannerContext { fn from_eval_stmt(stmt: &EvalStmt) -> Self { Self { @@ -1039,330 +864,6 @@ impl PromPlanner { }) } - async fn try_plan_binary_island( - &mut self, - binary_expr: &PromBinaryExpr, - ) -> Result> { - let original_ctx = self.ctx.clone(); - let mut collect_env = IslandCollectEnv::default(); - let Some(island_expr) = - IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) - else { - return Ok(None); - }; - - if collect_env.leaves.is_empty() - || collect_env.vector_occurrences <= collect_env.leaves.len() - { - return Ok(None); - } - - let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); - for (idx, leaf) in collect_env.leaves.iter().enumerate() { - let plan = self - .prom_vector_selector_to_plan(&leaf.selector, false) - .await?; - let ctx = self.ctx.clone(); - let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); - let plan = LogicalPlanBuilder::from(plan) - .alias(alias.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - planned_leaves.push(PlannedIslandLeaf { - plan, - ctx, - alias, - display_table: leaf.display_table.clone(), - }); - } - - if planned_leaves.iter().any(|leaf| { - Self::field_columns_contain_native_histogram( - leaf.plan.schema(), - &leaf.ctx.field_columns, - ) - }) { - self.ctx = original_ctx; - return Ok(None); - } - - if !Self::binary_island_join_contexts_supported(&planned_leaves) { - self.ctx = original_ctx; - return Ok(None); - } - - let mut input = planned_leaves[0].plan.clone(); - for right_idx in 1..planned_leaves.len() { - input = self.join_binary_island_leaf( - input, - &planned_leaves[0], - &planned_leaves[right_idx], - )?; - } - - let field_exprs = - Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; - if field_exprs.scalar || field_exprs.exprs.is_empty() { - self.ctx = original_ctx; - return Ok(None); - } - - let plan = self.project_binary_island( - input, - &planned_leaves[0].alias, - &planned_leaves[0].ctx, - field_exprs, - )?; - Ok(Some(plan)) - } - - fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { - if leaves - .iter() - .any(|leaf| leaf.ctx.time_index_column.is_none()) - { - return false; - } - - if leaves.len() <= 1 { - return true; - } - - let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); - - leaves.iter().skip(1).all(|leaf| { - (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) - || leaf.ctx.tag_columns.iter().collect::>() == first_tags - }) - } - - fn join_binary_island_leaf( - &self, - left: LogicalPlan, - first_leaf: &PlannedIslandLeaf, - right_leaf: &PlannedIslandLeaf, - ) -> Result { - let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() - || right_leaf.ctx.tag_columns.is_empty()) - && !first_leaf - .ctx - .tag_columns - .iter() - .chain(&right_leaf.ctx.tag_columns) - .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); - let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( - left.schema(), - right_leaf.plan.schema(), - &first_leaf.ctx, - &right_leaf.ctx, - only_join_time_index, - &None, - )?; - - if let (Some(left_time_index_column), Some(right_time_index_column)) = ( - first_leaf.ctx.time_index_column.clone(), - right_leaf.ctx.time_index_column.clone(), - ) { - left_keys.insert(left_time_index_column); - right_keys.insert(right_time_index_column); - } - - LogicalPlanBuilder::from(left) - .join_detailed( - right_leaf.plan.clone(), - JoinType::Inner, - ( - left_keys - .into_iter() - .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) - .collect::>(), - right_keys - .into_iter() - .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) - .collect::>(), - ), - force_empty_join.then_some(lit(false)), - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn build_binary_island_field_exprs( - expr: &IslandExpr, - leaves: &[PlannedIslandLeaf], - schema: &DFSchemaRef, - ) -> Result { - match expr { - IslandExpr::VectorLeaf(id) => { - let leaf = &leaves[*id]; - let exprs = leaf - .ctx - .field_columns - .iter() - .map(|field| { - schema - .qualified_field_with_name(Some(&leaf.alias), field) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }) - .collect::>>()?; - let names = leaf - .ctx - .field_columns - .iter() - .map(|field| format!("{}.{}", leaf.display_table, field)) - .collect(); - Ok(IslandFieldExprs { - exprs, - names, - scalar: false, - }) - } - IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { - exprs: vec![expr.clone()], - names: vec![expr.schema_name().to_string()], - scalar: true, - }), - IslandExpr::Unary { input } => { - let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; - let mut exprs = Vec::with_capacity(input.exprs.len()); - let mut names = Vec::with_capacity(input.names.len()); - for (expr, name) in input.exprs.into_iter().zip(input.names) { - exprs.push(DfExpr::Negative(Box::new(expr))); - names.push(format!("-{name}")); - } - Ok(IslandFieldExprs { - exprs, - names, - scalar: input.scalar, - }) - } - IslandExpr::Binary { op, lhs, rhs } => { - let same_leaf = match (&**lhs, &**rhs) { - (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) - if left == right => - { - Some(*left) - } - _ => None, - }; - let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; - let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; - let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let scalar = lhs.scalar && rhs.scalar; - let op = op.to_string(); - - let (exprs, names) = match (lhs.scalar, rhs.scalar) { - (true, true) => { - let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; - let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); - (vec![expr], vec![name]) - } - (true, false) => { - let mut exprs = Vec::with_capacity(rhs.exprs.len()); - let mut names = Vec::with_capacity(rhs.names.len()); - for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { - exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); - names.push(format!("{} {op} {rhs_name}", lhs.names[0])); - } - (exprs, names) - } - (false, true) => { - let mut exprs = Vec::with_capacity(lhs.exprs.len()); - let mut names = Vec::with_capacity(lhs.names.len()); - for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { - exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); - names.push(format!("{lhs_name} {op} {}", rhs.names[0])); - } - (exprs, names) - } - (false, false) => { - let mut exprs = Vec::new(); - let mut names = Vec::new(); - for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs - .exprs - .into_iter() - .zip(rhs.exprs) - .zip(lhs.names.into_iter().zip(rhs.names)) - .enumerate() - { - if let Some(leaf) = same_leaf { - let field = leaves[leaf] - .ctx - .field_columns - .get(idx) - .cloned() - .unwrap_or_else(|| lhs_name.clone()); - lhs_name = format!("lhs.{field}"); - rhs_name = format!("rhs.{field}"); - } - exprs.push(expr_builder(lhs_expr, rhs_expr)?); - names.push(format!("{lhs_name} {op} {rhs_name}")); - } - (exprs, names) - } - }; - - Ok(IslandFieldExprs { - exprs, - names, - scalar, - }) - } - } - } - - fn project_binary_island( - &mut self, - input: LogicalPlan, - base_alias: &TableReference, - base_ctx: &PromPlannerContext, - field_exprs: IslandFieldExprs, - ) -> Result { - self.ctx = base_ctx.clone(); - - let schema = input.schema(); - let non_field_exprs = base_ctx - .tag_columns - .iter() - .chain(base_ctx.time_index_column.iter()) - .map(|column| { - schema - .qualified_field_with_name(Some(base_alias), column) - .context(DataFusionPlanningSnafu) - .map(|field| DfExpr::Column(field.into())) - }); - let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) - .into_iter() - .map(Ok); - - self.ctx.field_columns = field_exprs.names; - let field_exprs = field_exprs - .exprs - .into_iter() - .zip(self.ctx.field_columns.iter()) - .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); - - let project_exprs = non_field_exprs - .chain(tsid_expr) - .chain(field_exprs) - .collect::>>()?; - - let plan = LogicalPlanBuilder::from(input) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - self.ctx.table_name = None; - self.ctx.schema_name = None; - - Ok(plan) - } - async fn prom_binary_expr_to_plan( &mut self, query_engine_state: &QueryEngineState, @@ -5000,423 +4501,6 @@ impl PromPlanner { Ok(normalized_exprs) } - /// Create a classic, native, or mixed histogram helper plan. - async fn create_histogram_plan( - &mut self, - function_name: &str, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - let float_literal = |param: &PromExpr| -> Result { - let value = (|| { - let expr = Self::get_param_as_literal_expr( - Some(param), - None, - Some(ArrowDataType::Float64), - ) - .ok()?; - let simplifier = ExprSimplifier::new(SimplifyContext::default()); - let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; - let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { - return None; - }; - let ScalarValue::Float64(Some(value)) = - value.cast_to(&ArrowDataType::Float64).ok()? - else { - return None; - }; - Some(value) - })() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - })?; - Ok(value) - }; - let (function, input) = match (function_name, args.args.as_slice()) { - (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( - HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), - input.as_ref().clone(), - ), - (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( - HistogramFoldOperation::Fraction { - lower: float_literal(lower)?.into(), - upper: float_literal(upper)?.into(), - }, - input.as_ref().clone(), - ), - _ => { - return FunctionInvalidArgumentSnafu { - fn_name: function_name.to_string(), - } - .fail(); - } - }; - - let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; - // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a - // stable series identifier anymore. HistogramFold must not treat it as a label column. - let input_plan = self.strip_tsid_column(input_plan)?; - self.ctx.use_tsid = false; - - if let Some((float_field, histogram_field)) = - Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) - .map(|(float, histogram)| (float.to_string(), histogram.to_string())) - { - if self.ctx.has_le_tag() { - return self.create_mixed_histogram_plan( - function, - input_plan, - float_field, - histogram_field, - ); - } - self.ctx.field_columns = vec![histogram_field]; - } - if self.all_field_columns_are_native_histograms(input_plan.schema()) { - return self.create_native_histogram_plan(function, input_plan); - } - - if !self.ctx.has_le_tag() { - // Return empty result instead of error when 'le' column is not found - // This handles the case when histogram metrics don't exist - return Ok(LogicalPlan::EmptyRelation( - datafusion::logical_expr::EmptyRelation { - produce_one_row: false, - schema: input_plan.schema().clone(), - }, - )); - } - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - // FIXME(ruihang): support multi fields - let field_column = self - .ctx - .field_columns - .first() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: function.function_name().to_string(), - })? - .clone(); - // remove le column from tag columns - self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); - - let fold = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - field_column, - time_index_column, - function, - None, - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(fold), - })) - } - - fn create_native_histogram_expr( - &self, - function: HistogramFoldOperation, - field_column: &str, - ) -> DfExpr { - let field = DfExpr::Column(Column::from_name(field_column)); - let (func, args) = match function { - HistogramFoldOperation::Quantile(quantile) => ( - Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(quantile))], - ), - HistogramFoldOperation::Fraction { lower, upper } => ( - Arc::new(NativeHistogramFraction::scalar_udf_with_collector( - self.promql_annotations.clone(), - )), - vec![field, lit(f64::from(lower)), lit(f64::from(upper))], - ), - }; - DfExpr::ScalarFunction(ScalarFunction { func, args }) - } - - fn create_native_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - ) -> Result { - ensure!( - self.ctx.field_columns.len() == 1, - MultiFieldsNotSupportedSnafu { - operator: function.function_name() - }, - ); - - let field_column = self.ctx.field_columns[0].clone(); - let function_expr = self.create_native_histogram_expr(function, &field_column); - let display_name = function_expr.schema_name().to_string(); - self.ctx.field_columns = vec![display_name.clone()]; - - let project_exprs = std::iter::once(self.create_time_index_column_expr()?) - .chain(std::iter::once(function_expr.alias(display_name))) - .chain(self.create_tag_column_exprs()?) - .collect::>(); - - LogicalPlanBuilder::from(input_plan) - .project(project_exprs) - .context(DataFusionPlanningSnafu)? - .filter(self.create_empty_values_filter_expr(false)?) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - fn create_mixed_histogram_plan( - &mut self, - function: HistogramFoldOperation, - input_plan: LogicalPlan, - float_field: String, - histogram_field: String, - ) -> Result { - let time_index_column = - self.ctx - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let tag_columns = self.ctx.tag_columns.clone(); - let folded = HistogramFold::new_with_operation( - LE_COLUMN_NAME.to_string(), - float_field.clone(), - time_index_column.clone(), - function, - Some(histogram_field.clone()), - input_plan, - ) - .context(DataFusionPlanningSnafu)?; - let record_collision = DfExpr::ScalarFunction(ScalarFunction { - func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( - "vector contains a mix of classic and native histograms".to_string(), - self.promql_annotations.clone(), - )), - args: vec![col(&float_field), col(&histogram_field)], - }); - let keep = when( - col(&float_field) - .is_not_null() - .and(col(&histogram_field).is_not_null()), - record_collision, - ) - .otherwise(lit(true)) - .context(DataFusionPlanningSnafu)?; - - let native_expr = self.create_native_histogram_expr(function, &histogram_field); - let output_field = native_expr.schema_name().to_string(); - let value = DfExpr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![col(&float_field), native_expr], - }); - self.ctx.field_columns = vec![output_field.clone()]; - LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { - node: Arc::new(folded), - })) - .filter(keep) - .context(DataFusionPlanningSnafu)? - .project( - std::iter::once(col(&time_index_column)) - .chain(std::iter::once(value.alias(output_field))) - .chain(tag_columns.iter().map(col)), - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu) - } - - /// Create a [SPECIAL_VECTOR_FUNCTION] plan - async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), - } - .fail(); - } - let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; - - // reuse `SPECIAL_TIME_FUNCTION` as name of time index column - self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); - self.ctx.reset_table_name_and_schema(); - self.ctx.tag_columns = vec![]; - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns = vec![greptime_value().to_string()]; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new( - EmptyMetric::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - SPECIAL_TIME_FUNCTION.to_string(), - greptime_value().to_string(), - Some(lit), - ) - .context(DataFusionPlanningSnafu)?, - ), - })) - } - - /// Create a [SCALAR_FUNCTION] plan - async fn create_scalar_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - ensure!( - args.len() == 1, - FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION - } - ); - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - let input_schema = input.schema().clone(); - let alternative_samples = - Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); - let histogram_fields = self - .ctx - .field_columns - .iter() - .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) - .count(); - ensure!( - self.ctx.field_columns.len() == 1 || alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: SCALAR_FUNCTION - }, - ); - let scalar_field = self - .ctx - .field_columns - .iter() - .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) - .or_else(|| self.ctx.field_columns.first()) - .cloned() - .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: SCALAR_FUNCTION, - })?; - let input = if histogram_fields == self.ctx.field_columns.len() { - // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN - // for every evaluation timestamp without attempting a Struct-to-Float64 cast. - LogicalPlanBuilder::from(input) - .filter(lit(false)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else if histogram_fields > 0 { - // A mixed vector contributes only its float samples to scalar(). - LogicalPlanBuilder::from(input) - .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - input - }; - let scalar_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - ScalarCalculate::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - input, - self.ctx.time_index_column.as_ref().unwrap(), - &self.ctx.tag_columns, - &scalar_field, - self.ctx.table_name.as_deref(), - ) - .context(PromqlPlanNodeSnafu)?, - ), - }); - // scalar plan have no tag columns - self.ctx.tag_columns.clear(); - self.ctx.aggregation_field_labels.clear(); - self.ctx.field_columns.clear(); - self.ctx - .field_columns - .push(scalar_plan.schema().field(1).name().clone()); - Ok(scalar_plan) - } - - /// Create a [SPECIAL_ABSENT_FUNCTION] plan - async fn create_absent_plan( - &mut self, - args: &PromFunctionArgs, - query_engine_state: &QueryEngineState, - ) -> Result { - if args.args.len() != 1 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), - } - .fail(); - } - let input = self - .prom_expr_to_plan(&args.args[0], query_engine_state) - .await?; - - let time_index_expr = self.create_time_index_column_expr()?; - let first_field_expr = - self.create_field_column_exprs()? - .pop() - .with_context(|| ValueNotFoundSnafu { - table: self.ctx.table_name.clone().unwrap_or_default(), - })?; - let first_value_expr = first_value(first_field_expr, vec![]); - - let ordered_aggregated_input = LogicalPlanBuilder::from(input) - .aggregate( - vec![time_index_expr.clone()], - vec![first_value_expr.clone()], - ) - .context(DataFusionPlanningSnafu)? - .sort(vec![time_index_expr.sort(true, false)]) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - let fake_labels = self - .ctx - .selector_matcher - .iter() - .filter_map(|matcher| match matcher.op { - MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), - _ => None, - }) - .collect::>(); - - // Create the absent plan - let absent_plan = LogicalPlan::Extension(Extension { - node: Arc::new( - Absent::try_new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - self.ctx.time_index_column.as_ref().unwrap().clone(), - self.ctx.field_columns[0].clone(), - fake_labels, - ordered_aggregated_input, - ) - .context(DataFusionPlanningSnafu)?, - ), - }); - - // The absent series carries the equality matchers as labels, not the input's - // tags or value fields, so the input's field grouping labels no longer apply. - self.ctx.aggregation_field_labels.clear(); - Ok(absent_plan) - } - /// Try to build a DataFusion Literal Expression from PromQL Expr, return /// `None` if the input is not a literal expression. fn try_build_literal_expr(expr: &PromExpr) -> Option { @@ -6618,203 +5702,6 @@ impl PromPlanner { matches!(plan, LogicalPlan::EmptyRelation(relation) if !relation.produce_one_row) } - /// Build a set operator (AND/OR/UNLESS) - fn set_op_on_non_field_columns( - &mut self, - mut left: LogicalPlan, - mut right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - op: TokenType, - modifier: &Option, - ) -> Result { - let left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - - if matches!(op.id(), token::T_LOR) { - return self.or_operator( - left, - right, - left_tag_col_set, - right_tag_col_set, - left_context, - right_context, - modifier, - ); - } - - if let Some(modifier) = modifier { - ensure!( - matches!( - modifier.card, - VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany - ), - UnsupportedVectorMatchSnafu { - name: modifier.card.clone(), - }, - ); - } - - let output_context = left_context.clone(); - let visible_left_schema = left.schema().clone(); - let mut left_context = left_context; - let mut right_context = right_context; - let added_marker_to_left = if Self::only_temporality_match_label_mismatches( - &left_context, - &right_context, - modifier, - ) { - let aligned = Self::align_temporality_match_column( - left, - right, - &mut left_context, - &mut right_context, - )?; - left = aligned.0; - right = aligned.1; - aligned.2 - } else { - false - }; - - let mut left_tag_col_set = left_context - .tag_columns - .iter() - .cloned() - .collect::>(); - let mut right_tag_col_set = right_context - .tag_columns - .iter() - .cloned() - .collect::>(); - if let Some(matching) = modifier - .as_ref() - .and_then(|modifier| modifier.matching.as_ref()) - { - match matching { - LabelModifier::Include(on) => { - let mask = on.labels.iter().cloned().collect::>(); - left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); - right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); - } - LabelModifier::Exclude(ignoring) => { - for label in &ignoring.labels { - let _ = left_tag_col_set.remove(label); - let _ = right_tag_col_set.remove(label); - } - } - } - } - ensure!( - left_tag_col_set == right_tag_col_set, - CombineTableColumnMismatchSnafu { - left: left_tag_col_set.iter().cloned().collect::>(), - right: right_tag_col_set.iter().cloned().collect::>(), - } - ); - - let left_time_index = left_context.time_index_column.clone().unwrap(); - let right_time_index = right_context.time_index_column.clone().unwrap(); - - // alias right time index column if necessary - if left_context.time_index_column != right_context.time_index_column { - let right_project_exprs = right - .schema() - .fields() - .iter() - .map(|field| { - if field.name() == &right_time_index { - DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) - } else { - DfExpr::Column(Column::from_name(field.name())) - } - }) - .collect::>(); - - right = LogicalPlanBuilder::from(right) - .project(right_project_exprs) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - } - - let join_keys = left_tag_col_set - .into_iter() - .chain([left_time_index]) - .collect::>(); - - ensure!( - left_context.field_columns.len() == 1 - || Self::field_columns_are_alternative_samples( - left.schema(), - &left_context.field_columns, - ), - MultiFieldsNotSupportedSnafu { - operator: "AND/UNLESS operator" - } - ); - // Generate join plan. - // All set operations in PromQL are "distinct" - let result = match op.id() { - token::T_LAND => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftSemi, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LUNLESS => LogicalPlanBuilder::from(left) - .distinct() - .context(DataFusionPlanningSnafu)? - .join_detailed( - right, - JoinType::LeftAnti, - (join_keys.clone(), join_keys), - None, - NullEquality::NullEqualsNull, - ) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu), - token::T_LOR => { - // OR is handled at the beginning of this function, as it cannot - // be expressed using JOIN like AND and UNLESS. - unreachable!() - } - _ => UnexpectedTokenSnafu { token: op }.fail(), - }?; - let result = if added_marker_to_left { - LogicalPlanBuilder::from(result) - .project(visible_left_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)? - } else { - result - }; - - // AND/UNLESS preserve the complete left operand's visible columns and values; encoded - // markers are decoded. - self.ctx = output_context; - Ok(result) - } - fn string_value_data_type(data_type: &ArrowDataType) -> Option<&ArrowDataType> { match data_type { data_type if data_type.is_string() => Some(data_type), @@ -6870,601 +5757,6 @@ impl PromPlanner { } } - // TODO(ruihang): change function name - #[allow(clippy::too_many_arguments)] - fn or_operator( - &mut self, - left: LogicalPlan, - right: LogicalPlan, - left_tag_cols_set: HashSet, - right_tag_cols_set: HashSet, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> Result { - let left_is_empty = Self::is_zero_row_empty_relation(&left); - let right_is_empty = Self::is_zero_row_empty_relation(&right); - match (left_is_empty, right_is_empty) { - (true, false) => { - self.ctx = right_context; - return Ok(right); - } - (false, true) => { - self.ctx = left_context; - return Ok(left); - } - (true, true) => { - self.ctx = left_context; - return Ok(left); - } - (false, false) => {} - } - - ensure!( - !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), - UnexpectedPlanExprSnafu { - desc: "OR operator input has zero columns", - } - ); - let left_has_alternative_samples = - Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); - let right_has_alternative_samples = Self::field_columns_are_alternative_samples( - right.schema(), - &right_context.field_columns, - ); - ensure!( - left_context.field_columns.len() == 1 || left_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - ensure!( - right_context.field_columns.len() == 1 || right_has_alternative_samples, - MultiFieldsNotSupportedSnafu { - operator: "OR operator" - } - ); - - // prepare hash sets - let all_tags = left_tag_cols_set - .union(&right_tag_cols_set) - .cloned() - .collect::>(); - let left_qualifier = left.schema().qualified_field(0).0.cloned(); - let right_qualifier = right.schema().qualified_field(0).0.cloned(); - let left_qualifier_string = left_qualifier - .as_ref() - .map(|l| l.to_string()) - .unwrap_or_default(); - let right_qualifier_string = right_qualifier - .as_ref() - .map(|r| r.to_string()) - .unwrap_or_default(); - let left_time_index_column = - left_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let right_time_index_column = - right_context - .time_index_column - .clone() - .with_context(|| TimeIndexNotFoundSnafu { - table: right_qualifier_string.clone(), - })?; - let native_histogram_type = Self::native_histogram_arrow_type(); - let is_numeric = |data_type: &ArrowDataType| { - matches!( - data_type, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) - }; - let left_fields = left_context - .field_columns - .iter() - .map(|name| { - left.schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let right_fields = right_context - .field_columns - .iter() - .map(|name| { - right - .schema() - .iter() - .find(|(_, field)| field.name() == name) - .map(|(qualifier, field)| { - (name.clone(), qualifier.cloned(), field.data_type().clone()) - }) - .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) - }) - .collect::>>()?; - let left_field = &left_fields[0]; - let right_field = &right_fields[0]; - let left_field_col = &left_field.0; - let right_field_col = &right_field.0; - let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { - fields.iter().all(|(_, _, data_type)| { - is_numeric(data_type) || data_type == &native_histogram_type - }) - }; - let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { - if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>(), - right_fields - .iter() - .map(|(_, _, data_type)| data_type) - .collect::>() - ), - } - .fail(); - } - true - } else { - (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) - || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) - }; - let target_field_type = if mixed_sample_types { - // Mixed vectors use the existing response representation: one nullable float column - // and one nullable native-histogram column. - ArrowDataType::Float64 - } else if left_field.2 == right_field.2 { - left_field.2.clone() - } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { - ArrowDataType::Float64 - } else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR value fields have incompatible types: {:?} and {:?}", - left_field.2, right_field.2 - ), - } - .fail(); - }; - let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { - let mut reserved_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - for (name, _, _) in left_fields.iter().chain(&right_fields) { - reserved_names.remove(name); - } - reserved_names.extend(all_tags.iter().cloned()); - let unique_name = |prefix: &str, reserved_names: &mut HashSet| { - let mut index = 0; - loop { - let name = format!("{prefix}{index}"); - index += 1; - if reserved_names.insert(name.clone()) { - break name; - } - } - }; - let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); - let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); - (float_field, histogram_field) - } else { - (left_field_col.clone(), String::new()) - }; - let left_tag_types = left_tag_cols_set - .iter() - .map(|label| { - left.schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let right_tag_types = right_tag_cols_set - .iter() - .map(|label| { - right - .schema() - .fields() - .iter() - .find(|field| field.name() == label) - .map(|field| (label.clone(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) - }) - .collect::>>()?; - let mut target_tag_types = HashMap::with_capacity(all_tags.len()); - for label in &all_tags { - let Some(data_type) = - Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) - else { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR label {label} has incompatible types: {:?} and {:?}", - left_tag_types.get(label), - right_tag_types.get(label) - ), - } - .fail(); - }; - target_tag_types.insert(label.clone(), data_type); - } - let left_has_tsid = left - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - let right_has_tsid = right - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); - - // step 0: fill all columns in output schema - let mut all_columns_set = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment - // (e.g. `unknown_metric or some_metric`). - if !(left_has_tsid && right_has_tsid) { - all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); - } - // remove time index column - all_columns_set.remove(&left_time_index_column); - all_columns_set.remove(&right_time_index_column); - if mixed_sample_types { - for (name, _, _) in left_fields.iter().chain(&right_fields) { - all_columns_set.remove(name); - } - all_columns_set.extend(all_tags.iter().cloned()); - all_columns_set.insert(mixed_float_field_col.clone()); - all_columns_set.insert(mixed_histogram_field_col.clone()); - } else if left_field_col != right_field_col { - // remove field column in the right - all_columns_set.remove(right_field_col); - } - let mut all_columns = all_columns_set.into_iter().collect::>(); - // sort to ensure the generated schema is not volatile - all_columns.sort_unstable(); - // use left time index column name as the result time index column name - all_columns.insert(0, left_time_index_column.clone()); - let mut occupied_column_names = left - .schema() - .fields() - .iter() - .chain(right.schema().fields().iter()) - .map(|field| field.name().clone()) - .collect::>(); - - // step 1: align schema using project, fill non-exist columns with null - let aligned_label_expr = |col: &String, source_types: &HashMap| { - let target_type = &target_tag_types[col]; - if let Some(source_type) = source_types.get(col) { - let expr = DfExpr::Column(Column::new(None::, col)); - if source_type == target_type { - expr - } else { - DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) - } - } else { - DfExpr::Literal( - Self::string_scalar_value(target_type, None) - .expect("target label type is a string"), - None, - ) - .alias(col.clone()) - } - }; - let null_histogram = - ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; - let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], - output_col: &String| { - if output_col == &mixed_float_field_col { - if let Some((name, qualifier, data_type)) = fields - .iter() - .find(|(_, _, data_type)| is_numeric(data_type)) - { - let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); - if data_type == &ArrowDataType::Float64 { - expr.alias(output_col) - } else { - DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) - .alias(output_col) - } - } else { - DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) - } - } else { - fields - .iter() - .find(|(_, _, data_type)| data_type == &native_histogram_type) - .map(|(name, qualifier, _)| { - DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) - }) - .unwrap_or_else(|| { - DfExpr::Literal(null_histogram.clone(), None).alias(output_col) - }) - } - }; - let left_proj_exprs = all_columns.iter().map(|col| { - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&left_fields, col) - } else if !mixed_sample_types - && col == left_field_col - && left_field.2 != target_field_type - { - DfExpr::Cast(Cast::new( - Box::new(DfExpr::Column(Column::new( - left_field.1.clone(), - left_field_col, - ))), - target_field_type.clone(), - )) - .alias(left_field_col.clone()) - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &left_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_time_index_expr = DfExpr::Column(Column::new( - right_qualifier.clone(), - right_time_index_column, - )) - .alias(left_time_index_column.clone()); - // The field column in right side may not have qualifier (it may be removed by join operation), - // so we need to find it from the schema. - // `skip（1)` to skip the time index column - let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { - // expr - if mixed_sample_types - && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) - { - mixed_value_expr(&right_fields, col) - } else if !mixed_sample_types && col == left_field_col { - let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); - if right_field.2 != target_field_type { - DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) - .alias(left_field_col.clone()) - } else if left_field_col != right_field_col { - expr.alias(left_field_col.clone()) - } else { - expr - } - } else if target_tag_types.contains_key(col) { - aligned_label_expr(col, &right_tag_types) - } else { - DfExpr::Column(Column::new(None::, col)) - } - }); - let right_proj_exprs = [right_time_index_expr] - .into_iter() - .chain(right_proj_exprs_without_time_index); - - let left_projected = LogicalPlanBuilder::from(left) - .project(left_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(left_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_projected = LogicalPlanBuilder::from(right) - .project(right_proj_exprs) - .context(DataFusionPlanningSnafu)? - .alias(right_qualifier_string.clone()) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 2: compute match columns - let mut match_columns = if let Some(modifier) = modifier - && let Some(matching) = &modifier.matching - { - match matching { - // keeps columns mentioned in `on` - LabelModifier::Include(on) => on.labels.clone(), - // removes columns memtioned in `ignoring` - LabelModifier::Exclude(ignoring) => { - let ignoring = ignoring.labels.iter().cloned().collect::>(); - all_tags.difference(&ignoring).cloned().collect() - } - } - } else { - all_tags.iter().cloned().collect() - }; - // sort to ensure the generated plan is not volatile - match_columns.sort_unstable(); - match_columns.dedup(); - occupied_column_names.extend( - left_projected - .schema() - .fields() - .iter() - .chain(right_projected.schema().fields().iter()) - .map(|field| field.name().clone()), - ); - - let visible_schema = left_projected.schema().clone(); - let visible_left_exprs = left_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let visible_right_exprs = right_projected - .schema() - .iter() - .map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - }) - .collect::>(); - let mut left_match_exprs = Vec::with_capacity(match_columns.len()); - let mut right_match_exprs = Vec::with_capacity(match_columns.len()); - let mut next_internal_column = 0; - - for label in &match_columns { - let left_field = if left_tag_cols_set.contains(label) { - Some( - left_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let right_field = if right_tag_cols_set.contains(label) { - Some( - right_projected - .schema() - .iter() - .find(|(_, field)| field.name() == label) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, - ) - } else { - None - }; - let data_type = match (left_field.as_ref(), right_field.as_ref()) { - (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { - left_type.clone() - } - (Some((_, left_type)), Some((_, right_type))) => { - return UnexpectedPlanExprSnafu { - desc: format!( - "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" - ), - } - .fail(); - } - (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), - (None, None) => ArrowDataType::Utf8, - }; - let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { - return UnexpectedPlanExprSnafu { - desc: format!("OR match label {label} must be a string"), - } - .fail(); - }; - let internal_name = loop { - let name = format!("__promql_or_match_{next_internal_column}"); - next_internal_column += 1; - if occupied_column_names.insert(name.clone()) { - break name; - } - }; - left_match_exprs.push(Self::normalized_match_key_expr( - label, - left_field, - &value_type, - &internal_name, - )); - right_match_exprs.push(Self::normalized_match_key_expr( - label, - right_field, - &value_type, - &internal_name, - )); - } - - let left_augmented = LogicalPlanBuilder::from(left_projected) - .project(visible_left_exprs.into_iter().chain(left_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - let right_augmented = LogicalPlanBuilder::from(right_projected) - .project(visible_right_exprs.into_iter().chain(right_match_exprs)) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 3: build `UnionDistinctOn` with normalized internal match keys. - let visible_field_count = visible_schema.fields().len(); - let compare_key_indices = - (visible_field_count..visible_field_count + match_columns.len()).collect::>(); - let (time_qualifier, _) = visible_schema - .iter() - .find(|(_, field)| field.name() == &left_time_index_column) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let ts_col_idx = left_augmented - .schema() - .iter() - .position(|(qualifier, field)| { - qualifier == time_qualifier && field.name() == &left_time_index_column - }) - .with_context(|| TimeIndexNotFoundSnafu { - table: left_qualifier_string.clone(), - })?; - let union_distinct_on = UnionDistinctOn::try_new( - left_augmented, - right_augmented, - compare_key_indices, - ts_col_idx, - ) - .context(DataFusionPlanningSnafu)?; - let augmented_result = LogicalPlan::Extension(Extension { - node: Arc::new(union_distinct_on), - }); - let result = LogicalPlanBuilder::from(augmented_result) - .project(visible_schema.iter().map(|(qualifier, field)| { - DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) - })) - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; - - // step 4: update context - let output_field_col = left_field_col.clone(); - let mut output_context = left_context; - let mut visible_tags = all_tags.into_iter().collect::>(); - visible_tags.sort_unstable(); - output_context.time_index_column = Some(left_time_index_column); - output_context.tag_columns = visible_tags; - output_context.field_columns = if mixed_sample_types { - vec![mixed_float_field_col, mixed_histogram_field_col] - } else { - vec![output_field_col] - }; - output_context.use_tsid = left_has_tsid && right_has_tsid; - self.ctx = output_context; - - Ok(result) - } - /// Build a projection that project and perform operation expr for every value columns. /// Non-value columns (tag and timestamp) will be preserved in the projection. /// @@ -7705,7661 +5997,4 @@ enum ScalarFunc { } #[cfg(test)] -mod test { - use std::time::{Duration, UNIX_EPOCH}; - - use catalog::RegisterTableRequest; - use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; - use common_base::Plugins; - use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_query::native_histogram::{ - CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, - }; - use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; - use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use common_query::test_util::DummyDecoder; - use common_recordbatch::RecordBatch as GreptimeRecordBatch; - use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, - }; - use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::datasource::{MemTable, provider_as_source}; - use datafusion::execution::context::SessionContext; - use datafusion::logical_expr::Extension; - use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{ColumnSchema, Schema}; - use promql_parser::label::Labels; - use promql_parser::parser; - use session::context::QueryContext; - use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use table::Table; - use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; - use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; - - use super::*; - use crate::QueryEngineContext; - use crate::options::QueryOptions; - use crate::parser::QueryLanguageParser; - use crate::query_engine::DefaultSerializer; - - mod delta; - - fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { - if let LogicalPlan::Extension(Extension { node }) = plan - && let Some(instant_manipulate) = node.as_any().downcast_ref::() - { - return Some(instant_manipulate); - } - - plan.inputs().into_iter().find_map(find_instant_manipulate) - } - - fn build_query_engine_state() -> QueryEngineState { - QueryEngineState::new( - new_memory_catalog_manager().unwrap(), - None, - None, - None, - None, - None, - false, - Plugins::default(), - QueryOptions::default(), - ) - } - - #[test] - fn common_label_type_preserves_only_shared_dictionary_encoding() { - let dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::UInt32), - Box::new(ArrowDataType::Utf8), - ); - let other_dictionary = ArrowDataType::Dictionary( - Box::new(ArrowDataType::Int32), - Box::new(ArrowDataType::Utf8), - ); - - assert_eq!( - Some(dictionary.clone()), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) - ); - assert_eq!( - Some(ArrowDataType::Utf8), - PromPlanner::common_label_data_type(Some(&dictionary), None) - ); - } - - async fn build_optimized_promql_plan( - table_provider: DfTableSourceProvider, - eval_stmt: &EvalStmt, - ) -> LogicalPlan { - let state = build_query_engine_state(); - let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) - .await - .unwrap(); - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - state - .optimize_by_extension_rules(raw_plan, &context) - .unwrap() - } - - async fn build_optimized_tsid_plan( - query: &str, - num_tag: usize, - num_field: usize, - end_secs: u64, - lookback_secs: u64, - ) -> String { - let eval_stmt = EvalStmt { - expr: parser::parse(query).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(end_secs)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(lookback_secs), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - num_tag, - num_field, - ) - .await; - - build_optimized_promql_plan(table_provider, &eval_stmt) - .await - .display_indent_schema() - .to_string() - } - - async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); - assert!(plan_str.contains("Distinct:")); - assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - } - - async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { - let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; - assert!(!plan_str.contains("Distinct:"), "{plan_str}"); - } - - fn build_eval_stmt(expr: &str) -> EvalStmt { - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - } - } - - enum DirectOrValue { - Float64(f64), - Int64(i64), - NativeHistogram(NativeHistogram), - Utf8(&'static str), - } - - impl DirectOrValue { - fn data_type(&self) -> ArrowDataType { - match self { - Self::Float64(_) => ArrowDataType::Float64, - Self::Int64(_) => ArrowDataType::Int64, - Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), - Self::Utf8(_) => ArrowDataType::Utf8, - } - } - fn array(&self) -> Arc { - match self { - Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), - Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), - Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), - Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), - } - } - } - - fn direct_or_histogram() -> NativeHistogram { - NativeHistogram { - schema: 0, - zero_threshold: 0.0, - sum: 1.0, - reset_hint: CounterResetHint::Unknown, - start_timestamp: None, - custom_values: vec![], - positive_spans: vec![], - negative_spans: vec![], - count: 1.0, - zero_count: 1.0, - positive_buckets: vec![], - negative_buckets: vec![], - } - } - - fn operator_metric_table( - name: &str, - table_id: u32, - tag: &str, - le: Option<&str>, - value: DirectOrValue, - ) -> table::TableRef { - let value_type = match &value { - DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), - DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), - DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), - DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), - }; - let tag_count = 1 + usize::from(le.is_some()); - let mut columns = vec![ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if le.is_some() { - columns.push(ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.extend([ - ColumnSchema::new( - "ts".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new("v".to_string(), value_type, true), - ]); - let schema = Arc::new(Schema::new(columns)); - let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; - if let Some(le) = le { - arrays.push(Arc::new(StringArray::from(vec![le]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, - value.array(), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - name, - GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), - table_id, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let value_index = tag_count + 1; - let meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![value_index]) - .next_column_id((value_index + 1) as u32) - .build() - .unwrap(); - let info = Arc::new( - TableInfoBuilder::default() - .table_id(table_id) - .name(name) - .meta(meta) - .build() - .unwrap(), - ); - Arc::new(Table::new( - info, - FilterPushDownType::Unsupported, - backing.data_source(), - )) - } - - fn operator_table_provider() -> DfTableSourceProvider { - let catalog = MemoryCatalogManager::with_default_setup(); - let tables = [ - operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), - operator_metric_table( - "lh", - 2_002, - "b", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), - operator_metric_table( - "rh", - 2_004, - "a", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), - operator_metric_table( - "bad_classic", - 2_006, - "d", - Some("broken"), - DirectOrValue::Float64(1.0), - ), - operator_metric_table( - "bad_native", - 2_007, - "d", - None, - DirectOrValue::NativeHistogram(direct_or_histogram()), - ), - ]; - for table in tables { - let info = table.table_info(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: info.name.clone(), - table_id: info.ident.table_id, - table, - }) - .unwrap(); - } - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn operator_eval_stmt(expr: &str) -> EvalStmt { - let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); - EvalStmt { - expr: parser::parse(expr).unwrap(), - start: time, - end: time, - interval: Duration::from_secs(1), - lookback_delta: Duration::from_secs(5), - } - } - - struct DirectOrSource { - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - } - - fn source( - name: &'static str, - empty: bool, - timestamp: i64, - tags: Vec<(&'static str, Option<&'static str>)>, - value: DirectOrValue, - ) -> DirectOrSource { - DirectOrSource { - name, - empty, - timestamp, - tags, - value, - } - } - - fn tagged_source( - name: &'static str, - empty: bool, - tag: (&'static str, Option<&'static str>), - value: DirectOrValue, - ) -> DirectOrSource { - source(name, empty, 1, vec![("job", Some("job")), tag], value) - } - - fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { - source(name, true, 1, vec![("job", Some("job"))], value) - } - - fn table(source: &DirectOrSource) -> Arc { - let mut fields = vec![Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )]; - fields.extend( - source - .tags - .iter() - .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), - ); - fields.push(Field::new("v", source.value.data_type(), true)); - let schema = Arc::new(ArrowSchema::new(fields)); - let partitions = if source.empty { - vec![vec![]] - } else { - let mut columns: Vec> = - vec![Arc::new(TimestampMillisecondArray::from(vec![ - source.timestamp, - ]))]; - columns.extend( - source - .tags - .iter() - .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), - ); - columns.push(source.value.array()); - vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] - }; - Arc::new(MemTable::try_new(schema, partitions).unwrap()) - } - - fn scan(source: &DirectOrSource) -> LogicalPlan { - LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) - .unwrap() - .build() - .unwrap() - } - - fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { - PromPlannerContext { - table_name: Some(qualifier.to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![field.to_string()], - tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), - ..Default::default() - } - } - - fn or_modifier(expr: &str) -> Option { - let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { - unreachable!() - }; - expr.modifier - } - - async fn plan_direct_or( - left: LogicalPlan, - right: LogicalPlan, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> LogicalPlan { - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - planner - .or_operator( - left, - right, - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - modifier, - ) - .unwrap() - } - - async fn execute( - plan: LogicalPlan, - state: &QueryEngineState, - ) -> (LogicalPlan, Vec) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - (optimized, batches) - } - - async fn run( - left: &DirectOrSource, - right: &DirectOrSource, - left_context: PromPlannerContext, - right_context: PromPlannerContext, - modifier: &Option, - ) -> (LogicalPlan, Vec) { - let plan = plan_direct_or( - scan(left), - scan(right), - left_context, - right_context, - modifier, - ) - .await; - execute(plan, &build_query_engine_state()).await - } - - async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { - let sample = |histogram: bool| { - if histogram { - DirectOrValue::NativeHistogram(direct_or_histogram()) - } else { - DirectOrValue::Float64(1.25) - } - }; - let left = tagged_source( - "lhs", - false, - ( - "k", - Some(if histogram_on_left { - "histogram" - } else { - "float" - }), - ), - sample(histogram_on_left), - ); - let right = tagged_source( - "rhs", - false, - ( - "k", - Some(if histogram_on_left { - "float" - } else { - "histogram" - }), - ), - sample(!histogram_on_left), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job", "k"], "v"); - let right_context = direct_or_context("rhs", &["job", "k"], "v"); - let plan = planner - .or_operator( - scan(&left), - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - (planner, plan) - } - - async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { - let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); - let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); - let row_count = histograms.len() + 1; - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "ts", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new("k", ArrowDataType::Utf8, false), - Field::new(&float_field, ArrowDataType::Float64, true), - Field::new( - &histogram_field, - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let mut histogram_values = Vec::with_capacity(row_count); - histogram_values.push(None); - histogram_values.extend(histograms.into_iter().map(Some)); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), - Arc::new(StringArray::from_iter_values( - (0..row_count).map(|row| format!("kind_{row}")), - )), - Arc::new(Float64Array::from_iter( - (0..row_count).map(|row| (row == 0).then_some(1.25)), - )), - build_histogram_array(&histogram_values), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let planner = PromPlanner { - table_provider, - ctx: PromPlannerContext { - table_name: Some("mixed".to_string()), - time_index_column: Some("ts".to_string()), - field_columns: vec![float_field, histogram_field], - tag_columns: vec!["k".to_string()], - ..Default::default() - }, - promql_annotations: None, - }; - (planner, plan) - } - - fn assert_no_internal_or_keys(schema: &DFSchema) { - assert!( - schema - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")), - "{schema:?}" - ); - } - - fn values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect() - } - - fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = datafusion::arrow::compute::cast( - batch.column_by_name(column).unwrap(), - &ArrowDataType::Float64, - ) - .unwrap(); - values - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect::>() - }) - .collect() - } - - fn histograms(batches: &[RecordBatch], column: &str) -> Vec { - batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name(column) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).filter_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .collect() - } - - fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { - let mut rows = batches - .iter() - .flat_map(|batch| { - let values = batch - .column_by_name("v") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let labels = batch - .column_by_name("k") - .map(|column| column.as_any().downcast_ref::().unwrap()); - (0..batch.num_rows()).map(move |i| { - ( - values.value(i), - labels.and_then(|labels| { - (!labels.is_null(i)).then(|| labels.value(i).to_string()) - }), - ) - }) - }) - .collect::>(); - rows.sort_by(|left, right| left.0.total_cmp(&right.0)); - rows - } - - fn matrix_source( - name: &'static str, - k: Option>, - timestamp: i64, - value: f64, - ) -> DirectOrSource { - let mut tags = vec![("job", Some("job"))]; - if let Some(k) = k { - tags.push(("k", k)); - } - source(name, false, timestamp, tags, DirectOrValue::Float64(value)) - } - - fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { - direct_or_context( - name, - if k.is_some() { &["job", "k"] } else { &["job"] }, - "v", - ) - } - - async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), - ], - &["pod", "instance"], - ) - .await - } - - fn assert_normal_metric_schema(plan: &LogicalPlan) { - let fields = plan.schema().fields(); - assert_eq!(fields.len(), 4, "{fields:?}"); - assert!( - fields.iter().any(|field| field.name() == "pod"), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| field.name() == "instance"), - "{fields:?}" - ); - assert!( - fields - .iter() - .any(|field| field.name() == greptime_timestamp()), - "{fields:?}" - ); - assert!( - fields.iter().any(|field| { - field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 - }), - "{fields:?}" - ); - } - - async fn build_test_table_provider_with_distinct_tags( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn contains_histogram_fold(plan: &LogicalPlan) -> bool { - matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) - || plan.inputs().into_iter().any(contains_histogram_fold) - } - - async fn build_set_op_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("bucket_metric", &["job", "le"]), - ("normal_metric", &["job"]), - ("fallback_metric", &["instance"]), - ]) - .await - } - - async fn build_or_context_table_provider() -> DfTableSourceProvider { - build_test_table_provider_with_distinct_tags(&[ - ("normal_metric", &["job"]), - ("other_metric", &["instance"]), - ("non_hist_metric", &["instance"]), - ]) - .await - } - - async fn optimize_and_create_physical_plan( - state: &QueryEngineState, - plan: LogicalPlan, - ) -> ( - LogicalPlan, - Arc, - ) { - let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); - let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); - let physical = state - .session_state() - .create_physical_plan(&optimized) - .await - .unwrap(); - (optimized, physical) - } - - async fn build_test_table_provider( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - for i in 0..num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - build_test_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - "native_histogram_2".to_string(), - native_histogram_value_type().clone(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2, 3]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_mixed_native_histogram_table_provider( - table_name: &str, - ) -> DfTableSourceProvider { - build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await - } - - async fn build_test_mixed_native_histogram_table_provider_with_marker( - table_name: &str, - temporality_marker: bool, - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - if temporality_marker { - columns.push(ColumnSchema::new( - OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), - ConcreteDataType::string_datatype(), - true, - )); - } - let tag_count = columns.len(); - columns.extend([ - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices((0..tag_count).collect()) - .value_indices(vec![tag_count + 1, tag_count + 2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let mut arrays: Vec> = - vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; - if temporality_marker { - arrays.push(Arc::new(StringArray::from(vec![ - Some(GREPTIME_TEMPORALITY_DELTA), - Some(GREPTIME_TEMPORALITY_DELTA), - ]))); - } - arrays.extend([ - Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, - build_histogram_array(&[None, Some(direct_or_histogram())]), - Arc::new(Float64Array::from(vec![Some(2.0), None])), - ]); - let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 1024, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - fn classic_and_native_histogram_table_provider( - native_tag: &str, - native_le: Option<&str>, - native_histogram: NativeHistogram, - ) -> DfTableSourceProvider { - let table_name = "mixed_histogram"; - let catalog = MemoryCatalogManager::with_default_setup(); - let schema = Arc::new(Schema::new(vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - LE_COLUMN_NAME.to_string(), - ConcreteDataType::string_datatype(), - true, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - greptime_native_histogram().to_string(), - native_histogram_value_type().clone(), - true, - ), - ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ])); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices(vec![0, 1]) - .value_indices(vec![3, 4]) - .next_column_id(5) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(), - ); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from(vec![ - "classic", "classic", native_tag, "classic", "classic", native_tag, - ])), - Arc::new(StringArray::from(vec![ - Some("1"), - Some("+Inf"), - native_le, - Some("1"), - Some("+Inf"), - native_le, - ])), - Arc::new(TimestampMillisecondArray::from(vec![ - 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, - ])), - build_histogram_array(&[ - None, - None, - Some(native_histogram.clone()), - None, - None, - Some(native_histogram), - ]), - Arc::new(Float64Array::from(vec![ - Some(2.0), - Some(4.0), - None, - Some(2.0), - Some(4.0), - None, - ])), - ], - ) - .unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - table_name, - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 2_200, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id: 2_200, - table, - }) - .unwrap(); - - DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_tsid( - table_name_tuples: &[(String, String)], - num_tag: usize, - num_field: usize, - ) -> DfTableSourceProvider { - let table_specs = table_name_tuples - .iter() - .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) - .collect::>(); - build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await - } - - async fn build_test_table_provider_with_tsid_fields( - table_specs: &[((String, String), usize)], - num_tag: usize, - ) -> DfTableSourceProvider { - let table_specs = table_specs - .iter() - .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) - .collect::>(); - build_test_table_provider_with_tsid_tag_fields(&table_specs).await - } - - async fn build_test_table_provider_with_tsid_tag_fields( - table_specs: &[((String, String), usize, usize)], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - - let physical_table_name = "phy"; - let physical_table_id = 999u32; - let physical_num_tag = table_specs - .iter() - .map(|(_, num_tag, _)| *num_tag) - .max() - .unwrap_or(0); - let physical_num_field = table_specs - .iter() - .map(|(_, _, num_field)| *num_field) - .max() - .unwrap_or(0); - - // Register a metric engine physical table with internal columns. - { - let mut columns = vec![ - ColumnSchema::new( - DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), - ConcreteDataType::uint32_datatype(), - false, - ), - ColumnSchema::new( - DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), - ConcreteDataType::uint64_datatype(), - false, - ), - ]; - for i in 0..physical_num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..physical_num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(primary_key_indices) - .value_indices( - (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), - ) - .engine(METRIC_ENGINE_NAME.to_string()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(physical_table_id) - .name(physical_table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: physical_table_name.to_string(), - table_id: physical_table_id, - table, - }) - .is_ok() - ); - } - - // Register metric engine logical tables without `__tsid`, referencing the physical table. - for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() - { - let mut columns = vec![]; - for i in 0..*num_tag { - columns.push(ColumnSchema::new( - format!("tag_{i}"), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - for i in 0..*num_field { - columns.push(ColumnSchema::new( - format!("field_{i}"), - ConcreteDataType::float64_datatype(), - true, - )); - } - - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options.extra_options.insert( - LOGICAL_TABLE_METADATA_KEY.to_string(), - physical_table_name.to_string(), - ); - let table_id = 1024u32 + idx as u32; - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..*num_tag).collect()) - .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(table_id) - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_test_table_provider_with_fields( - table_name_tuples: &[(String, String)], - tags: &[&str], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (schema_name, table_name) in table_name_tuples { - let mut columns = vec![]; - let num_tag = tags.len(); - for tag in tags { - columns.push(ColumnSchema::new( - tag.to_string(), - ConcreteDataType::string_datatype(), - false, - )); - } - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices((0..num_tag).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name.clone()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: schema_name.clone(), - table_name: table_name.clone(), - table_id: 1024, - table, - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - // { - // input: `abs(some_metric{foo!="bar"})`, - // expected: &Call{ - // Func: MustGetFunction("abs"), - // Args: Expressions{ - // &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // }, - // }, - // }, - async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { - let prom_expr = - parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" - ).replace("TEMPLATE", plan_name); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_abs() { - do_single_instant_function_call("abs", "abs").await; - } - - #[tokio::test] - #[should_panic] - async fn single_absent() { - do_single_instant_function_call("absent", "").await; - } - - #[tokio::test] - async fn single_ceil() { - do_single_instant_function_call("ceil", "ceil").await; - } - - #[tokio::test] - async fn single_exp() { - do_single_instant_function_call("exp", "exp").await; - } - - #[tokio::test] - async fn single_ln() { - do_single_instant_function_call("ln", "ln").await; - } - - #[tokio::test] - async fn single_log2() { - do_single_instant_function_call("log2", "log2").await; - } - - #[tokio::test] - async fn single_log10() { - do_single_instant_function_call("log10", "log10").await; - } - - #[tokio::test] - #[should_panic] - async fn single_scalar() { - do_single_instant_function_call("scalar", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sgn() { - do_single_instant_function_call("sgn", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort() { - do_single_instant_function_call("sort", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sort_desc() { - do_single_instant_function_call("sort_desc", "").await; - } - - #[tokio::test] - async fn single_sqrt() { - do_single_instant_function_call("sqrt", "sqrt").await; - } - - #[tokio::test] - async fn single_timestamp_plan_preserves_source_value() { - let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn single_acos() { - do_single_instant_function_call("acos", "acos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_acosh() { - do_single_instant_function_call("acosh", "").await; - } - - #[tokio::test] - async fn single_asin() { - do_single_instant_function_call("asin", "asin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_asinh() { - do_single_instant_function_call("asinh", "").await; - } - - #[tokio::test] - async fn single_atan() { - do_single_instant_function_call("atan", "atan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_atanh() { - do_single_instant_function_call("atanh", "").await; - } - - #[tokio::test] - async fn single_cos() { - do_single_instant_function_call("cos", "cos").await; - } - - #[tokio::test] - #[should_panic] - async fn single_cosh() { - do_single_instant_function_call("cosh", "").await; - } - - #[tokio::test] - async fn single_sin() { - do_single_instant_function_call("sin", "sin").await; - } - - #[tokio::test] - #[should_panic] - async fn single_sinh() { - do_single_instant_function_call("sinh", "").await; - } - - #[tokio::test] - async fn single_tan() { - do_single_instant_function_call("tan", "tan").await; - } - - #[tokio::test] - #[should_panic] - async fn single_tanh() { - do_single_instant_function_call("tanh", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_deg() { - do_single_instant_function_call("deg", "").await; - } - - #[tokio::test] - #[should_panic] - async fn single_rad() { - do_single_instant_function_call("rad", "").await; - } - - // { - // input: "avg by (foo)(some_metric)", - // expected: &AggregateExpr{ - // Op: AVG, - // Expr: &VectorSelector{ - // Name: "some_metric", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // PosRange: PositionRange{ - // Start: 13, - // End: 24, - // }, - // }, - // Grouping: []string{"foo"}, - // PosRange: PositionRange{ - // Start: 0, - // End: 25, - // }, - // }, - // }, - async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { - let prom_expr = parser::parse(&format!( - "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", - )) - .unwrap(); - let mut eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // test group by - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_no_without = String::from( - "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!( - plan.display_indent_schema().to_string(), - expected_no_without - ); - - // test group without - if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { - *modifier = Some(LabelModifier::Exclude(Labels { - labels: vec![String::from("tag_1")].into_iter().collect(), - })); - } - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 2, - 2, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected_without = String::from( - "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" - ).replace("TEMPLATE", plan_name); - assert_eq!(plan.display_indent_schema().to_string(), expected_without); - } - - #[tokio::test] - async fn aggregate_sum() { - do_aggregate_expr_plan("sum", "sum").await; - } - - #[tokio::test] - async fn tsid_is_used_for_series_divide_when_available() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - assert!(plan_str.contains("__tsid ASC NULLS FIRST")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - let manipulate = find_instant_manipulate(&plan).unwrap(); - let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new( - &[], - Arc::new( - datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0] - .schema() - .as_arrow() - .clone(), - ), - None, - ) - .unwrap(), - )))); - assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); - } - - #[tokio::test] - async fn default_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn reject_binary_fill_modifiers() { - let state = build_query_engine_state(); - - for query in [ - "some_metric + fill(0) some_alt_metric", - "some_metric + fill_left(0) some_alt_metric", - "some_metric + fill_right(0) some_alt_metric", - "(some_metric + fill(0) some_alt_metric) + some_metric", - ] { - let eval_stmt = build_eval_stmt(query); - let table_provider = build_test_table_provider(&[], 0, 0).await; - let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) - .await - .unwrap_err(); - - assert!( - matches!( - &err, - crate::promql::error::Error::UnsupportedExpr { name, .. } - if name == "PromQL fill modifiers" - ), - "{err}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { - for query in [ - "timestamp(some_metric) / some_metric", - "some_metric / timestamp(some_metric)", - ] { - let eval_stmt = build_eval_stmt(query); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); - assert!( - plan_str.contains("lhs.tag_0 = rhs.tag_0"), - "{query}: {plan_str}" - ); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { - let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), - 2, - 1, - ), - ( - (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - - assert!( - plan_str.contains("Boolean(false)") || plan_str.contains("false"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_is_preserved_for_nested_default_binary_joins() { - let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_third_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_leaf_plan() { - let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { - let eval_stmt = - build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1025)") - .count(), - 1, - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - } - - #[tokio::test] - async fn binary_island_reuses_self_operand_without_join() { - let eval_stmt = build_eval_stmt("some_metric / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 1, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_reuses_leaf_across_two_branches() { - let eval_stmt = - build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str - .matches("Filter: phy.__table_id = UInt32(1024)") - .count(), - 1, - "{plan_str}" - ); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_generated_alias_avoids_user_column_names() { - let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - &["prom_v0", "__prom_v0"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan.schema().field_names(); - assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); - assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_clears_qualifier_for_nested_unary_projection() { - let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 2, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_distinct_matcher_leaves() { - let eval_stmt = build_eval_stmt( - "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_keeps_offset_leaves_distinct() { - let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_group_modifier() { - let eval_stmt = build_eval_stmt( - "(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric", - ); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn binary_island_falls_back_for_comparison_filter() { - let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); - assert_eq!( - plan_str.matches("PromInstantManipulate").count(), - 3, - "{plan_str}" - ); - } - - #[tokio::test] - async fn tsid_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - let value_columns = field_names - .iter() - .filter(|name| { - *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME - }) - .count(); - assert_eq!(value_columns, 1, "{field_names:?}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_shorter_field_side() { - let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); - - let table_provider = build_test_table_provider_with_tsid_fields( - &[ - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "two_field_metric".to_string(), - ), - 2, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "one_field_metric".to_string(), - ), - 1, - ), - ], - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let field_names = plan - .schema() - .fields() - .iter() - .map(|field| field.name().clone()) - .collect::>(); - assert!( - field_names.iter().any(|name| name == "field_0"), - "{field_names:?}" - ); - assert!( - !field_names.iter().any(|name| name == "field_1"), - "{field_names:?}" - ); - } - - #[tokio::test] - async fn label_matching_modifier_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn ignoring_absent_label_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { - let eval_stmt = - build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_full_label_set_keeps_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_partial_label_set_disables_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { - let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ - ( - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - 2, - 1, - ), - ( - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - 1, - 1, - ), - ]) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("__tsid ="), "{plan_str}"); - assert!( - plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { - let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), - promql_annotations: None, - }; - let plan = planner - .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), - "{plan_str}" - ); - assert!(planner.ctx.use_tsid, "{plan_str}"); - } - - #[tokio::test] - async fn comparison_bool_binary_join_uses_tsid_when_available() { - let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); - - let table_provider = build_test_table_provider_with_tsid( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "some_alt_metric".to_string(), - ), - ], - 2, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), - "{plan_str}" - ); - assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); - assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); - } - - #[tokio::test] - async fn scalar_count_count_range_keeps_full_window() { - let plan_str = build_optimized_tsid_plan( - "scalar(count(count(some_metric) by (tag_0)))", - 1, - 1, - 100_000, - 1, - ) - .await; - assert!(plan_str.contains("ScalarCalculate: tags=[]")); - assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); - assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); - } - - #[tokio::test] - async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { - let plan_str = build_optimized_tsid_plan( - "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", - 2, - 1, - 10, - 300, - ) - .await; - assert!(plan_str.contains("Distinct:"), "{plan_str}"); - } - - #[tokio::test] - async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(count(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { - assert_nested_count_rewrite_applies( - "count(sum(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" - ) - .await; - } - - #[tokio::test] - async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { - for (query, expected_outer_agg) in [ - ( - "count(avg(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", - ), - ( - "count(min(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", - ), - ( - "count(max(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", - ), - ( - "count(stddev(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", - ), - ( - "count(stdvar(some_metric) by (tag_0))", - "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", - ), - ] { - assert_nested_count_rewrite_applies(query, expected_outer_agg).await; - } - } - - #[tokio::test] - async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { - let count_plan = - build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1) - .await; - assert!( - !count_plan.contains("some_metric.field_0 IS NOT NULL"), - "{count_plan}" - ); - - for query in [ - "count(sum(some_metric) by (tag_0))", - "count(avg(some_metric) by (tag_0))", - "count(min(some_metric) by (tag_0))", - "count(max(some_metric) by (tag_0))", - "count(stddev(some_metric) by (tag_0))", - "count(stdvar(some_metric) by (tag_0))", - ] { - let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; - assert!( - plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), - "{query}: {plan_str}" - ); - } - } - - #[tokio::test] - async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { - assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; - assert_nested_count_rewrite_missing( - "count(sum(irate(some_metric[1h])) by (tag_0))", - 2, - 300, - ) - .await; - } - - #[tokio::test] - async fn physical_table_name_is_not_leaked_in_plan() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("TableScan: phy"), "{plan}"); - assert!(plan_str.contains("SubqueryAlias: some_metric")); - assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); - assert!(!plan_str.contains("TableScan: some_metric")); - } - - #[tokio::test] - async fn sum_without_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn topk_without_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn sum_by_does_not_group_by_tsid() { - let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn aggregate_over_binary_time_function_expr() { - for op in ["sum", "min", "max", "avg"] { - let prom_expr = parser::parse(&format!( - "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" - )) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - let aggr_line = plan_str - .lines() - .find(|line| line.contains("Aggregate: groupBy=")) - .unwrap(); - assert!(aggr_line.contains(op), "{plan_str}"); - assert!(aggr_line.contains("first_value"), "{plan_str}"); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) - ); - } - } - - #[tokio::test] - async fn topk_by_does_not_partition_by_tsid() { - let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - - let window_line = plan_str - .lines() - .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) - .unwrap(); - let partition_by = window_line - .split("PARTITION BY [") - .nth(1) - .and_then(|s| s.split("] ORDER BY").next()) - .unwrap(); - assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); - } - - #[tokio::test] - async fn selector_matcher_on_tsid_does_not_use_internal_column() { - let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { - if let LogicalPlan::Filter(filter) = plan { - datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); - } - for input in plan.inputs() { - collect_filter_cols(input, out); - } - } - - let mut filter_cols = HashSet::new(); - collect_filter_cols(&plan, &mut filter_cols); - assert!( - !filter_cols - .iter() - .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn tsid_is_not_used_when_physical_table_is_missing() { - let prom_expr = parser::parse("some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let catalog_list = MemoryCatalogManager::with_default_setup(); - - // Register a metric engine logical table referencing a missing physical table. - let mut columns = vec![ColumnSchema::new( - "tag_0".to_string(), - ConcreteDataType::string_datatype(), - false, - )]; - columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - "field_0".to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - let schema = Arc::new(Schema::new(columns)); - let mut options = table::requests::TableOptions::default(); - options - .extra_options - .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .engine(METRIC_ENGINE_NAME.to_string()) - .options(options) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .table_id(1024) - .name("some_metric") - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "some_metric".to_string(), - table_id: 1024, - table, - }) - .unwrap(); - - let table_provider = DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); - assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); - } - - #[tokio::test] - async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { - let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - - // Merging aggregate: label set is reduced, tsid should not be carried. - let prom_expr = parser::parse("sum(some_metric)").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(!plan_str.contains("first_value")); - } - - #[tokio::test] - async fn or_operator_with_unknown_metric_does_not_require_tsid() { - let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_tsid( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert!( - !plan - .schema() - .fields() - .iter() - .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) - ); - } - - #[tokio::test] - async fn aggregate_avg() { - do_aggregate_expr_plan("avg", "avg").await; - } - - #[tokio::test] - #[should_panic] // output type doesn't match - async fn aggregate_count() { - do_aggregate_expr_plan("count", "count").await; - } - - #[tokio::test] - async fn aggregate_min() { - do_aggregate_expr_plan("min", "min").await; - } - - #[tokio::test] - async fn aggregate_max() { - do_aggregate_expr_plan("max", "max").await; - } - - #[tokio::test] - async fn aggregate_group() { - // Regression test for `group()` aggregator. - // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) - // should be plannable, and `group()` should produce constant 1 for each group. - let prom_expr = parser::parse( - "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "kubernetes_build_info".to_string(), - )], - &["cluster", "service", "job"], - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("max(Float64(1")); - } - - #[tokio::test] - async fn aggregate_stddev() { - do_aggregate_expr_plan("stddev", "stddev_pop").await; - } - - #[tokio::test] - async fn aggregate_stdvar() { - do_aggregate_expr_plan("stdvar", "var_pop").await; - } - - // TODO(ruihang): add range fn tests once exprs are ready. - - // { - // input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", - // expected: &BinaryExpr{ - // Op: ADD, - // LHS: &VectorSelector{ - // Name: "a", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // RHS: &VectorSelector{ - // Name: "sum", - // LabelMatchers: []*labels.Matcher{ - // MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), - // MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), - // }, - // }, - // VectorMatching: &VectorMatching{}, - // }, - // }, - #[tokio::test] - async fn binary_op_column_column() { - let prom_expr = - parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = String::from( - "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ - \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - async fn indie_query_plan_compare>(query: &str, expected: T) { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); - } - - #[tokio::test] - async fn binary_op_literal_column() { - let query = r#"1 + some_metric{tag_0="bar"}"#; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn binary_op_literal_literal() { - let query = r#"1 + 1"#; - let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] - TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_bool_grammar() { - let query = "some_metric != bool 1.2345"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn bool_with_additional_arithmetic() { - let query = "some_metric + (1 == bool 2)"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn simple_unary() { - let query = "-some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn increase_aggr() { - let query = "increase(some_metric[5m])"; - let expected = String::from( - "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - async fn native_histogram_plan(query: &str) -> String { - let table_provider = build_test_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent_schema().to_string() - } - - #[tokio::test] - async fn native_histogram_count_uses_native_udf() { - let plan = native_histogram_plan("histogram_count(some_metric)").await; - - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - } - - #[tokio::test] - async fn timestamp_filters_native_histogram_stale_marker_before_projection() { - let mut stale = direct_or_histogram(); - stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); - let table = operator_metric_table( - "stale_histogram", - 2_100, - "a", - None, - DirectOrValue::NativeHistogram(stale), - ); - let catalog = MemoryCatalogManager::with_default_setup(); - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "stale_histogram".to_string(), - table_id: 2_100, - table, - }) - .unwrap(); - let provider = DfTableSourceProvider::new( - catalog, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - provider, - &operator_eval_stmt("timestamp(stale_histogram)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { - let histograms = build_histogram_array(&[None]); - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new( - greptime_native_histogram(), - histograms.data_type().clone(), - true, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1_000])), - histograms, - Arc::new(Float64Array::from(vec![f64::from_bits( - PROMETHEUS_STALE_NAN_BITS, - )])), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(SeriesDivide::new( - Vec::new(), - "timestamp".to_string(), - input, - )), - }); - let input = LogicalPlan::Extension(Extension { - node: Arc::new(InstantManipulate::new( - 1_000, - 1_000, - 5_000, - 1_000, - 0, - "timestamp".to_string(), - Vec::new(), - Some(greptime_native_histogram().to_string()), - input, - )), - }); - // Match timestamp()'s parent projection, which otherwise prunes the companion lane. - let plan = LogicalPlanBuilder::from(input) - .project([col("timestamp")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - } - - #[tokio::test] - async fn native_histogram_rate_can_feed_count() { - let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; - - assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); - assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_quantile_skips_classic_fold() { - let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); - assert!(!plan.contains("HistogramFold:"), "{plan}"); - assert!(plan.contains("some_metric.le"), "{plan}"); - // The phi literal is threaded into the native quantile UDF as its second argument. - assert!(plan.contains("Float64(0.9)"), "{plan}"); - // The empty-values filter drops NULL quantile results so the output is empty - // when all native histogram samples are dropped. - assert!(plan.contains("IS NOT NULL"), "{plan}"); - } - - #[tokio::test] - async fn mixed_native_histogram_quantile_uses_histogram_field() { - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - - assert!( - plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), - "{plan}" - ); - assert!(!plan.contains("EmptyRelation"), "{plan}"); - } - - #[tokio::test] - async fn mixed_histogram_helpers_execute_classic_and_native_samples() { - let state = build_query_engine_state(); - for (query, expected) in [ - ( - "histogram_quantile(0.5, mixed_histogram)", - vec![("classic", 1.0), ("native", 0.0)], - ), - ( - "histogram_fraction(-Inf, +Inf, mixed_histogram)", - vec![("classic", 1.0), ("native", 1.0)], - ), - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); - assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let tags = batch - .column_by_name("tag") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) - }) - .collect::>(); - actual.sort_by_key(|(tag, _)| *tag); - assert_eq!(actual, expected, "{query}"); - } - } - - #[tokio::test] - async fn mixed_histogram_helpers_report_annotations() { - let state = build_query_engine_state(); - let mut native_histogram = direct_or_histogram(); - native_histogram.count = 2.0; - native_histogram.sum = f64::NAN; - for (native_tag, expected_rows, expected_warnings, expected_infos) in [ - ( - "classic", - 0, - vec!["vector contains a mix of classic and native histograms"], - vec![], - ), - ( - "native", - 2, - vec![], - vec!["input to histogram_quantile has NaN observations, result is skewed higher"], - ), - ] { - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - classic_and_native_histogram_table_provider( - native_tag, - None, - native_histogram.clone(), - ), - &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows - ); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!(warnings, expected_warnings); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { - let state = build_query_engine_state(); - let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); - stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "classic", - Some("native"), - direct_or_histogram(), - ), - &stmt, - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert_eq!( - plan_text.matches("TableScan: mixed_histogram").count(), - 1, - "{plan_text}" - ); - - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - let mut actual = batches - .iter() - .flat_map(|batch| { - let le = batch - .column_by_name(LE_COLUMN_NAME) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let timestamps = batch - .column_by_name("timestamp") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column_by_name(&value_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).map(|row| { - ( - timestamps.value(row), - (!le.is_null(row)).then(|| le.value(row).to_string()), - values.value(row), - ) - }) - }) - .collect::>(); - actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); - assert_eq!( - actual, - vec![ - (1_000, None, 1.0), - (1_000, Some("native".to_string()), 0.0), - (2_000, None, 1.0), - (2_000, Some("native".to_string()), 0.0), - ] - ); - } - - #[tokio::test] - async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { - let state = build_query_engine_state(); - for native_le in [None, Some("native")] { - for query in [ - "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", - "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", - ] { - let plan = PromPlanner::stmt_to_plan( - classic_and_native_histogram_table_provider( - "native", - native_le, - direct_or_histogram(), - ), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "native_le={native_le:?}, query={query}" - ); - } - } - } - - #[tokio::test] - async fn native_histogram_quantile_rejects_multi_field_input() { - let table_provider = build_test_multi_histogram_table_provider("some_metric").await; - let result = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_quantile(0.9, some_metric)"), - &build_query_engine_state(), - ) - .await; - - let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); - assert!( - err.to_string() - .contains("Multi fields calculation is not supported in histogram_quantile"), - "{err}" - ); - } - - #[tokio::test] - async fn native_histogram_topk_uses_drop_udf() { - let plan = native_histogram_plan("topk(1, some_metric)").await; - - assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); - assert!( - plan.contains("Filter: prom_native_histogram_drop_float") - && plan.contains("IS NOT NULL"), - "{plan}" - ); - } - - #[tokio::test] - async fn mixed_or_topk_bottomk_ignore_native_histograms() { - for op in ["topk", "bottomk"] { - let collector = PromqlAnnotationCollector::default(); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), - "{plan:?}" - ); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!( - infos, - vec![format!( - "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" - )] - ); - } - } - - #[tokio::test] - async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { - let plan = native_histogram_plan("scalar(some_metric)").await; - - assert!(plan.contains("ScalarCalculate"), "{plan}"); - assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - #[tokio::test] - async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { - for function in ["sort", "sort_desc"] { - let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; - - assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - assert!( - !plan.contains(&format!("Sort: {}", greptime_native_histogram())), - "{plan}" - ); - assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); - } - - for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { - let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; - - assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); - assert!(plan.contains(greptime_native_histogram()), "{plan}"); - assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); - } - } - - #[tokio::test] - async fn unsupported_native_histogram_functions_use_drop_udf() { - for query in [ - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "round(some_metric)", - "abs(some_metric)", - ] { - let plan = native_histogram_plan(query).await; - - assert!( - plan.contains("prom_native_histogram_drop_float"), - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn native_histogram_absent_over_time_uses_native_udf() { - let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; - - assert!( - plan.contains("prom_native_histogram_absent_over_time"), - "{plan}" - ); - } - - #[tokio::test] - async fn native_histogram_all_function_arms_route_correctly() { - // Every native-histogram match arm in `create_function_expr` must route to the - // expected UDF when all field columns are native histograms. `holt_winters` shares - // the `double_exponential_smoothing` arm but is not registered in the promql - // parser (0.10), so it cannot be exercised through a query string. - let cases = [ - // Range functions routed to native histogram UDFs. - ( - "increase(some_metric[5m])", - "prom_native_histogram_increase", - ), - ("rate(some_metric[5m])", "prom_native_histogram_rate"), - ("delta(some_metric[5m])", "prom_native_histogram_delta"), - ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), - ("irate(some_metric[5m])", "prom_native_histogram_irate"), - ("resets(some_metric[5m])", "prom_native_histogram_resets"), - ("changes(some_metric[5m])", "prom_native_histogram_changes"), - ( - "avg_over_time(some_metric[5m])", - "prom_native_histogram_avg_over_time", - ), - ( - "sum_over_time(some_metric[5m])", - "prom_native_histogram_sum_over_time", - ), - ( - "count_over_time(some_metric[5m])", - "prom_native_histogram_count_over_time", - ), - ( - "last_over_time(some_metric[5m])", - "prom_native_histogram_last_over_time", - ), - ( - "present_over_time(some_metric[5m])", - "prom_native_histogram_present_over_time", - ), - // Unsupported functions dropped with the float-null UDF. - ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), - ( - "min_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "max_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stddev_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "stdvar_over_time(some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "quantile_over_time(0.9, some_metric[5m])", - "prom_native_histogram_drop_float", - ), - ( - "predict_linear(some_metric[5m], 60)", - "prom_native_histogram_drop_float", - ), - ( - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - "prom_native_histogram_drop_float", - ), - ("round(some_metric)", "prom_native_histogram_drop_float"), - ("rad(some_metric)", "prom_native_histogram_drop_float"), - ("deg(some_metric)", "prom_native_histogram_drop_float"), - ("sgn(some_metric)", "prom_native_histogram_drop_float"), - // Instant helper functions routed to native histogram UDFs. - ( - "histogram_count(some_metric)", - "prom_native_histogram_count", - ), - ("histogram_sum(some_metric)", "prom_native_histogram_sum"), - ("histogram_avg(some_metric)", "prom_native_histogram_avg"), - ( - "histogram_stddev(some_metric)", - "prom_native_histogram_stddev", - ), - ( - "histogram_stdvar(some_metric)", - "prom_native_histogram_stdvar", - ), - ( - "histogram_fraction(-2 + 1, 2 / 2, some_metric)", - "prom_native_histogram_fraction", - ), - ]; - - for (query, expected_udf) in cases { - let plan = native_histogram_plan(query).await; - assert!(plan.contains(expected_udf), "{query}\n{plan}"); - if query.starts_with("histogram_fraction") { - assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); - } - } - } - - #[tokio::test] - async fn mixed_native_histogram_ranges_use_coordinated_udfs() { - let dual_output = [ - "increase(some_metric[5m])", - "rate(some_metric[5m])", - "delta(some_metric[5m])", - "idelta(some_metric[5m])", - "irate(some_metric[5m])", - "avg_over_time(some_metric[5m])", - "sum_over_time(some_metric[5m])", - "last_over_time(some_metric[5m])", - ]; - let float_output = [ - "resets(some_metric[5m])", - "changes(some_metric[5m])", - "deriv(some_metric[5m])", - "min_over_time(some_metric[5m])", - "max_over_time(some_metric[5m])", - "count_over_time(some_metric[5m])", - "absent_over_time(some_metric[5m])", - "present_over_time(some_metric[5m])", - "stddev_over_time(some_metric[5m])", - "stdvar_over_time(some_metric[5m])", - "quantile_over_time(0.9, some_metric[5m])", - "predict_linear(some_metric[5m], 60)", - "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", - ]; - - for query in dual_output.iter().chain(float_output.iter()) { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); - assert_eq!( - plan.contains("prom_mixed_range_histogram"), - dual_output.contains(query), - "{query}\n{plan}" - ); - } - - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), - &build_query_engine_state(), - ) - .await - .unwrap() - .display_indent_schema() - .to_string(); - let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; - assert_eq!(plan, expected); - } - - #[tokio::test] - async fn mixed_native_histogram_rate_executes_real_ranges() { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - ), - Field::new(greptime_value(), ArrowDataType::Float64, true), - Field::new( - greptime_native_histogram(), - native_histogram_value_type().as_arrow_type(), - true, - ), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), - Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), - build_histogram_array(&[None, Some(direct_or_histogram()), None]), - ], - ) - .unwrap(); - let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); - let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) - .unwrap() - .build() - .unwrap(); - let collector = PromqlAnnotationCollector::default(); - let mut planner = PromPlanner { - table_provider: build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await, - ctx: PromPlannerContext { - start: 3000, - end: 3000, - interval: 1000, - range: Some(3000), - time_index_column: Some("timestamp".to_string()), - field_columns: vec![ - greptime_native_histogram().to_string(), - greptime_value().to_string(), - ], - ..Default::default() - }, - promql_annotations: Some(collector.clone()), - }; - let input = LogicalPlan::Extension(Extension { - node: Arc::new( - RangeManipulate::new( - 3000, - 3000, - 1000, - 0, - 3000, - "timestamp".to_string(), - planner.ctx.field_columns.clone(), - input, - ) - .unwrap(), - ), - }); - let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { - unreachable!() - }; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = Vec::new(); - collector.append_to(&mut warnings, &mut Vec::new()); - assert!( - warnings - .iter() - .any(|warning| warning.contains("mix of float and native histogram")) - ); - } - - #[tokio::test] - async fn native_histogram_mixed_field_table_behaves() { - // Exercise function planning after float and histogram samples have already been - // represented as alternative nullable fields. Histogram functions must select the - // histogram field without adding a NULL float field that would reject every row. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("histogram_count(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("prom_native_histogram_count"), - "{plan_str}" - ); - assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); - assert!( - plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), - "{plan_str}" - ); - - // Value sorting keeps the float column and never sorts by the histogram column. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("sort(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!( - plan_str.contains("greptime_value ASC NULLS FIRST"), - "{plan_str}" - ); - assert!( - !plan_str.contains("greptime_native_histogram ASC"), - "{plan_str}" - ); - - // scalar() ignores histogram samples and evaluates only the float field. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt("scalar(some_metric)"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); - assert!( - plan_str.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - - // Functions that preserve both alternative fields keep rows with either sample type. - let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), - &build_query_engine_state(), - ) - .await - .unwrap(); - let plan_str = plan.display_indent_schema().to_string(); - let filter = plan_str.lines().next().unwrap(); - assert!( - filter.starts_with("Filter: ") - && filter.contains("greptime_native_histogram IS NOT NULL") - && filter.contains(" OR ") - && filter.contains("greptime_value IS NOT NULL"), - "{plan_str}" - ); - } - - #[tokio::test] - async fn less_filter_on_value() { - let query = "some_metric < 1.2345"; - let expected = String::from( - "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn count_over_time() { - let query = "count_over_time(some_metric[5m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - /// The outer `PromRangeManipulate` from a subquery must be preceded by - /// `Sort` + `PromSeriesDivide`. - #[tokio::test] - async fn count_over_time_subquery() { - let query = "count_over_time(some_metric[10m:1m])"; - let expected = String::from( - "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn test_hash_join() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_sum".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["uri", "kubernetes_namespace", "kubernetes_pod_name"], - ) - .await; - // Should be ok - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ - \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ - \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ - \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ - \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ - \n SubqueryAlias: http_server_requests_seconds_sum\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_sum\ - \n SubqueryAlias: http_server_requests_seconds_count\ - \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ - \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ - \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ - \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ - \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ - \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ - \n TableScan: http_server_requests_seconds_count"; - assert_eq!(plan.to_string(), expected); - } - - #[tokio::test] - async fn test_nested_histogram_quantile() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "greptime_servers_grpc_requests_elapsed_bucket".to_string(), - )], - &["pod", "le", "path", "code", "container"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_histogram_quantile_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Arithmetic applied to a histogram_quantile() result. Regression for #8144: - // HistogramFold used to drop the input column qualifiers, so the binary-op - // projection failed to resolve the qualified tag column. - let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "http_request_duration_seconds_bucket".to_string(), - )], - &["pod", "le"], - ) - .await; - // Should plan without a "No field named ..." error. - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_and_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, - ]; - - for case in cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_used_bytes".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "kubelet_volume_stats_capacity_bytes".to_string(), - ), - ], - &["namespace", "persistentvolumeclaim"], - ) - .await; - // Should be ok - let _ = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_nested_binary_op() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - - ( - sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) - or - vector(0) - )"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "nginx_ingress_controller_requests".to_string(), - )], - &["namespace", "job"], - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_parse_or_operator() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#" - sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / - (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) - or - 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / - sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; - - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 - or - sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / - (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + - sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or - (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "background_waitevent_cnt".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "foreground_waitevent_cnt".to_string(), - ), - ], - &["tenant_name", "cluster_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; - let table_provider = build_test_table_provider_with_fields( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_cpu_load_average_10s".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "container_spec_cpu_quota".to_string(), - ), - ], - &["cluster_name", "host_name"], - ) - .await; - eval_stmt.expr = parser::parse(case).unwrap(); - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn value_matcher() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let cases = [ - // single equal matcher - ( - r#"some_metric{__field__="field_1"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two equal matchers - ( - r#"some_metric{__field__="field_1", __field__="field_0"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single not_eq matcher - ( - r#"some_metric{__field__!="field_1"}"#, - vec![ - "some_metric.field_0", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // two not_eq matchers - ( - r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (no conflict) - ( - r#"some_metric{__field__="field_1", __field__!="field_0"}"#, - vec![ - "some_metric.field_1", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // equal and not_eq matchers (conflict) - ( - r#"some_metric{__field__="field_2", __field__!="field_2"}"#, - vec![ - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex eq matcher - ( - r#"some_metric{__field__=~"field_1|field_2"}"#, - vec![ - "some_metric.field_1", - "some_metric.field_2", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - // single regex not_eq matcher - ( - r#"some_metric{__field__!~"field_1|field_2"}"#, - vec![ - "some_metric.field_0", - "some_metric.tag_0", - "some_metric.tag_1", - "some_metric.tag_2", - "some_metric.timestamp", - ], - ), - ]; - - for case in cases { - let prom_expr = parser::parse(case.0).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let mut fields = plan.schema().field_names(); - let mut expected = case.1.into_iter().map(String::from).collect::>(); - fields.sort(); - expected.sort(); - assert_eq!(fields, expected, "case: {:?}", case.0); - } - - let bad_cases = [ - r#"some_metric{__field__="nonexistent"}"#, - r#"some_metric{__field__!="nonexistent"}"#, - ]; - - for case in bad_cases { - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "case: {:?}", case); - } - } - - #[tokio::test] - async fn custom_schema() { - let query = "some_alt_metric{__schema__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__database__=\"greptime_private\"}"; - let expected = String::from( - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - - let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; - let expected = String::from( - "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ - \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", - ); - - indie_query_plan_compare(query, expected).await; - } - - #[tokio::test] - async fn only_equals_is_supported_for_special_matcher() { - let queries = &[ - "some_alt_metric{__schema__!=\"greptime_private\"}", - "some_alt_metric{__schema__=~\"lalala\"}", - "some_alt_metric{__database__!=\"greptime_private\"}", - "some_alt_metric{__database__=~\"lalala\"}", - ]; - - for query in queries { - let prom_expr = parser::parse(query).unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = build_test_table_provider( - &[ - (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), - ( - "greptime_private".to_string(), - "some_alt_metric".to_string(), - ), - ], - 1, - 1, - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - assert!(plan.is_err(), "query: {:?}", query); - } - } - - #[tokio::test] - async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), - promql_annotations: None, - }; - planner.ctx.time_index_column = Some("timestamp".to_string()); - planner.ctx.start = 1_000; - planner.ctx.lookback_delta = 0; - let schema = Arc::new( - DFSchema::try_from(ArrowSchema::new(vec![Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - false, - )])) - .unwrap(), - ); - for (end, interval, windows) in [ - (1_000, 1_000, 1), - (2_000, 1_000, 1), - (7_201_000, 7_200_000, 2), - ] { - planner.ctx.end = end; - planner.ctx.interval = interval; - let filter = planner - .build_time_index_filter(0, &schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!(filter.matches(">=").count(), windows, "{filter}"); - assert!( - filter.contains("TimestampNanosecond(1000000000, None)"), - "{filter}" - ); - } - planner.ctx.end = i64::MAX; - let filter = planner - .build_time_index_filter(0, &schema) - .unwrap() - .unwrap() - .to_string(); - assert!( - filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"), - "{filter}" - ); - - // A lookback subtraction can underflow milliseconds while the upper bound remains - // representable. Keep that upper bound so LastRow cannot select a future sample. - let ms_schema = Arc::new( - DFSchema::try_from(ArrowSchema::new(vec![Field::new( - "timestamp", - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )])) - .unwrap(), - ); - planner.ctx.start = i64::MIN + 100; - planner.ctx.end = planner.ctx.start; - planner.ctx.lookback_delta = 200; - let filter = planner - .build_time_index_filter(0, &ms_schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!( - filter, - format!( - "timestamp <= TimestampMillisecond({}, None)", - i64::MIN + 100 - ) - ); - - // The lower bound can also overflow while converting milliseconds to native nanoseconds. - // Its representable upper bound still has to reach the scan. - planner.ctx.start = 0; - planner.ctx.end = 0; - planner.ctx.lookback_delta = 300_000; - let filter = planner - .build_time_index_filter(9_223_372_036_854, &schema) - .unwrap() - .unwrap() - .to_string(); - assert_eq!( - filter, - "timestamp <= TimestampNanosecond(-9223372036854000000, None)" - ); - } - - #[tokio::test] - async fn test_non_ms_precision() { - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new( - "tag".to_string(), - ConcreteDataType::string_datatype(), - false, - ), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_nanosecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new( - "field".to_string(), - ConcreteDataType::float64_datatype(), - true, - ), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name("metrics".to_string()) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "metrics".to_string(), - table_id: 1024, - table, - }) - .is_ok() - ); - - let plan = PromPlanner::stmt_to_plan( - DfTableSourceProvider::new( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("metrics{tag = \"1\"}").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert_eq!( - plan.display_indent_schema().to_string(), - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" - ); - let plan = PromPlanner::stmt_to_plan( - DfTableSourceProvider::new( - catalog_list.clone(), - false, - QueryContext::arc(), - DummyDecoder::arc(), - true, - ), - &EvalStmt { - expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert_eq!( - plan.display_indent_schema().to_string(), - "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" - ); - } - - #[tokio::test] - async fn test_nonexistent_label() { - // template - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let case = r#"some_metric{nonexistent="hi"}"#; - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 3, - 3, - ) - .await; - // Should be ok - let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_label_join() { - let prom_expr = parser::parse( - "label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn test_label_replace() { - let prom_expr = parser::parse( - "label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")", - ) - .unwrap(); - let eval_stmt = EvalStmt { - expr: prom_expr, - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - let table_provider = - build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - - let expected = r#" -Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] - TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; - - let ret = plan.display_indent_schema().to_string(); - assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); - } - - #[tokio::test] - async fn label_replace_aggregation_queries_plan_successfully() { - let aggregate = - r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; - let queries = [ - aggregate.to_string(), - format!("{aggregate} <= 10"), - format!("{aggregate} * 0.8"), - format!("0.8 * {aggregate}"), - format!("{aggregate} <= {aggregate} * 0.8"), - ]; - let state = build_query_engine_state(); - let mut failures = Vec::new(); - - for query in queries { - let table_provider = build_test_table_provider( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - ) - .await; - if let Err(error) = - PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await - { - failures.push(format!("{query}: {error:?}")); - } - } - - assert!(failures.is_empty(), "{}", failures.join("\n")); - } - - #[tokio::test] - async fn test_matchers_to_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = - r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - )], - 3, - 3, - ) - .await; - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_topk_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ - \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_count_values_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_value_alias() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = r#" -Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; - assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); - } - - #[tokio::test] - async fn test_quantile_expr() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "http_server_requests_seconds_count".to_string(), - ), - ], - &["ip"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; - - assert_eq!(plan.display_indent_schema().to_string(), expected); - } - - #[tokio::test] - async fn test_or_not_exists_table_label() { - let state = build_query_engine_state(); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], - &["job"], - ) - .await; - let raw = PromPlanner::stmt_to_plan( - provider, - &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), - &state, - ) - .await - .unwrap(); - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_0@") - ); - let (optimized, batches) = execute(raw, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - } - - #[tokio::test] - async fn test_histogram_quantile_missing_le_column() { - let mut eval_stmt = EvalStmt { - expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), - start: UNIX_EPOCH, - end: UNIX_EPOCH - .checked_add(Duration::from_secs(100_000)) - .unwrap(), - interval: Duration::from_secs(5), - lookback_delta: Duration::from_secs(1), - }; - - // Test case: histogram_quantile with a table that doesn't have 'le' column - let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; - - let prom_expr = parser::parse(case).unwrap(); - eval_stmt.expr = prom_expr; - - // Create a table provider with a table that doesn't have 'le' column - let table_provider = build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "non_existent_histogram_bucket".to_string(), - )], - &["pod", "instance"], // Note: no 'le' column - ) - .await; - - // Should return empty result instead of error - let result = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await; - - // This should succeed now (returning empty result) instead of failing with "Cannot find column le" - assert!( - result.is_ok(), - "Expected successful plan creation with empty result, but got error: {:?}", - result.err() - ); - - // Verify that the result is an EmptyRelation - let plan = result.unwrap(); - match plan { - LogicalPlan::EmptyRelation(_) => { - // This is what we expect - } - _ => panic!("Expected EmptyRelation, but got: {:?}", plan), - } - } - - #[tokio::test] - async fn test_direct_or_normalizes_missing_match_labels() { - type Case<'a> = ( - Option>, - Option>, - i64, - i64, - &'a [(f64, Option<&'a str>)], - ); - - let modifier = or_modifier("lhs or on(k) rhs"); - #[rustfmt::skip] - let cases: &[Case<'_>] = &[ - (None, None, 1, 1, &[(1.0, None)]), - (None, Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), - (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), - (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), - (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), - (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), - ]; - for &(left, right, left_ts, right_ts, expected) in cases { - let (optimized, batches) = run( - &matrix_source("lhs", left, left_ts, 1.0), - &matrix_source("rhs", right, right_ts, 2.0), - matrix_context("lhs", left), - matrix_context("rhs", right), - &modifier, - ) - .await; - assert_no_internal_or_keys(optimized.schema()); - assert_eq!( - rows(&batches), - expected - .iter() - .map(|(value, label)| (*value, label.map(str::to_string))) - .collect::>() - ); - } - } - - #[tokio::test] - async fn test_direct_or_match_modifiers() { - for (modifier, left, right, expected) in [ - (None, "left", "right", 2), - (or_modifier("lhs or on(k) rhs"), "same", "same", 1), - (or_modifier("lhs or on() rhs"), "left", "right", 1), - (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), - ] { - let (_, batches) = run( - &matrix_source("lhs", Some(Some(left)), 1, 1.0), - &matrix_source("rhs", Some(Some(right)), 1, 2.0), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &modifier, - ) - .await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected - ); - } - } - - #[tokio::test] - async fn test_direct_or_nested_projection_uses_left_context() { - let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); - let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert!(raw.schema().iter().any(|(qualifier, field)| { - qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" - })); - let nested = LogicalPlanBuilder::from(raw) - .project(vec![ - DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(DfExpr::Column(Column::new( - Some(TableReference::bare("lhs")), - "v", - ))), - op: Operator::Plus, - right: Box::new(lit(1.0)), - }) - .alias("v_plus"), - ]) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(nested, &build_query_engine_state()).await; - assert_eq!(values(&batches, "v_plus"), vec![2.0]); - } - - #[tokio::test] - async fn test_direct_or_skips_user_internal_key_name() { - const USER_TAG: &str = "__promql_or_match_0"; - let left = tagged_source( - "lhs", - false, - (USER_TAG, Some("left")), - DirectOrValue::Float64(1.0), - ); - let right = tagged_source( - "rhs", - false, - (USER_TAG, Some("right")), - DirectOrValue::Float64(2.0), - ); - let raw = plan_direct_or( - scan(&left), - scan(&right), - direct_or_context("lhs", &["job", USER_TAG], "v"), - direct_or_context("rhs", &["job", USER_TAG], "v"), - &or_modifier("lhs or on(missing_label) rhs"), - ) - .await; - assert!( - raw.display_indent_schema() - .to_string() - .contains("__promql_or_match_1@") - ); - let (_, batches) = execute(raw, &build_query_engine_state()).await; - assert!( - batches - .iter() - .all(|batch| batch.column_by_name(USER_TAG).is_some()) - ); - } - - #[tokio::test] - async fn test_direct_or_substrait_round_trip_with_normalized_key() { - let state = build_query_engine_state(); - let ctx = SessionContext::new_with_state(state.session_state()); - let catalog = Arc::new(MemoryCatalogProvider::new()); - catalog - .register_schema("public", Arc::new(MemorySchemaProvider::new())) - .unwrap(); - ctx.register_catalog("datafusion", catalog); - let left = matrix_source("lhs", Some(Some("")), 1, 1.0); - let right = matrix_source("rhs", None, 1, 2.0); - ctx.register_table( - TableReference::full("datafusion", "public", "lhs"), - table(&left), - ) - .unwrap(); - ctx.register_table( - TableReference::full("datafusion", "public", "rhs"), - table(&right), - ) - .unwrap(); - let raw = plan_direct_or( - ctx.table("datafusion.public.lhs") - .await - .unwrap() - .into_unoptimized_plan(), - ctx.table("datafusion.public.rhs") - .await - .unwrap() - .into_unoptimized_plan(), - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - let decoded = DFLogicalSubstraitConvertor - .decode( - DFLogicalSubstraitConvertor - .encode(&raw, DefaultSerializer) - .unwrap(), - ctx.state(), - ) - .await - .unwrap(); - let (optimized, batches) = execute(decoded, &state).await; - assert_no_internal_or_keys(optimized.schema()); - assert!(batches.iter().all(|batch| { - batch - .schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_or_match_")) - })); - assert_eq!(values(&batches, "v"), vec![1.0]); - } - - #[tokio::test] - async fn test_direct_or_numeric_value_types() { - let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); - let right = tagged_source( - "rhs", - false, - ("k", Some("rhs")), - DirectOrValue::Float64(0.5), - ); - let (optimized, batches) = run( - &left, - &right, - direct_or_context("lhs", &["job", "k"], "v"), - direct_or_context("rhs", &["job", "k"], "v"), - &or_modifier("lhs or on(k) rhs"), - ) - .await; - assert_eq!( - optimized - .schema() - .field_with_name(None, "v") - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!(values(&batches, "v"), vec![0.5]); - let provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider: provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left_context = direct_or_context("lhs", &["job"], "v"); - let right_context = direct_or_context("rhs", &["job"], "v"); - let error = planner - .or_operator( - scan(&job_source("lhs", DirectOrValue::Utf8("x"))), - scan(&job_source("rhs", DirectOrValue::Float64(1.0))), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on() rhs"), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("OR value fields have incompatible types") - ); - } - - #[tokio::test] - async fn test_or_with_histogram_quantile_missing_le_column() { - let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - - #[tokio::test] - async fn test_or_with_right_empty_histogram_restores_left_context() { - let eval_stmt = build_eval_stmt( - r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, - ); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_or_with_both_empty_histograms() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, - ); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - match plan { - LogicalPlan::EmptyRelation(relation) => { - assert!(!relation.produce_one_row); - assert!(!relation.schema.fields().is_empty()); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64) - ); - assert!( - relation - .schema - .fields() - .iter() - .any(|field| field.name() == "pod") - ); - assert!( - !relation - .schema - .fields() - .iter() - .any(|field| field.name() == "instance") - ); - } - _ => panic!("Expected EmptyRelation, but got: {plan:?}"), - } - } - - #[tokio::test] - async fn test_nested_or_with_both_empty_histograms() { - for case in [ - r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, - r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_test_table_provider_with_fields( - &[ - ( - DEFAULT_SCHEMA_NAME.to_string(), - "left_histogram_bucket".to_string(), - ), - ( - DEFAULT_SCHEMA_NAME.to_string(), - "right_histogram_bucket".to_string(), - ), - ], - &["pod", "instance"], - ) - .await; - - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - } - } - - #[tokio::test] - async fn test_or_with_empty_histogram_modifiers() { - for case in [ - r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, - r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, - ] { - let eval_stmt = build_eval_stmt(case); - let table_provider = build_missing_le_or_normal_metric_table_provider().await; - - let plan = - PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) - .await - .unwrap(); - assert_normal_metric_schema(&plan); - } - } - - #[tokio::test] - async fn test_unless_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, - ); - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &state, - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; - assert!(contains_histogram_fold(&optimized), "{optimized:?}"); - let batches = - datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) - .await - .unwrap(); - assert!(batches.iter().all(|batch| batch.num_rows() == 0)); - } - - #[tokio::test] - async fn test_and_preserves_left_context_for_histogram() { - let eval_stmt = build_eval_stmt( - r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, - ); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(contains_histogram_fold(&plan), "{plan:?}"); - } - - #[tokio::test] - async fn test_and_preserves_left_context_when_le_is_missing() { - let eval_stmt = - build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); - let plan = PromPlanner::stmt_to_plan( - build_set_op_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); - assert!(!plan.schema().fields().is_empty()); - assert!(!contains_histogram_fold(&plan), "{plan:?}"); - } - - async fn build_matching_filter_plan(query: &str) -> String { - let table_provider = build_test_table_provider_with_distinct_tags(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - /// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value - /// field that is neither a primary key nor a value column of the metric, so - /// `count by(status) (...)` still reports it among the aggregation's tag columns. - async fn build_test_table_provider_with_string_field( - table_tags: &[(&str, &[&str])], - ) -> DfTableSourceProvider { - let catalog_list = MemoryCatalogManager::with_default_setup(); - for (table_name, tags) in table_tags { - let mut columns = tags - .iter() - .map(|tag| { - ColumnSchema::new( - (*tag).to_string(), - ConcreteDataType::string_datatype(), - false, - ) - }) - .collect::>(); - columns.push( - ColumnSchema::new( - greptime_timestamp().to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ); - columns.push(ColumnSchema::new( - greptime_value().to_string(), - ConcreteDataType::float64_datatype(), - true, - )); - columns.push(ColumnSchema::new( - "status".to_string(), - ConcreteDataType::string_datatype(), - true, - )); - let table_meta = TableMetaBuilder::empty() - .schema(Arc::new(Schema::new(columns))) - .primary_key_indices((0..tags.len()).collect()) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = TableInfoBuilder::default() - .name((*table_name).to_string()) - .meta(table_meta) - .build() - .unwrap(); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: (*table_name).to_string(), - table_id: 1024, - table: EmptyTable::from_table_info(&table_info), - }) - .is_ok() - ); - } - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - async fn build_matching_filter_plan_with_string_field(query: &str) -> String { - let table_provider = build_test_table_provider_with_string_field(&[ - ("metric_a", &["host", "device"]), - ("metric_b", &["host", "device"]), - ]) - .await; - let plan = PromPlanner::stmt_to_plan( - table_provider, - &build_eval_stmt(query), - &build_query_engine_state(), - ) - .await - .unwrap(); - plan.display_indent().to_string() - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_both_operands() { - for query in [ - r#"metric_a / metric_b{host="foo"}"#, - r#"metric_a / on(host, device) metric_b{host="foo"}"#, - r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, - r#"metric_a / ignoring(device) metric_b{host="foo"}"#, - r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, - r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 2, - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { - for query in [ - r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, - r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, - r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 2, - "{query}\n{plan}" - ); - } - // A global ranking one-side must see every host, so the matcher stays put. - let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; - let plan = build_matching_filter_plan(query).await; - assert_eq!( - plan.matches(r#"host = Utf8("foo")"#).count(), - 1, - "{query}\n{plan}" - ); - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_selecting_aggregations() { - // `topk` ranks its input, so filtering before it changes the candidate set. - let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_value_field_matcher_stays_on_its_own_operand() { - let value = greptime_value(); - for query in [ - format!(r#"metric_a / metric_b{{{value}="2"}}"#), - format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), - ] { - let plan = build_matching_filter_plan(&query).await; - assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn binary_value_field_matcher_is_not_copied_across_aggregations() { - // `status` varies between the samples of one series, so filtering the other operand by it - // would drop the newest sample before sample selection (#9242). - for query in [ - r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, - r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, - r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, - ] { - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!( - plan.matches(r#"Utf8("ready")"#).count(), - 1, - "{query}\n{plan}" - ); - } - } - - #[tokio::test] - async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { - // `count`, not `sum`: the string value field is not summable. - let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; - let plan = build_matching_filter_plan_with_string_field(query).await; - assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); - } - - #[tokio::test] - async fn binary_matching_label_filter_skips_unproven_expressions() { - for query in [ - r#"metric_a > on(host, device) metric_b{host="foo"}"#, - r#"metric_a / on(host) group_left metric_b{host="foo"}"#, - r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, - r#"metric_a / on(device) metric_b{host="foo"}"#, - ] { - let plan = build_matching_filter_plan(query).await; - assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); - } - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_output() { - let case = r#"(normal_metric or other_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let state = build_query_engine_state(); - let plan = - PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; - } - - #[tokio::test] - async fn test_or_context_uses_left_qualified_empty_histogram_output() { - let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; - let eval_stmt = build_eval_stmt(case); - let plan = PromPlanner::stmt_to_plan( - build_or_context_table_provider().await, - &eval_stmt, - &build_query_engine_state(), - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .any(|field| field.data_type() == &ArrowDataType::Float64), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_direct_or_preserves_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (planner, plan) = mixed_direct_or(histogram_on_left).await; - - let float_field = &planner.ctx.field_columns[0]; - let histogram_field = &planner.ctx.field_columns[1]; - assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); - assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); - assert_eq!( - plan.schema() - .field_with_name(None, float_field) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - plan.schema() - .field_with_name(None, histogram_field) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - - let (optimized, batches) = execute(plan, &build_query_engine_state()).await; - assert_no_internal_or_keys(optimized.schema()); - let mut sample_kinds = batches - .iter() - .flat_map(|batch| { - let values = batch.column_by_name(float_field).unwrap(); - let histograms = batch.column_by_name(histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (values.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - sample_kinds.sort_unstable(); - assert_eq!(sample_kinds, vec![(false, true), (true, false)]); - } - } - - #[tokio::test] - async fn malformed_classic_bucket_does_not_drop_native_histogram() { - let state = build_query_engine_state(); - let collector = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), - &state, - Some(collector.clone()), - ) - .await - .unwrap(); - let value_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &value_field), vec![0.0]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert!(infos.is_empty()); - } - - #[tokio::test] - async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), - &state, - ) - .await - .unwrap(); - let plan_text = plan.display_indent_schema().to_string(); - assert!( - plan_text.contains("prom_native_histogram_mul_scalar"), - "{plan_text}" - ); - assert!( - plan_text.contains("prom_native_histogram_scalar_mul"), - "{plan_text}" - ); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - assert!(values(&batches, &float_field).is_empty()); - let mut sums = histograms(&batches, &histogram_field) - .into_iter() - .map(|histogram| histogram.sum) - .collect::>(); - sums.sort_by(f64::total_cmp); - assert_eq!(sums, vec![2.0, 3.0]); - } - - #[tokio::test] - async fn test_mixed_binary_operator_reports_only_dropped_samples() { - for (query, expected_rows, expected_infos) in [ - ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), - ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), - ("(lf or on(tag) lh) % on(tag) lh", 0, 1), - ] { - let state = build_query_engine_state(); - let annotations = PromqlAnnotationCollector::default(); - let plan = PromPlanner::stmt_to_plan_with_annotations( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - Some(annotations.clone()), - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - expected_rows, - "{query}" - ); - let mut warnings = vec![]; - let mut infos = vec![]; - annotations.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty(), "{query}: {warnings:?}"); - assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); - } - } - - #[tokio::test] - async fn test_histogram_only_min_drops_empty_aggregate_group() { - // `min` over native-histogram-only input drops every sample in the group, so the - // NULL-valued aggregate row must be filtered out. Otherwise an outer expression - // like `group()` resurrects the group Prometheus considers unseen. - let state = build_query_engine_state(); - for query in ["min(lh)", "group(min(lh))"] { - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - let (_, batches) = execute(plan, &state).await; - assert_eq!( - batches.iter().map(RecordBatch::num_rows).sum::(), - 0, - "{query}" - ); - } - } - - #[tokio::test] - async fn test_mixed_min_drops_histogram_only_group() { - // With alternative float/histogram fields, `min by (tag)` keeps float-only groups - // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of - // emitting a NULL-valued row for them. - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - assert_eq!(values(&batches, &float_field), vec![2.0]); - } - - #[tokio::test] - async fn test_mixed_or_can_feed_another_or() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); - let mut float_values = values(&batches, &float_field); - float_values.sort_by(f64::total_cmp); - assert_eq!(float_values, vec![2.0, 7.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_mixed_fields_align_with_single_float_vector() { - let (planner, mixed) = mixed_direct_or(false).await; - let scale = tagged_source( - "scale", - false, - ("k", Some("float")), - DirectOrValue::Float64(2.0), - ); - let scale = scan(&scale); - let scale_fields = vec!["v".to_string()]; - let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { - unreachable!() - }; - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - mixed.schema(), - scale.schema(), - &planner.ctx.field_columns, - &scale_fields, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(_, right)| *right == &scale_fields[0]) - ); - - let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( - scale.schema(), - mixed.schema(), - &scale_fields, - &planner.ctx.field_columns, - binary.op, - false, - false, - ); - assert!(invalid_pairs.is_empty()); - assert_eq!( - groups - .iter() - .map(|(output, _)| output.clone()) - .collect::>(), - planner.ctx.field_columns - ); - assert_eq!(groups.len(), 2); - assert!( - groups - .iter() - .flat_map(|(_, pairs)| pairs) - .all(|(left, _)| *left == &scale_fields[0]) - ); - } - - #[tokio::test] - async fn test_non_bool_comparison_filters_mixed_sample_lanes() { - let (planner, input) = mixed_direct_or(false).await; - let input_schema = input.schema().clone(); - let plan = planner - .filter_on_field_column(input, |field| { - if PromPlanner::field_column_is_native_histogram(&input_schema, field) { - Ok(lit(false)) - } else { - Ok(col(field).gt(lit(0.0))) - } - }) - .unwrap(); - let float_field = planner.ctx.field_columns[0].clone(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_field), vec![1.25]); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - } - - #[tokio::test] - async fn test_mixed_left_and_unless_preserve_sample_lanes() { - for (expression, expected_sample_kind) in [ - ("lhs and on(k) mask", (false, true)), - ("lhs unless on(k) mask", (true, false)), - ] { - let (mut planner, left) = mixed_direct_or(false).await; - let left_context = planner.ctx.clone(); - let float_field = left_context.field_columns[0].clone(); - let histogram_field = left_context.field_columns[1].clone(); - let mask = tagged_source( - "mask", - false, - ("k", Some("histogram")), - DirectOrValue::Float64(1.0), - ); - let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { - unreachable!() - }; - let plan = planner - .set_op_on_non_field_columns( - left, - scan(&mask), - left_context, - direct_or_context("mask", &["job", "k"], "v"), - binary.op, - &binary.modifier, - ) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let sample_kinds = batches - .iter() - .flat_map(|batch| { - let floats = batch.column_by_name(&float_field).unwrap(); - let histograms = batch.column_by_name(&histogram_field).unwrap(); - (0..batch.num_rows()) - .map(|row| (floats.is_valid(row), histograms.is_valid(row))) - }) - .collect::>(); - assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); - } - } - - #[tokio::test] - async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &build_eval_stmt("some_metric * scalar(vector(2))"), - &build_query_engine_state(), - ) - .await - .unwrap(); - let schema = plan.schema(); - assert_eq!( - schema - .field_with_unqualified_name(greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Float64 - ); - assert_eq!( - schema - .field_with_unqualified_name(greptime_native_histogram()) - .unwrap() - .data_type(), - &native_histogram_value_type().as_arrow_type() - ); - assert!( - plan.display_indent_schema() - .to_string() - .contains("prom_native_histogram_mul_scalar"), - "{plan:?}" - ); - } - - #[tokio::test] - async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - operator_table_provider(), - &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), - &state, - ) - .await - .unwrap(); - let float_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &ArrowDataType::Float64) - .unwrap() - .name() - .clone(); - let histogram_field = plan - .schema() - .fields() - .iter() - .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) - .unwrap() - .name() - .clone(); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(values(&batches, &float_field), vec![0.0]); - assert_eq!(histograms(&batches, &histogram_field).len(), 1); - } - - #[tokio::test] - async fn test_unary_negates_mixed_float_and_native_histogram_samples() { - for histogram_on_left in [false, true] { - let (mut planner, input) = mixed_direct_or(histogram_on_left).await; - let plan = planner.negate_field_columns(input).unwrap(); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let float_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) - .unwrap(); - let histogram_field = planner - .ctx - .field_columns - .iter() - .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, float_field), vec![-1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(histogram_field) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, -1.0); - assert_eq!(histogram.sum, -1.0); - assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); - } - } - - #[tokio::test] - async fn test_native_histogram_sum_and_avg_execute_real_batches() { - for op_name in ["sum", "avg"] { - for incompatible in [false, true] { - let mut second = direct_or_histogram(); - if incompatible { - second.schema = CUSTOM_BUCKETS_SCHEMA; - second.custom_values = vec![1.0]; - } - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = - mixed_aggregate_input(vec![direct_or_histogram(), second]).await; - planner.promql_annotations = Some(collector.clone()); - let histogram_column = planner.ctx.field_columns[1].clone(); - planner.ctx.field_columns = vec![histogram_column.clone()]; - let input = LogicalPlanBuilder::from(input) - .project([col("ts"), col(&histogram_column)]) - .unwrap() - .build() - .unwrap(); - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse(&format!("{op_name}(mixed)")).unwrap() - else { - unreachable!() - }; - let (aggregate_exprs, _) = - planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter(planner.create_empty_values_filter_expr(false).unwrap()) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(infos.is_empty()); - if incompatible { - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - assert!(warnings.iter().any(|warning| { - warning - == &format!( - "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" - ) - })); - } else { - let histograms = histograms(&batches, &histogram_column); - assert_eq!(histograms.len(), 1); - let expected = if op_name == "sum" { 2.0 } else { 1.0 }; - assert_eq!(histograms[0].count, expected); - assert_eq!(histograms[0].sum, expected); - assert!(warnings.is_empty()); - } - } - } - } - - #[tokio::test] - async fn test_canonical_mixed_count_group_and_count_values_execute() { - let state = build_query_engine_state(); - for (query, expected) in [ - ("count(some_metric)", vec![2.0]), - ("group(some_metric)", vec![1.0]), - (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), - ] { - let plan = PromPlanner::stmt_to_plan( - build_test_mixed_native_histogram_table_provider("some_metric").await, - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap(); - assert!( - plan.schema() - .fields() - .iter() - .all(|field| !field.name().starts_with("__promql_sample_count")), - "{query}: {plan:?}" - ); - let value_fields = plan - .schema() - .fields() - .iter() - .filter(|field| { - matches!( - field.data_type(), - ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 - ) || field.data_type() == &native_histogram_value_type().as_arrow_type() - }) - .collect::>(); - assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); - assert_ne!( - value_fields[0].data_type(), - &native_histogram_value_type().as_arrow_type(), - "{query}: {plan:?}" - ); - let value_column = value_fields[0].name().clone(); - - let (_, batches) = execute(plan, &state).await; - let mut actual = numeric_values(&batches, &value_column); - actual.sort_by(f64::total_cmp); - assert_eq!(actual, expected, "{query}"); - - if query.starts_with("count_values") { - let mut sample_labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name("sample") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .map(str::to_string) - }) - .collect::>(); - sample_labels.sort(); - let mut expected_labels = - vec!["2".to_string(), direct_or_histogram().promql_string()]; - expected_labels.sort(); - assert_eq!(sample_labels, expected_labels); - } - } - } - - #[tokio::test] - async fn test_mixed_or_sum_aggregates_each_sample_type() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts"), col("k")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([ - col(&float_column), - col(&histogram_column), - col("ts"), - col("k"), - ]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let histogram = batches - .iter() - .find_map(|batch| { - let values = batch - .column_by_name(&histogram_column)? - .as_any() - .downcast_ref::()?; - (0..values.len()).find_map(|row| { - common_query::native_histogram::read_histogram(values, row).unwrap() - }) - }) - .unwrap(); - assert_eq!(histogram.count, 1.0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_direct_or(false).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert_eq!( - warnings, - vec![ - "sum: dropped aggregation result containing both float and native histogram samples" - ] - ); - } - - #[tokio::test] - async fn test_mixed_or_sum_drops_incompatible_mixed_group() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("sum(lhs)").unwrap() - else { - unreachable!() - }; - let mut custom = direct_or_histogram(); - custom.schema = CUSTOM_BUCKETS_SCHEMA; - custom.custom_values = vec![1.0]; - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col(&histogram_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.iter().any(|warning| { - warning - == "sum: dropped aggregation result containing both float and native histogram samples" - })); - } - - #[tokio::test] - async fn test_mixed_or_min_records_only_present_histograms() { - let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = - parser::parse("min(lhs)").unwrap() - else { - unreachable!() - }; - let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; - - for (histograms, expected_infos) in [ - (vec![], vec![]), - (vec![direct_or_histogram()], vec![expected_info]), - ] { - let collector = PromqlAnnotationCollector::default(); - let (mut planner, input) = mixed_aggregate_input(histograms).await; - planner.promql_annotations = Some(collector.clone()); - let float_column = planner.ctx.field_columns[0].clone(); - let histogram_column = planner.ctx.field_columns[1].clone(); - let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); - let plan = LogicalPlanBuilder::from(input) - .aggregate(vec![col("ts")], aggregate_exprs) - .unwrap() - .filter( - planner - .mixed_ignored_histogram_filter_expr(op, &histogram_column) - .unwrap(), - ) - .unwrap() - .project([col(&float_column), col("ts")]) - .unwrap() - .build() - .unwrap(); - - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(values(&batches, &float_column), vec![1.25]); - let mut warnings = vec![]; - let mut infos = vec![]; - collector.append_to(&mut warnings, &mut infos); - assert!(warnings.is_empty()); - assert_eq!(infos, expected_infos); - } - } - - #[tokio::test] - async fn test_mixed_or_value_aliases_do_not_replace_labels() { - let left = source( - "lhs", - false, - 1, - vec![("job", Some("job")), ("k", Some("float"))], - DirectOrValue::Float64(1.0), - ); - let right = source( - "rhs", - false, - 1, - vec![ - ("job", Some("job")), - ("k", Some("histogram")), - (greptime_value(), Some("value-label")), - ], - DirectOrValue::NativeHistogram(direct_or_histogram()), - ); - let table_provider = build_test_table_provider_with_fields( - &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], - &[], - ) - .await; - let mut planner = PromPlanner { - table_provider, - ctx: PromPlannerContext::default(), - promql_annotations: None, - }; - let left = LogicalPlanBuilder::from(scan(&left)) - .project(vec![ - col("ts"), - col("job"), - col("k"), - col("v").alias(greptime_value()), - ]) - .unwrap() - .build() - .unwrap(); - let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); - let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); - let plan = planner - .or_operator( - left, - scan(&right), - left_context.tag_columns.iter().cloned().collect(), - right_context.tag_columns.iter().cloned().collect(), - left_context, - right_context, - &or_modifier("lhs or on(k) rhs"), - ) - .unwrap(); - - assert_eq!( - plan.schema() - .field_with_name(None, greptime_value()) - .unwrap() - .data_type(), - &ArrowDataType::Utf8 - ); - assert!( - planner - .ctx - .field_columns - .iter() - .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) - ); - assert!(PromPlanner::field_columns_are_alternative_samples( - plan.schema(), - &planner.ctx.field_columns - )); - let (_, batches) = execute(plan, &build_query_engine_state()).await; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); - let labels = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name(greptime_value()) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .collect::>(); - assert_eq!(labels, vec!["value-label"]); - } - - #[tokio::test] - async fn test_mixed_or_routes_float_histogram_and_label_functions() { - for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { - unreachable!() - }; - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let values = batches - .iter() - .flat_map(|batch| { - batch - .schema() - .fields() - .iter() - .position(|field| field.data_type() == &ArrowDataType::Float64) - .map(|index| { - batch - .column(index) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - }) - .into_iter() - .flatten() - }) - .collect::>(); - assert_eq!(values, vec![expected], "{function}"); - } - - let (mut planner, input) = mixed_direct_or(false).await; - let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( - input.schema(), - &planner.ctx.field_columns, - ); - let PromExpr::Call(call) = - parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() - else { - unreachable!() - }; - let args = planner.create_function_args(&call.args.args).unwrap(); - let state = build_query_engine_state(); - let (mut exprs, _) = planner - .create_function_expr(&call.func, args.literals, input.schema(), &state) - .unwrap(); - exprs.insert(0, planner.create_time_index_column_expr().unwrap()); - exprs.extend(planner.create_tag_column_exprs().unwrap()); - let plan = LogicalPlanBuilder::from(input) - .project(exprs) - .unwrap() - .filter( - planner - .create_empty_values_filter_expr(preserve_any_value) - .unwrap(), - ) - .unwrap() - .build() - .unwrap(); - let (_, batches) = execute(plan, &state).await; - let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); - assert_eq!(sample_count, 2); - } - - /// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: - /// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. - async fn build_count_values_table_provider() -> DfTableSourceProvider { - build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await - } - - /// Table provider with a single metric `cv_metric` holding one series per given value at - /// `ts=1000` (`k` names the series, the sample value is the given value). - async fn build_count_values_table_provider_with_values( - values: &[f64], - ) -> DfTableSourceProvider { - build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( - values.to_vec(), - ))) - .await - } - - /// Like [`build_count_values_table_provider_with_values`], but with a caller provided value - /// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). - async fn build_count_values_table_provider_with_value_array( - values: ArrayRef, - ) -> DfTableSourceProvider { - let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); - let catalog_list = MemoryCatalogManager::with_default_setup(); - let columns = vec![ - ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), - ColumnSchema::new(greptime_value().to_string(), value_data_type, true), - ]; - let schema = Arc::new(Schema::new(columns)); - let table_meta = TableMetaBuilder::empty() - .schema(schema.clone()) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(1024) - .build() - .unwrap(); - let table_info = Arc::new( - TableInfoBuilder::default() - .table_id(3_001) - .name("cv_metric") - .meta(table_meta) - .build() - .unwrap(), - ); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from( - (0..values.len()) - .map(|index| format!("k{index}")) - .collect::>(), - )), - Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), - values.clone(), - ], - ) - .unwrap(); - let backing = GreptimeMemTable::new_with_catalog( - "cv_metric", - GreptimeRecordBatch::from_df_record_batch(schema, batch), - 3_001, - DEFAULT_CATALOG_NAME.to_string(), - DEFAULT_SCHEMA_NAME.to_string(), - ); - let table = Arc::new(Table::new( - table_info, - FilterPushDownType::Unsupported, - backing.data_source(), - )); - - assert!( - catalog_list - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: "cv_metric".to_string(), - table_id: 3_001, - table, - }) - .is_ok() - ); - - DfTableSourceProvider::new( - catalog_list, - false, - QueryContext::arc(), - DummyDecoder::arc(), - false, - ) - } - - /// Collects `(label, value)` pairs of a `count_values` result, where `label` is the - /// PromQL label generated by `count_values` and `value` is the aggregated sample value. - /// - /// The generated label holds the original sample value in PromQL's textual form, so it - /// is asserted as a string: comparing it as a number would not catch formatting bugs - /// (`1.0` instead of `1`, scientific notation, ...). - fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { - let mut rows = batches - .iter() - .flat_map(|batch| { - // The aggregated value is the only numeric column that is not the generated label. - let value_index = batch - .schema() - .fields() - .iter() - .position(|field| { - field.name() != label - && matches!( - field.data_type(), - ArrowDataType::Float64 - | ArrowDataType::Int64 - | ArrowDataType::UInt64 - ) - }) - .expect("no aggregated value column"); - let labels = batch - .column_by_name(label) - .expect("no generated label column") - .as_any() - .downcast_ref::() - .expect("the generated label must be a string column"); - let values = datafusion::arrow::compute::cast( - batch.column(value_index), - &ArrowDataType::Float64, - ) - .unwrap(); - let values = values.as_any().downcast_ref::().unwrap(); - labels - .iter() - .zip(values.iter()) - .map(|(label, value)| (label.unwrap(), value.unwrap())) - .collect::>() - }) - .collect::>(); - rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); - rows - } - - /// Asserts that a `count_values` result holds one sample per label set and evaluation - /// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat - /// a label set (that would mean the samples were still grouped by the overwritten input - /// label). - fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { - let mut seen = HashMap::>::new(); - for batch in batches { - let timestamp_index = batch - .schema() - .fields() - .iter() - .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) - .expect("no timestamp column"); - let timestamps = batch - .column(timestamp_index) - .as_any() - .downcast_ref::() - .expect("timestamp column is not a millisecond timestamp"); - let labels = batch - .column_by_name(label) - .expect("no generated label column") - .as_any() - .downcast_ref::() - .expect("the generated label must be a string column"); - for (timestamp, label) in timestamps.iter().zip(labels.iter()) { - let timestamp = timestamp.unwrap(); - let label = label.unwrap(); - assert!( - seen.entry(timestamp).or_default().insert(label.to_string()), - "duplicated label set `{label}` at timestamp {timestamp}" - ); - } - } - } - - #[tokio::test] - async fn test_count_values_generated_label_survives_enclosing_expr() { - // https://github.com/GreptimeTeam/greptimedb/issues/9181 - for (case, label) in [ - (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), - ( - r#"abs(count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"round(count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), - ( - r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, - "v", - ), - ( - r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, - "v", - ), - ( - r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, - "v", - ), - // The generated label overwrites an input label with the same name. - ( - r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, - "ip", - ), - ( - r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, - "ip", - ), - ] { - let plan = PromPlanner::stmt_to_plan( - build_test_table_provider_with_fields( - &[( - DEFAULT_SCHEMA_NAME.to_string(), - "prometheus_tsdb_head_series".to_string(), - )], - &["ip"], - ) - .await, - &build_eval_stmt(case), - &build_query_engine_state(), - ) - .await - .unwrap(); - - let label_columns = plan - .schema() - .fields() - .iter() - .filter(|field| field.name() == label) - .count(); - assert_eq!( - label_columns, - 1, - "{case}: the `{label}` label must survive: {}", - plan.display_indent() - ); - } - } - - #[tokio::test] - async fn test_count_values_generated_label_in_enclosing_expr_execute() { - // https://github.com/GreptimeTeam/greptimedb/issues/9181 - let state = build_query_engine_state(); - // (query, generated label, expected `(label value, aggregated value)` pairs) - for (query, label, expected) in [ - ( - r#"count_values("v", cv_metric)"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"abs(count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"round(count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"count_values("v", cv_metric) + 1"#, - "v", - vec![("1", 3.0), ("2", 2.0)], - ), - ( - r#"sum by (v) (count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - ( - r#"topk(10, count_values("v", cv_metric))"#, - "v", - vec![("1", 2.0), ("2", 1.0)], - ), - // The generated label overwrites the input label with the same name, and the - // samples are grouped by the generated label only: `{k="1"}` holds the two - // samples of value `1.0` instead of one row per (overwritten label, value). - ( - r#"count_values("k", cv_metric) by (k)"#, - "k", - vec![("1", 2.0), ("2", 1.0)], - ), - ] { - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider().await, - &operator_eval_stmt(query), - &state, - ) - .await - .unwrap_or_else(|err| panic!("{query}: {err}")); - - assert_eq!( - plan.schema() - .fields() - .iter() - .filter(|field| field.name() == label) - .count(), - 1, - "{query}: {}", - plan.display_indent() - ); - - let (_, batches) = execute(plan, &state).await; - assert_eq!(count_values_rows(&batches, label), expected, "{query}"); - assert_unique_label_set_per_timestamp(&batches, label); - } - } - - #[tokio::test] - async fn test_count_values_label_is_prometheus_formatted_value() { - // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: - // the shortest decimal form of the sample value without an exponent. The label is a - // label, so it must be a string column holding exactly that text: arrow's - // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. - let state = build_query_engine_state(); - // Samples are grouped by that formatted text, exactly like Prometheus groups by the - // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that - // round to the same text share one group. `0.0` formats as "0". - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider_with_values(&[ - -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, - ]) - .await, - &operator_eval_stmt(r#"count_values("v", cv_metric)"#), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - assert_unique_label_set_per_timestamp(&batches, "v"); - let mut labels = count_values_rows(&batches, "v") - .into_iter() - .map(|(label, _)| label) - .collect::>(); - labels.sort(); - assert_eq!( - labels, - vec![ - "-0", - "0", - "0.0000001", - "0.5", - "1", - "1000000000000000000000", - "2.5", - "200", - ] - ); - } - - #[tokio::test] - async fn test_count_values_groups_by_formatted_value_for_bigint_input() { - // The grouping key of `count_values` is the formatted sample value, not the raw input - // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and - // `2^53 + 1`) cast and format to the same label, so they must share one group and one - // count, exactly like Prometheus, which groups by the generated label text. - let state = build_query_engine_state(); - let plan = PromPlanner::stmt_to_plan( - build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ - 9_007_199_254_740_992_i64, - 9_007_199_254_740_993_i64, - ]))) - .await, - &operator_eval_stmt(r#"count_values("v", cv_metric)"#), - &state, - ) - .await - .unwrap(); - - let (_, batches) = execute(plan, &state).await; - // One timestamp must never carry the same label set twice. - assert_unique_label_set_per_timestamp(&batches, "v"); - assert_eq!( - count_values_rows(&batches, "v"), - vec![("9007199254740992", 2.0)] - ); - } -} +mod test; diff --git a/src/query/src/promql/planner/function_plans.rs b/src/query/src/promql/planner/function_plans.rs new file mode 100644 index 00000000000..68a254dd2bc --- /dev/null +++ b/src/query/src/promql/planner/function_plans.rs @@ -0,0 +1,473 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Dedicated helper plans for PromQL functions that cannot be expressed as a +//! plain scalar function over the input plan: the histogram helpers, `vector()`, +//! `scalar()`, and `absent()`. + +use std::sync::Arc; + +use common_query::prelude::greptime_value; +use datafusion::functions_aggregate::expr_fn::first_value; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; +use datafusion::prelude::{Column, Expr as DfExpr}; +use datafusion::scalar::ScalarValue; +use datafusion_common::DFSchema; +use datafusion_expr::expr_fn::when; +use datafusion_expr::simplify::SimplifyContext; +use datafusion_expr::{col, lit}; +use datafusion_functions::core::coalesce; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::{ + Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, ScalarCalculate, +}; +use promql::functions::{NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile}; +use promql_parser::label::MatchOp; +use promql_parser::parser::{Expr as PromExpr, FunctionArgs as PromFunctionArgs}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::promql::error::{ + DataFusionPlanningSnafu, FunctionInvalidArgumentSnafu, MultiFieldsNotSupportedSnafu, + PromqlPlanNodeSnafu, Result, TimeIndexNotFoundSnafu, ValueNotFoundSnafu, +}; +use crate::promql::planner::{ + LE_COLUMN_NAME, PromPlanner, SCALAR_FUNCTION, SPECIAL_ABSENT_FUNCTION, + SPECIAL_HISTOGRAM_FRACTION, SPECIAL_HISTOGRAM_QUANTILE, SPECIAL_TIME_FUNCTION, + SPECIAL_VECTOR_FUNCTION, +}; +use crate::query_engine::QueryEngineState; + +impl PromPlanner { + /// Create a classic, native, or mixed histogram helper plan. + pub(super) async fn create_histogram_plan( + &mut self, + function_name: &str, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + let float_literal = |param: &PromExpr| -> Result { + let value = (|| { + let expr = Self::get_param_as_literal_expr( + Some(param), + None, + Some(ArrowDataType::Float64), + ) + .ok()?; + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; + let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { + return None; + }; + let ScalarValue::Float64(Some(value)) = + value.cast_to(&ArrowDataType::Float64).ok()? + else { + return None; + }; + Some(value) + })() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + })?; + Ok(value) + }; + let (function, input) = match (function_name, args.args.as_slice()) { + (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( + HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), + input.as_ref().clone(), + ), + (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( + HistogramFoldOperation::Fraction { + lower: float_literal(lower)?.into(), + upper: float_literal(upper)?.into(), + }, + input.as_ref().clone(), + ), + _ => { + return FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + } + .fail(); + } + }; + + let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; + // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a + // stable series identifier anymore. HistogramFold must not treat it as a label column. + let input_plan = self.strip_tsid_column(input_plan)?; + self.ctx.use_tsid = false; + + if let Some((float_field, histogram_field)) = + Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + { + if self.ctx.has_le_tag() { + return self.create_mixed_histogram_plan( + function, + input_plan, + float_field, + histogram_field, + ); + } + self.ctx.field_columns = vec![histogram_field]; + } + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_plan(function, input_plan); + } + + if !self.ctx.has_le_tag() { + // Return empty result instead of error when 'le' column is not found + // This handles the case when histogram metrics don't exist + return Ok(LogicalPlan::EmptyRelation( + datafusion::logical_expr::EmptyRelation { + produce_one_row: false, + schema: input_plan.schema().clone(), + }, + )); + } + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + // FIXME(ruihang): support multi fields + let field_column = self + .ctx + .field_columns + .first() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function.function_name().to_string(), + })? + .clone(); + // remove le column from tag columns + self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); + + let fold = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + field_column, + time_index_column, + function, + None, + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(fold), + })) + } + + fn create_native_histogram_expr( + &self, + function: HistogramFoldOperation, + field_column: &str, + ) -> DfExpr { + let field = DfExpr::Column(Column::from_name(field_column)); + let (func, args) = match function { + HistogramFoldOperation::Quantile(quantile) => ( + Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(quantile))], + ), + HistogramFoldOperation::Fraction { lower, upper } => ( + Arc::new(NativeHistogramFraction::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(lower)), lit(f64::from(upper))], + ), + }; + DfExpr::ScalarFunction(ScalarFunction { func, args }) + } + + fn create_native_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + ) -> Result { + ensure!( + self.ctx.field_columns.len() == 1, + MultiFieldsNotSupportedSnafu { + operator: function.function_name() + }, + ); + + let field_column = self.ctx.field_columns[0].clone(); + let function_expr = self.create_native_histogram_expr(function, &field_column); + let display_name = function_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + + let project_exprs = std::iter::once(self.create_time_index_column_expr()?) + .chain(std::iter::once(function_expr.alias(display_name))) + .chain(self.create_tag_column_exprs()?) + .collect::>(); + + LogicalPlanBuilder::from(input_plan) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .filter(self.create_empty_values_filter_expr(false)?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn create_mixed_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + float_field: String, + histogram_field: String, + ) -> Result { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let tag_columns = self.ctx.tag_columns.clone(); + let folded = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + float_field.clone(), + time_index_column.clone(), + function, + Some(histogram_field.clone()), + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + let record_collision = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + "vector contains a mix of classic and native histograms".to_string(), + self.promql_annotations.clone(), + )), + args: vec![col(&float_field), col(&histogram_field)], + }); + let keep = when( + col(&float_field) + .is_not_null() + .and(col(&histogram_field).is_not_null()), + record_collision, + ) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu)?; + + let native_expr = self.create_native_histogram_expr(function, &histogram_field); + let output_field = native_expr.schema_name().to_string(); + let value = DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![col(&float_field), native_expr], + }); + self.ctx.field_columns = vec![output_field.clone()]; + LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { + node: Arc::new(folded), + })) + .filter(keep) + .context(DataFusionPlanningSnafu)? + .project( + std::iter::once(col(&time_index_column)) + .chain(std::iter::once(value.alias(output_field))) + .chain(tag_columns.iter().map(col)), + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + /// Create a [SPECIAL_VECTOR_FUNCTION] plan + pub(super) async fn create_vector_plan( + &mut self, + args: &PromFunctionArgs, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_VECTOR_FUNCTION.to_string(), + } + .fail(); + } + let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; + + // reuse `SPECIAL_TIME_FUNCTION` as name of time index column + self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); + self.ctx.reset_table_name_and_schema(); + self.ctx.tag_columns = vec![]; + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns = vec![greptime_value().to_string()]; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new( + EmptyMetric::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + SPECIAL_TIME_FUNCTION.to_string(), + greptime_value().to_string(), + Some(lit), + ) + .context(DataFusionPlanningSnafu)?, + ), + })) + } + + /// Create a [SCALAR_FUNCTION] plan + pub(super) async fn create_scalar_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + ensure!( + args.len() == 1, + FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION + } + ); + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + let input_schema = input.schema().clone(); + let alternative_samples = + Self::field_columns_are_alternative_samples(&input_schema, &self.ctx.field_columns); + let histogram_fields = self + .ctx + .field_columns + .iter() + .filter(|field| Self::field_column_is_native_histogram(&input_schema, field)) + .count(); + ensure!( + self.ctx.field_columns.len() == 1 || alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: SCALAR_FUNCTION + }, + ); + let scalar_field = self + .ctx + .field_columns + .iter() + .find(|field| !Self::field_column_is_native_histogram(&input_schema, field)) + .or_else(|| self.ctx.field_columns.first()) + .cloned() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: SCALAR_FUNCTION, + })?; + let input = if histogram_fields == self.ctx.field_columns.len() { + // scalar() ignores histogram samples. An empty input makes ScalarCalculate emit NaN + // for every evaluation timestamp without attempting a Struct-to-Float64 cast. + LogicalPlanBuilder::from(input) + .filter(lit(false)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else if histogram_fields > 0 { + // A mixed vector contributes only its float samples to scalar(). + LogicalPlanBuilder::from(input) + .filter(DfExpr::Column(Column::from_name(&scalar_field)).is_not_null()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + input + }; + let scalar_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + ScalarCalculate::new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + input, + self.ctx.time_index_column.as_ref().unwrap(), + &self.ctx.tag_columns, + &scalar_field, + self.ctx.table_name.as_deref(), + ) + .context(PromqlPlanNodeSnafu)?, + ), + }); + // scalar plan have no tag columns + self.ctx.tag_columns.clear(); + self.ctx.aggregation_field_labels.clear(); + self.ctx.field_columns.clear(); + self.ctx + .field_columns + .push(scalar_plan.schema().field(1).name().clone()); + Ok(scalar_plan) + } + + /// Create a [SPECIAL_ABSENT_FUNCTION] plan + pub(super) async fn create_absent_plan( + &mut self, + args: &PromFunctionArgs, + query_engine_state: &QueryEngineState, + ) -> Result { + if args.args.len() != 1 { + return FunctionInvalidArgumentSnafu { + fn_name: SPECIAL_ABSENT_FUNCTION.to_string(), + } + .fail(); + } + let input = self + .prom_expr_to_plan(&args.args[0], query_engine_state) + .await?; + + let time_index_expr = self.create_time_index_column_expr()?; + let first_field_expr = + self.create_field_column_exprs()? + .pop() + .with_context(|| ValueNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let first_value_expr = first_value(first_field_expr, vec![]); + + let ordered_aggregated_input = LogicalPlanBuilder::from(input) + .aggregate( + vec![time_index_expr.clone()], + vec![first_value_expr.clone()], + ) + .context(DataFusionPlanningSnafu)? + .sort(vec![time_index_expr.sort(true, false)]) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + let fake_labels = self + .ctx + .selector_matcher + .iter() + .filter_map(|matcher| match matcher.op { + MatchOp::Equal => Some((matcher.name.clone(), matcher.value.clone())), + _ => None, + }) + .collect::>(); + + // Create the absent plan + let absent_plan = LogicalPlan::Extension(Extension { + node: Arc::new( + Absent::try_new( + self.ctx.start, + self.ctx.end, + self.ctx.interval, + self.ctx.time_index_column.as_ref().unwrap().clone(), + self.ctx.field_columns[0].clone(), + fake_labels, + ordered_aggregated_input, + ) + .context(DataFusionPlanningSnafu)?, + ), + }); + + // The absent series carries the equality matchers as labels, not the input's + // tags or value fields, so the input's field grouping labels no longer apply. + self.ctx.aggregation_field_labels.clear(); + Ok(absent_plan) + } +} diff --git a/src/query/src/promql/planner/island.rs b/src/query/src/promql/planner/island.rs new file mode 100644 index 00000000000..d7c04f75015 --- /dev/null +++ b/src/query/src/promql/planner/island.rs @@ -0,0 +1,536 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The binary-island fast path of the PromQL planner. + +use std::collections::{BTreeSet, HashMap}; + +use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; +use datafusion::common::DFSchemaRef; +use datafusion::logical_expr::expr::Alias; +use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion_common::{NullEquality, TableReference}; +use datafusion_expr::lit; +use promql_parser::label::{METRIC_NAME, MatchOp, Matcher}; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{ + BinaryExpr as PromBinaryExpr, Expr as PromExpr, Offset, ParenExpr, UnaryExpr, + VectorMatchCardinality, VectorSelector, +}; +use snafu::ResultExt; + +use crate::promql::error::{DataFusionPlanningSnafu, Result}; +use crate::promql::planner::{PromPlanner, PromPlannerContext}; + +/// Prefix for generated binary island leaf aliases. +const BINARY_ISLAND_LEAF_ALIAS_PREFIX: &str = "__prom_v"; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct VectorLeafKey { + metric_name: String, + matchers: Vec<(String, String, String)>, + or_matchers: Vec>, + offset_ms: i128, + at: String, +} + +#[derive(Debug, Clone)] +struct IslandLeaf { + selector: VectorSelector, + display_table: String, +} + +#[derive(Debug, Clone)] +enum IslandExpr { + VectorLeaf(usize), + Scalar(DfExpr), + Unary { + input: Box, + }, + Binary { + op: TokenType, + lhs: Box, + rhs: Box, + }, +} + +impl IslandExpr { + fn try_new(expr: &PromExpr, env: &mut IslandCollectEnv) -> Option { + if let Some(expr) = PromPlanner::try_build_literal_expr(expr) { + return Some(Self::Scalar(expr)); + } + + match expr { + PromExpr::Paren(ParenExpr { expr }) => Self::try_new(expr, env), + PromExpr::VectorSelector(selector) => { + let leaf = env.intern_leaf(selector)?; + Some(Self::VectorLeaf(leaf)) + } + PromExpr::Unary(UnaryExpr { expr }) => { + let input = Self::try_new(expr, env)?; + Some(Self::Unary { + input: Box::new(input), + }) + } + PromExpr::Binary(PromBinaryExpr { + lhs, + rhs, + op, + modifier, + }) if matches!( + op.id(), + token::T_ADD + | token::T_SUB + | token::T_MUL + | token::T_DIV + | token::T_MOD + | token::T_POW + | token::T_ATAN2 + ) && modifier.as_ref().is_none_or(|modifier| { + !modifier.return_bool + && modifier.matching.is_none() + && matches!(modifier.card, VectorMatchCardinality::OneToOne) + && modifier.fill_values.lhs.is_none() + && modifier.fill_values.rhs.is_none() + }) => + { + let lhs = Self::try_new(lhs, env)?; + let rhs = Self::try_new(rhs, env)?; + Some(Self::Binary { + op: *op, + lhs: Box::new(lhs), + rhs: Box::new(rhs), + }) + } + _ => None, + } + } +} + +#[derive(Debug, Default)] +struct IslandCollectEnv { + leaf_by_key: HashMap, + leaves: Vec, + vector_occurrences: usize, +} + +#[derive(Debug)] +struct PlannedIslandLeaf { + plan: LogicalPlan, + ctx: PromPlannerContext, + alias: TableReference, + display_table: String, +} + +#[derive(Debug)] +struct IslandFieldExprs { + exprs: Vec, + names: Vec, + scalar: bool, +} + +impl VectorLeafKey { + fn from_selector(selector: &VectorSelector) -> Option { + let mut metric_name = selector.name.clone(); + let mut matchers = Vec::with_capacity(selector.matchers.matchers.len()); + let matcher_key = |matcher: &Matcher| { + ( + matcher.name.clone(), + matcher.op.to_string(), + matcher.value.clone(), + ) + }; + + for matcher in &selector.matchers.matchers { + if matcher.name == METRIC_NAME { + if matcher.op != MatchOp::Equal || metric_name.is_some() { + return None; + } + metric_name = Some(matcher.value.clone()); + } else { + matchers.push(matcher_key(matcher)); + } + } + matchers.sort(); + + let mut or_matchers = selector + .matchers + .or_matchers + .iter() + .map(|group| { + let mut group = group.iter().map(matcher_key).collect::>(); + group.sort(); + group + }) + .collect::>(); + or_matchers.sort(); + + Some(Self { + metric_name: metric_name?, + matchers, + or_matchers, + offset_ms: match &selector.offset { + Some(Offset::Pos(duration)) => duration.as_millis() as i128, + Some(Offset::Neg(duration)) => -(duration.as_millis() as i128), + None => 0, + }, + at: format!("{:?}", selector.at), + }) + } +} + +impl IslandCollectEnv { + fn intern_leaf(&mut self, selector: &VectorSelector) -> Option { + self.vector_occurrences += 1; + let key = VectorLeafKey::from_selector(selector)?; + if let Some(id) = self.leaf_by_key.get(&key) { + return Some(*id); + } + + let id = self.leaves.len(); + self.leaves.push(IslandLeaf { + selector: selector.clone(), + display_table: key.metric_name.clone(), + }); + self.leaf_by_key.insert(key, id); + Some(id) + } +} + +impl PromPlanner { + pub(super) async fn try_plan_binary_island( + &mut self, + binary_expr: &PromBinaryExpr, + ) -> Result> { + let original_ctx = self.ctx.clone(); + let mut collect_env = IslandCollectEnv::default(); + let Some(island_expr) = + IslandExpr::try_new(&PromExpr::Binary(binary_expr.clone()), &mut collect_env) + else { + return Ok(None); + }; + + if collect_env.leaves.is_empty() + || collect_env.vector_occurrences <= collect_env.leaves.len() + { + return Ok(None); + } + + let mut planned_leaves = Vec::with_capacity(collect_env.leaves.len()); + for (idx, leaf) in collect_env.leaves.iter().enumerate() { + let plan = self + .prom_vector_selector_to_plan(&leaf.selector, false) + .await?; + let ctx = self.ctx.clone(); + let alias = TableReference::bare(format!("{BINARY_ISLAND_LEAF_ALIAS_PREFIX}{idx}")); + let plan = LogicalPlanBuilder::from(plan) + .alias(alias.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + planned_leaves.push(PlannedIslandLeaf { + plan, + ctx, + alias, + display_table: leaf.display_table.clone(), + }); + } + + if planned_leaves.iter().any(|leaf| { + Self::field_columns_contain_native_histogram( + leaf.plan.schema(), + &leaf.ctx.field_columns, + ) + }) { + self.ctx = original_ctx; + return Ok(None); + } + + if !Self::binary_island_join_contexts_supported(&planned_leaves) { + self.ctx = original_ctx; + return Ok(None); + } + + let mut input = planned_leaves[0].plan.clone(); + for right_idx in 1..planned_leaves.len() { + input = self.join_binary_island_leaf( + input, + &planned_leaves[0], + &planned_leaves[right_idx], + )?; + } + + let field_exprs = + Self::build_binary_island_field_exprs(&island_expr, &planned_leaves, input.schema())?; + if field_exprs.scalar || field_exprs.exprs.is_empty() { + self.ctx = original_ctx; + return Ok(None); + } + + let plan = self.project_binary_island( + input, + &planned_leaves[0].alias, + &planned_leaves[0].ctx, + field_exprs, + )?; + Ok(Some(plan)) + } + + fn binary_island_join_contexts_supported(leaves: &[PlannedIslandLeaf]) -> bool { + if leaves + .iter() + .any(|leaf| leaf.ctx.time_index_column.is_none()) + { + return false; + } + + if leaves.len() <= 1 { + return true; + } + + let first_tags = leaves[0].ctx.tag_columns.iter().collect::>(); + + leaves.iter().skip(1).all(|leaf| { + (Self::plan_has_tsid_column(&leaves[0].plan) && Self::plan_has_tsid_column(&leaf.plan)) + || leaf.ctx.tag_columns.iter().collect::>() == first_tags + }) + } + + fn join_binary_island_leaf( + &self, + left: LogicalPlan, + first_leaf: &PlannedIslandLeaf, + right_leaf: &PlannedIslandLeaf, + ) -> Result { + let only_join_time_index = (first_leaf.ctx.tag_columns.is_empty() + || right_leaf.ctx.tag_columns.is_empty()) + && !first_leaf + .ctx + .tag_columns + .iter() + .chain(&right_leaf.ctx.tag_columns) + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL); + let (mut left_keys, mut right_keys, force_empty_join) = self.binary_join_key_columns( + left.schema(), + right_leaf.plan.schema(), + &first_leaf.ctx, + &right_leaf.ctx, + only_join_time_index, + &None, + )?; + + if let (Some(left_time_index_column), Some(right_time_index_column)) = ( + first_leaf.ctx.time_index_column.clone(), + right_leaf.ctx.time_index_column.clone(), + ) { + left_keys.insert(left_time_index_column); + right_keys.insert(right_time_index_column); + } + + LogicalPlanBuilder::from(left) + .join_detailed( + right_leaf.plan.clone(), + JoinType::Inner, + ( + left_keys + .into_iter() + .map(|name| Column::new(Some(first_leaf.alias.clone()), name)) + .collect::>(), + right_keys + .into_iter() + .map(|name| Column::new(Some(right_leaf.alias.clone()), name)) + .collect::>(), + ), + force_empty_join.then_some(lit(false)), + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + fn build_binary_island_field_exprs( + expr: &IslandExpr, + leaves: &[PlannedIslandLeaf], + schema: &DFSchemaRef, + ) -> Result { + match expr { + IslandExpr::VectorLeaf(id) => { + let leaf = &leaves[*id]; + let exprs = leaf + .ctx + .field_columns + .iter() + .map(|field| { + schema + .qualified_field_with_name(Some(&leaf.alias), field) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }) + .collect::>>()?; + let names = leaf + .ctx + .field_columns + .iter() + .map(|field| format!("{}.{}", leaf.display_table, field)) + .collect(); + Ok(IslandFieldExprs { + exprs, + names, + scalar: false, + }) + } + IslandExpr::Scalar(expr) => Ok(IslandFieldExprs { + exprs: vec![expr.clone()], + names: vec![expr.schema_name().to_string()], + scalar: true, + }), + IslandExpr::Unary { input } => { + let input = Self::build_binary_island_field_exprs(input, leaves, schema)?; + let mut exprs = Vec::with_capacity(input.exprs.len()); + let mut names = Vec::with_capacity(input.names.len()); + for (expr, name) in input.exprs.into_iter().zip(input.names) { + exprs.push(DfExpr::Negative(Box::new(expr))); + names.push(format!("-{name}")); + } + Ok(IslandFieldExprs { + exprs, + names, + scalar: input.scalar, + }) + } + IslandExpr::Binary { op, lhs, rhs } => { + let same_leaf = match (&**lhs, &**rhs) { + (IslandExpr::VectorLeaf(left), IslandExpr::VectorLeaf(right)) + if left == right => + { + Some(*left) + } + _ => None, + }; + let lhs = Self::build_binary_island_field_exprs(lhs, leaves, schema)?; + let rhs = Self::build_binary_island_field_exprs(rhs, leaves, schema)?; + let expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let scalar = lhs.scalar && rhs.scalar; + let op = op.to_string(); + + let (exprs, names) = match (lhs.scalar, rhs.scalar) { + (true, true) => { + let expr = expr_builder(lhs.exprs[0].clone(), rhs.exprs[0].clone())?; + let name = format!("{} {op} {}", lhs.names[0], rhs.names[0]); + (vec![expr], vec![name]) + } + (true, false) => { + let mut exprs = Vec::with_capacity(rhs.exprs.len()); + let mut names = Vec::with_capacity(rhs.names.len()); + for (rhs_expr, rhs_name) in rhs.exprs.into_iter().zip(rhs.names) { + exprs.push(expr_builder(lhs.exprs[0].clone(), rhs_expr)?); + names.push(format!("{} {op} {rhs_name}", lhs.names[0])); + } + (exprs, names) + } + (false, true) => { + let mut exprs = Vec::with_capacity(lhs.exprs.len()); + let mut names = Vec::with_capacity(lhs.names.len()); + for (lhs_expr, lhs_name) in lhs.exprs.into_iter().zip(lhs.names) { + exprs.push(expr_builder(lhs_expr, rhs.exprs[0].clone())?); + names.push(format!("{lhs_name} {op} {}", rhs.names[0])); + } + (exprs, names) + } + (false, false) => { + let mut exprs = Vec::new(); + let mut names = Vec::new(); + for (idx, ((lhs_expr, rhs_expr), (mut lhs_name, mut rhs_name))) in lhs + .exprs + .into_iter() + .zip(rhs.exprs) + .zip(lhs.names.into_iter().zip(rhs.names)) + .enumerate() + { + if let Some(leaf) = same_leaf { + let field = leaves[leaf] + .ctx + .field_columns + .get(idx) + .cloned() + .unwrap_or_else(|| lhs_name.clone()); + lhs_name = format!("lhs.{field}"); + rhs_name = format!("rhs.{field}"); + } + exprs.push(expr_builder(lhs_expr, rhs_expr)?); + names.push(format!("{lhs_name} {op} {rhs_name}")); + } + (exprs, names) + } + }; + + Ok(IslandFieldExprs { + exprs, + names, + scalar, + }) + } + } + } + + fn project_binary_island( + &mut self, + input: LogicalPlan, + base_alias: &TableReference, + base_ctx: &PromPlannerContext, + field_exprs: IslandFieldExprs, + ) -> Result { + self.ctx = base_ctx.clone(); + + let schema = input.schema(); + let non_field_exprs = base_ctx + .tag_columns + .iter() + .chain(base_ctx.time_index_column.iter()) + .map(|column| { + schema + .qualified_field_with_name(Some(base_alias), column) + .context(DataFusionPlanningSnafu) + .map(|field| DfExpr::Column(field.into())) + }); + let tsid_expr = Self::optional_tsid_projection(schema, Some(base_alias), base_ctx.use_tsid) + .into_iter() + .map(Ok); + + self.ctx.field_columns = field_exprs.names; + let field_exprs = field_exprs + .exprs + .into_iter() + .zip(self.ctx.field_columns.iter()) + .map(|(expr, name)| Ok(DfExpr::Alias(Alias::new(expr, None::, name)))); + + let project_exprs = non_field_exprs + .chain(tsid_expr) + .chain(field_exprs) + .collect::>>()?; + + let plan = LogicalPlanBuilder::from(input) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + self.ctx.table_name = None; + self.ctx.schema_name = None; + + Ok(plan) + } +} diff --git a/src/query/src/promql/planner/set_operator.rs b/src/query/src/promql/planner/set_operator.rs new file mode 100644 index 00000000000..51cc100408f --- /dev/null +++ b/src/query/src/promql/planner/set_operator.rs @@ -0,0 +1,832 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Planning of the PromQL set operators (`and`, `unless`, and `or`). + +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::logical_expr::{Cast, Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; +use datafusion::scalar::ScalarValue; +use datafusion_common::{NullEquality, TableReference}; +use datatypes::arrow::datatypes::DataType as ArrowDataType; +use promql::extension_plan::UnionDistinctOn; +use promql_parser::parser::token::{self, TokenType}; +use promql_parser::parser::{BinModifier, LabelModifier, VectorMatchCardinality}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; + +use crate::promql::error::{ + ColumnNotFoundSnafu, CombineTableColumnMismatchSnafu, DataFusionPlanningSnafu, + MultiFieldsNotSupportedSnafu, Result, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, + UnexpectedTokenSnafu, UnsupportedVectorMatchSnafu, +}; +use crate::promql::planner::{ + OR_FLOAT_FIELD_PREFIX, OR_HISTOGRAM_FIELD_PREFIX, PromPlanner, PromPlannerContext, +}; + +impl PromPlanner { + // TODO(ruihang): change function name + #[allow(clippy::too_many_arguments)] + pub(super) fn or_operator( + &mut self, + left: LogicalPlan, + right: LogicalPlan, + left_tag_cols_set: HashSet, + right_tag_cols_set: HashSet, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, + ) -> Result { + let left_is_empty = Self::is_zero_row_empty_relation(&left); + let right_is_empty = Self::is_zero_row_empty_relation(&right); + match (left_is_empty, right_is_empty) { + (true, false) => { + self.ctx = right_context; + return Ok(right); + } + (false, true) => { + self.ctx = left_context; + return Ok(left); + } + (true, true) => { + self.ctx = left_context; + return Ok(left); + } + (false, false) => {} + } + + ensure!( + !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), + UnexpectedPlanExprSnafu { + desc: "OR operator input has zero columns", + } + ); + let left_has_alternative_samples = + Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); + let right_has_alternative_samples = Self::field_columns_are_alternative_samples( + right.schema(), + &right_context.field_columns, + ); + ensure!( + left_context.field_columns.len() == 1 || left_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + ensure!( + right_context.field_columns.len() == 1 || right_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" + } + ); + + // prepare hash sets + let all_tags = left_tag_cols_set + .union(&right_tag_cols_set) + .cloned() + .collect::>(); + let left_qualifier = left.schema().qualified_field(0).0.cloned(); + let right_qualifier = right.schema().qualified_field(0).0.cloned(); + let left_qualifier_string = left_qualifier + .as_ref() + .map(|l| l.to_string()) + .unwrap_or_default(); + let right_qualifier_string = right_qualifier + .as_ref() + .map(|r| r.to_string()) + .unwrap_or_default(); + let left_time_index_column = + left_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let right_time_index_column = + right_context + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: right_qualifier_string.clone(), + })?; + let native_histogram_type = Self::native_histogram_arrow_type(); + let is_numeric = |data_type: &ArrowDataType| { + matches!( + data_type, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) + }; + let left_fields = left_context + .field_columns + .iter() + .map(|name| { + left.schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let right_fields = right_context + .field_columns + .iter() + .map(|name| { + right + .schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let left_field = &left_fields[0]; + let right_field = &right_fields[0]; + let left_field_col = &left_field.0; + let right_field_col = &right_field.0; + let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { + fields.iter().all(|(_, _, data_type)| { + is_numeric(data_type) || data_type == &native_histogram_type + }) + }; + let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { + if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>(), + right_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>() + ), + } + .fail(); + } + true + } else { + (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) + || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) + }; + let target_field_type = if mixed_sample_types { + // Mixed vectors use the existing response representation: one nullable float column + // and one nullable native-histogram column. + ArrowDataType::Float64 + } else if left_field.2 == right_field.2 { + left_field.2.clone() + } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { + ArrowDataType::Float64 + } else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_field.2, right_field.2 + ), + } + .fail(); + }; + let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { + let mut reserved_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + for (name, _, _) in left_fields.iter().chain(&right_fields) { + reserved_names.remove(name); + } + reserved_names.extend(all_tags.iter().cloned()); + let unique_name = |prefix: &str, reserved_names: &mut HashSet| { + let mut index = 0; + loop { + let name = format!("{prefix}{index}"); + index += 1; + if reserved_names.insert(name.clone()) { + break name; + } + } + }; + let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); + let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); + (float_field, histogram_field) + } else { + (left_field_col.clone(), String::new()) + }; + let left_tag_types = left_tag_cols_set + .iter() + .map(|label| { + left.schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let right_tag_types = right_tag_cols_set + .iter() + .map(|label| { + right + .schema() + .fields() + .iter() + .find(|field| field.name() == label) + .map(|field| (label.clone(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() }) + }) + .collect::>>()?; + let mut target_tag_types = HashMap::with_capacity(all_tags.len()); + for label in &all_tags { + let Some(data_type) = + Self::common_label_data_type(left_tag_types.get(label), right_tag_types.get(label)) + else { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR label {label} has incompatible types: {:?} and {:?}", + left_tag_types.get(label), + right_tag_types.get(label) + ), + } + .fail(); + }; + target_tag_types.insert(label.clone(), data_type); + } + let left_has_tsid = left + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + let right_has_tsid = right + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME); + + // step 0: fill all columns in output schema + let mut all_columns_set = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + // Keep `__tsid` only when both sides contain it, otherwise it may break schema alignment + // (e.g. `unknown_metric or some_metric`). + if !(left_has_tsid && right_has_tsid) { + all_columns_set.remove(DATA_SCHEMA_TSID_COLUMN_NAME); + } + // remove time index column + all_columns_set.remove(&left_time_index_column); + all_columns_set.remove(&right_time_index_column); + if mixed_sample_types { + for (name, _, _) in left_fields.iter().chain(&right_fields) { + all_columns_set.remove(name); + } + all_columns_set.extend(all_tags.iter().cloned()); + all_columns_set.insert(mixed_float_field_col.clone()); + all_columns_set.insert(mixed_histogram_field_col.clone()); + } else if left_field_col != right_field_col { + // remove field column in the right + all_columns_set.remove(right_field_col); + } + let mut all_columns = all_columns_set.into_iter().collect::>(); + // sort to ensure the generated schema is not volatile + all_columns.sort_unstable(); + // use left time index column name as the result time index column name + all_columns.insert(0, left_time_index_column.clone()); + let mut occupied_column_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + + // step 1: align schema using project, fill non-exist columns with null + let aligned_label_expr = |col: &String, source_types: &HashMap| { + let target_type = &target_tag_types[col]; + if let Some(source_type) = source_types.get(col) { + let expr = DfExpr::Column(Column::new(None::, col)); + if source_type == target_type { + expr + } else { + DfExpr::Cast(Cast::new(Box::new(expr), target_type.clone())).alias(col.clone()) + } + } else { + DfExpr::Literal( + Self::string_scalar_value(target_type, None) + .expect("target label type is a string"), + None, + ) + .alias(col.clone()) + } + }; + let null_histogram = + ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; + let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], + output_col: &String| { + if output_col == &mixed_float_field_col { + if let Some((name, qualifier, data_type)) = fields + .iter() + .find(|(_, _, data_type)| is_numeric(data_type)) + { + let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); + if data_type == &ArrowDataType::Float64 { + expr.alias(output_col) + } else { + DfExpr::Cast(Cast::new(Box::new(expr), ArrowDataType::Float64)) + .alias(output_col) + } + } else { + DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) + } + } else { + fields + .iter() + .find(|(_, _, data_type)| data_type == &native_histogram_type) + .map(|(name, qualifier, _)| { + DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) + }) + .unwrap_or_else(|| { + DfExpr::Literal(null_histogram.clone(), None).alias(output_col) + }) + } + }; + let left_proj_exprs = all_columns.iter().map(|col| { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&left_fields, col) + } else if !mixed_sample_types + && col == left_field_col + && left_field.2 != target_field_type + { + DfExpr::Cast(Cast::new( + Box::new(DfExpr::Column(Column::new( + left_field.1.clone(), + left_field_col, + ))), + target_field_type.clone(), + )) + .alias(left_field_col.clone()) + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &left_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_time_index_expr = DfExpr::Column(Column::new( + right_qualifier.clone(), + right_time_index_column, + )) + .alias(left_time_index_column.clone()); + // The field column in right side may not have qualifier (it may be removed by join operation), + // so we need to find it from the schema. + // `skip（1)` to skip the time index column + let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { + // expr + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&right_fields, col) + } else if !mixed_sample_types && col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); + if right_field.2 != target_field_type { + DfExpr::Cast(Cast::new(Box::new(expr), target_field_type.clone())) + .alias(left_field_col.clone()) + } else if left_field_col != right_field_col { + expr.alias(left_field_col.clone()) + } else { + expr + } + } else if target_tag_types.contains_key(col) { + aligned_label_expr(col, &right_tag_types) + } else { + DfExpr::Column(Column::new(None::, col)) + } + }); + let right_proj_exprs = [right_time_index_expr] + .into_iter() + .chain(right_proj_exprs_without_time_index); + + let left_projected = LogicalPlanBuilder::from(left) + .project(left_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(left_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_projected = LogicalPlanBuilder::from(right) + .project(right_proj_exprs) + .context(DataFusionPlanningSnafu)? + .alias(right_qualifier_string.clone()) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 2: compute match columns + let mut match_columns = if let Some(modifier) = modifier + && let Some(matching) = &modifier.matching + { + match matching { + // keeps columns mentioned in `on` + LabelModifier::Include(on) => on.labels.clone(), + // removes columns memtioned in `ignoring` + LabelModifier::Exclude(ignoring) => { + let ignoring = ignoring.labels.iter().cloned().collect::>(); + all_tags.difference(&ignoring).cloned().collect() + } + } + } else { + all_tags.iter().cloned().collect() + }; + // sort to ensure the generated plan is not volatile + match_columns.sort_unstable(); + match_columns.dedup(); + occupied_column_names.extend( + left_projected + .schema() + .fields() + .iter() + .chain(right_projected.schema().fields().iter()) + .map(|field| field.name().clone()), + ); + + let visible_schema = left_projected.schema().clone(); + let visible_left_exprs = left_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let visible_right_exprs = right_projected + .schema() + .iter() + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + let mut left_match_exprs = Vec::with_capacity(match_columns.len()); + let mut right_match_exprs = Vec::with_capacity(match_columns.len()); + let mut next_internal_column = 0; + + for label in &match_columns { + let left_field = if left_tag_cols_set.contains(label) { + Some( + left_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let right_field = if right_tag_cols_set.contains(label) { + Some( + right_projected + .schema() + .iter() + .find(|(_, field)| field.name() == label) + .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) + .with_context(|| ColumnNotFoundSnafu { col: label.clone() })?, + ) + } else { + None + }; + let data_type = match (left_field.as_ref(), right_field.as_ref()) { + (Some((_, left_type)), Some((_, right_type))) if left_type == right_type => { + left_type.clone() + } + (Some((_, left_type)), Some((_, right_type))) => { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR match label {label} has incompatible types: {left_type:?} and {right_type:?}" + ), + } + .fail(); + } + (Some((_, data_type)), None) | (None, Some((_, data_type))) => data_type.clone(), + (None, None) => ArrowDataType::Utf8, + }; + let Some(value_type) = Self::string_value_data_type(&data_type).cloned() else { + return UnexpectedPlanExprSnafu { + desc: format!("OR match label {label} must be a string"), + } + .fail(); + }; + let internal_name = loop { + let name = format!("__promql_or_match_{next_internal_column}"); + next_internal_column += 1; + if occupied_column_names.insert(name.clone()) { + break name; + } + }; + left_match_exprs.push(Self::normalized_match_key_expr( + label, + left_field, + &value_type, + &internal_name, + )); + right_match_exprs.push(Self::normalized_match_key_expr( + label, + right_field, + &value_type, + &internal_name, + )); + } + + let left_augmented = LogicalPlanBuilder::from(left_projected) + .project(visible_left_exprs.into_iter().chain(left_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + let right_augmented = LogicalPlanBuilder::from(right_projected) + .project(visible_right_exprs.into_iter().chain(right_match_exprs)) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 3: build `UnionDistinctOn` with normalized internal match keys. + let visible_field_count = visible_schema.fields().len(); + let compare_key_indices = + (visible_field_count..visible_field_count + match_columns.len()).collect::>(); + let (time_qualifier, _) = visible_schema + .iter() + .find(|(_, field)| field.name() == &left_time_index_column) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let ts_col_idx = left_augmented + .schema() + .iter() + .position(|(qualifier, field)| { + qualifier == time_qualifier && field.name() == &left_time_index_column + }) + .with_context(|| TimeIndexNotFoundSnafu { + table: left_qualifier_string.clone(), + })?; + let union_distinct_on = UnionDistinctOn::try_new( + left_augmented, + right_augmented, + compare_key_indices, + ts_col_idx, + ) + .context(DataFusionPlanningSnafu)?; + let augmented_result = LogicalPlan::Extension(Extension { + node: Arc::new(union_distinct_on), + }); + let result = LogicalPlanBuilder::from(augmented_result) + .project(visible_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + + // step 4: update context + let output_field_col = left_field_col.clone(); + let mut output_context = left_context; + let mut visible_tags = all_tags.into_iter().collect::>(); + visible_tags.sort_unstable(); + output_context.time_index_column = Some(left_time_index_column); + output_context.tag_columns = visible_tags; + output_context.field_columns = if mixed_sample_types { + vec![mixed_float_field_col, mixed_histogram_field_col] + } else { + vec![output_field_col] + }; + output_context.use_tsid = left_has_tsid && right_has_tsid; + self.ctx = output_context; + + Ok(result) + } + + /// Build a set operator (AND/OR/UNLESS) + pub(super) fn set_op_on_non_field_columns( + &mut self, + mut left: LogicalPlan, + mut right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + op: TokenType, + modifier: &Option, + ) -> Result { + let left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + + if matches!(op.id(), token::T_LOR) { + return self.or_operator( + left, + right, + left_tag_col_set, + right_tag_col_set, + left_context, + right_context, + modifier, + ); + } + + if let Some(modifier) = modifier { + ensure!( + matches!( + modifier.card, + VectorMatchCardinality::OneToOne | VectorMatchCardinality::ManyToMany + ), + UnsupportedVectorMatchSnafu { + name: modifier.card.clone(), + }, + ); + } + + let output_context = left_context.clone(); + let visible_left_schema = left.schema().clone(); + let mut left_context = left_context; + let mut right_context = right_context; + let added_marker_to_left = if Self::only_temporality_match_label_mismatches( + &left_context, + &right_context, + modifier, + ) { + let aligned = Self::align_temporality_match_column( + left, + right, + &mut left_context, + &mut right_context, + )?; + left = aligned.0; + right = aligned.1; + aligned.2 + } else { + false + }; + + let mut left_tag_col_set = left_context + .tag_columns + .iter() + .cloned() + .collect::>(); + let mut right_tag_col_set = right_context + .tag_columns + .iter() + .cloned() + .collect::>(); + if let Some(matching) = modifier + .as_ref() + .and_then(|modifier| modifier.matching.as_ref()) + { + match matching { + LabelModifier::Include(on) => { + let mask = on.labels.iter().cloned().collect::>(); + left_tag_col_set = left_tag_col_set.intersection(&mask).cloned().collect(); + right_tag_col_set = right_tag_col_set.intersection(&mask).cloned().collect(); + } + LabelModifier::Exclude(ignoring) => { + for label in &ignoring.labels { + let _ = left_tag_col_set.remove(label); + let _ = right_tag_col_set.remove(label); + } + } + } + } + ensure!( + left_tag_col_set == right_tag_col_set, + CombineTableColumnMismatchSnafu { + left: left_tag_col_set.iter().cloned().collect::>(), + right: right_tag_col_set.iter().cloned().collect::>(), + } + ); + + let left_time_index = left_context.time_index_column.clone().unwrap(); + let right_time_index = right_context.time_index_column.clone().unwrap(); + + // alias right time index column if necessary + if left_context.time_index_column != right_context.time_index_column { + let right_project_exprs = right + .schema() + .fields() + .iter() + .map(|field| { + if field.name() == &right_time_index { + DfExpr::Column(Column::from_name(&right_time_index)).alias(&left_time_index) + } else { + DfExpr::Column(Column::from_name(field.name())) + } + }) + .collect::>(); + + right = LogicalPlanBuilder::from(right) + .project(right_project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + } + + let join_keys = left_tag_col_set + .into_iter() + .chain([left_time_index]) + .collect::>(); + + ensure!( + left_context.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + left.schema(), + &left_context.field_columns, + ), + MultiFieldsNotSupportedSnafu { + operator: "AND/UNLESS operator" + } + ); + // Generate join plan. + // All set operations in PromQL are "distinct" + let result = match op.id() { + token::T_LAND => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftSemi, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LUNLESS => LogicalPlanBuilder::from(left) + .distinct() + .context(DataFusionPlanningSnafu)? + .join_detailed( + right, + JoinType::LeftAnti, + (join_keys.clone(), join_keys), + None, + NullEquality::NullEqualsNull, + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu), + token::T_LOR => { + // OR is handled at the beginning of this function, as it cannot + // be expressed using JOIN like AND and UNLESS. + unreachable!() + } + _ => UnexpectedTokenSnafu { token: op }.fail(), + }?; + let result = if added_marker_to_left { + LogicalPlanBuilder::from(result) + .project(visible_left_schema.iter().map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + })) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + result + }; + + // AND/UNLESS preserve the complete left operand's visible columns and values; encoded + // markers are decoded. + self.ctx = output_context; + Ok(result) + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs new file mode 100644 index 00000000000..1634e6d9e1d --- /dev/null +++ b/src/query/src/promql/planner/test.rs @@ -0,0 +1,7585 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::time::{Duration, UNIX_EPOCH}; + +use catalog::RegisterTableRequest; +use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; +use common_base::Plugins; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_query::native_histogram::{ + CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, +}; +use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_query::test_util::DummyDecoder; +use common_recordbatch::RecordBatch as GreptimeRecordBatch; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::datatypes::{Field, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::{CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::{MemTable, provider_as_source}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::Extension; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use promql::extension_plan::HistogramFold; +use promql_parser::label::Labels; +use promql_parser::parser; +use session::context::QueryContext; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use table::Table; +use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; +use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; + +use super::*; +use crate::QueryEngineContext; +use crate::options::QueryOptions; +use crate::parser::QueryLanguageParser; +use crate::query_engine::DefaultSerializer; + +mod delta; + +fn find_instant_manipulate(plan: &LogicalPlan) -> Option<&InstantManipulate> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(instant_manipulate) = node.as_any().downcast_ref::() + { + return Some(instant_manipulate); + } + + plan.inputs().into_iter().find_map(find_instant_manipulate) +} + +fn build_query_engine_state() -> QueryEngineState { + QueryEngineState::new( + new_memory_catalog_manager().unwrap(), + None, + None, + None, + None, + None, + false, + Plugins::default(), + QueryOptions::default(), + ) +} + +#[test] +fn common_label_type_preserves_only_shared_dictionary_encoding() { + let dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Utf8), + ); + let other_dictionary = ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ); + + assert_eq!( + Some(dictionary.clone()), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&ArrowDataType::Utf8)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), Some(&other_dictionary)) + ); + assert_eq!( + Some(ArrowDataType::Utf8), + PromPlanner::common_label_data_type(Some(&dictionary), None) + ); +} + +async fn build_optimized_promql_plan( + table_provider: DfTableSourceProvider, + eval_stmt: &EvalStmt, +) -> LogicalPlan { + let state = build_query_engine_state(); + let raw_plan = PromPlanner::stmt_to_plan(table_provider, eval_stmt, &state) + .await + .unwrap(); + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + state + .optimize_by_extension_rules(raw_plan, &context) + .unwrap() +} + +async fn build_optimized_tsid_plan( + query: &str, + num_tag: usize, + num_field: usize, + end_secs: u64, + lookback_secs: u64, +) -> String { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(end_secs)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(lookback_secs), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + num_tag, + num_field, + ) + .await; + + build_optimized_promql_plan(table_provider, &eval_stmt) + .await + .display_indent_schema() + .to_string() +} + +async fn assert_nested_count_rewrite_applies(query: &str, expected_outer_agg: &str) { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("Projection: some_metric.timestamp, some_metric.tag_0")); + assert!(plan_str.contains("Distinct:")); + assert!(plan_str.contains(expected_outer_agg), "{plan_str}"); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); +} + +async fn assert_nested_count_rewrite_missing(query: &str, num_tag: usize, lookback_secs: u64) { + let plan_str = build_optimized_tsid_plan(query, num_tag, 1, 100_000, lookback_secs).await; + assert!(!plan_str.contains("Distinct:"), "{plan_str}"); +} + +fn build_eval_stmt(expr: &str) -> EvalStmt { + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } +} + +enum DirectOrValue { + Float64(f64), + Int64(i64), + NativeHistogram(NativeHistogram), + Utf8(&'static str), +} + +impl DirectOrValue { + fn data_type(&self) -> ArrowDataType { + match self { + Self::Float64(_) => ArrowDataType::Float64, + Self::Int64(_) => ArrowDataType::Int64, + Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), + Self::Utf8(_) => ArrowDataType::Utf8, + } + } + fn array(&self) -> Arc { + match self { + Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), + Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), + Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), + } + } +} + +fn direct_or_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + } +} + +fn operator_metric_table( + name: &str, + table_id: u32, + tag: &str, + le: Option<&str>, + value: DirectOrValue, +) -> table::TableRef { + let value_type = match &value { + DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), + DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), + DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), + DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), + }; + let tag_count = 1 + usize::from(le.is_some()); + let mut columns = vec![ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if le.is_some() { + columns.push(ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.extend([ + ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("v".to_string(), value_type, true), + ]); + let schema = Arc::new(Schema::new(columns)); + let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; + if let Some(le) = le { + arrays.push(Arc::new(StringArray::from(vec![le]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, + value.array(), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let value_index = tag_count + 1; + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![value_index]) + .next_column_id((value_index + 1) as u32) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(Table::new( + info, + FilterPushDownType::Unsupported, + backing.data_source(), + )) +} + +fn operator_table_provider() -> DfTableSourceProvider { + let catalog = MemoryCatalogManager::with_default_setup(); + let tables = [ + operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), + operator_metric_table( + "lh", + 2_002, + "b", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), + operator_metric_table( + "rh", + 2_004, + "a", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), + operator_metric_table( + "bad_classic", + 2_006, + "d", + Some("broken"), + DirectOrValue::Float64(1.0), + ), + operator_metric_table( + "bad_native", + 2_007, + "d", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + ]; + for table in tables { + let info = table.table_info(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table, + }) + .unwrap(); + } + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn operator_eval_stmt(expr: &str) -> EvalStmt { + let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: time, + end: time, + interval: Duration::from_secs(1), + lookback_delta: Duration::from_secs(5), + } +} + +struct DirectOrSource { + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +} + +fn source( + name: &'static str, + empty: bool, + timestamp: i64, + tags: Vec<(&'static str, Option<&'static str>)>, + value: DirectOrValue, +) -> DirectOrSource { + DirectOrSource { + name, + empty, + timestamp, + tags, + value, + } +} + +fn tagged_source( + name: &'static str, + empty: bool, + tag: (&'static str, Option<&'static str>), + value: DirectOrValue, +) -> DirectOrSource { + source(name, empty, 1, vec![("job", Some("job")), tag], value) +} + +fn job_source(name: &'static str, value: DirectOrValue) -> DirectOrSource { + source(name, true, 1, vec![("job", Some("job"))], value) +} + +fn table(source: &DirectOrSource) -> Arc { + let mut fields = vec![Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )]; + fields.extend( + source + .tags + .iter() + .map(|(name, _)| Field::new(*name, ArrowDataType::Utf8, true)), + ); + fields.push(Field::new("v", source.value.data_type(), true)); + let schema = Arc::new(ArrowSchema::new(fields)); + let partitions = if source.empty { + vec![vec![]] + } else { + let mut columns: Vec> = + vec![Arc::new(TimestampMillisecondArray::from(vec![ + source.timestamp, + ]))]; + columns.extend( + source + .tags + .iter() + .map(|(_, value)| Arc::new(StringArray::from(vec![*value])) as Arc), + ); + columns.push(source.value.array()); + vec![vec![RecordBatch::try_new(schema.clone(), columns).unwrap()]] + }; + Arc::new(MemTable::try_new(schema, partitions).unwrap()) +} + +fn scan(source: &DirectOrSource) -> LogicalPlan { + LogicalPlanBuilder::scan(source.name, provider_as_source(table(source)), None) + .unwrap() + .build() + .unwrap() +} + +fn direct_or_context(qualifier: &str, tags: &[&str], field: &str) -> PromPlannerContext { + PromPlannerContext { + table_name: Some(qualifier.to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![field.to_string()], + tag_columns: tags.iter().map(|tag| (*tag).to_string()).collect(), + ..Default::default() + } +} + +fn or_modifier(expr: &str) -> Option { + let PromExpr::Binary(expr) = parser::parse(expr).unwrap() else { + unreachable!() + }; + expr.modifier +} + +async fn plan_direct_or( + left: LogicalPlan, + right: LogicalPlan, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> LogicalPlan { + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + planner + .or_operator( + left, + right, + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + modifier, + ) + .unwrap() +} + +async fn execute(plan: LogicalPlan, state: &QueryEngineState) -> (LogicalPlan, Vec) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + (optimized, batches) +} + +async fn run( + left: &DirectOrSource, + right: &DirectOrSource, + left_context: PromPlannerContext, + right_context: PromPlannerContext, + modifier: &Option, +) -> (LogicalPlan, Vec) { + let plan = plan_direct_or( + scan(left), + scan(right), + left_context, + right_context, + modifier, + ) + .await; + execute(plan, &build_query_engine_state()).await +} + +async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { + let sample = |histogram: bool| { + if histogram { + DirectOrValue::NativeHistogram(direct_or_histogram()) + } else { + DirectOrValue::Float64(1.25) + } + }; + let left = tagged_source( + "lhs", + false, + ( + "k", + Some(if histogram_on_left { + "histogram" + } else { + "float" + }), + ), + sample(histogram_on_left), + ); + let right = tagged_source( + "rhs", + false, + ( + "k", + Some(if histogram_on_left { + "float" + } else { + "histogram" + }), + ), + sample(!histogram_on_left), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job", "k"], "v"); + let right_context = direct_or_context("rhs", &["job", "k"], "v"); + let plan = planner + .or_operator( + scan(&left), + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + (planner, plan) +} + +async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { + let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); + let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); + let row_count = histograms.len() + 1; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new("k", ArrowDataType::Utf8, false), + Field::new(&float_field, ArrowDataType::Float64, true), + Field::new( + &histogram_field, + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let mut histogram_values = Vec::with_capacity(row_count); + histogram_values.push(None); + histogram_values.extend(histograms.into_iter().map(Some)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|row| format!("kind_{row}")), + )), + Arc::new(Float64Array::from_iter( + (0..row_count).map(|row| (row == 0).then_some(1.25)), + )), + build_histogram_array(&histogram_values), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let planner = PromPlanner { + table_provider, + ctx: PromPlannerContext { + table_name: Some("mixed".to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![float_field, histogram_field], + tag_columns: vec!["k".to_string()], + ..Default::default() + }, + promql_annotations: None, + }; + (planner, plan) +} + +fn assert_no_internal_or_keys(schema: &DFSchema) { + assert!( + schema + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")), + "{schema:?}" + ); +} + +fn values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect() +} + +fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = datafusion::arrow::compute::cast( + batch.column_by_name(column).unwrap(), + &ArrowDataType::Float64, + ) + .unwrap(); + values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect::>() + }) + .collect() +} + +fn histograms(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).filter_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .collect() +} + +fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name("v") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = batch + .column_by_name("k") + .map(|column| column.as_any().downcast_ref::().unwrap()); + (0..batch.num_rows()).map(move |i| { + ( + values.value(i), + labels.and_then(|labels| { + (!labels.is_null(i)).then(|| labels.value(i).to_string()) + }), + ) + }) + }) + .collect::>(); + rows.sort_by(|left, right| left.0.total_cmp(&right.0)); + rows +} + +fn matrix_source( + name: &'static str, + k: Option>, + timestamp: i64, + value: f64, +) -> DirectOrSource { + let mut tags = vec![("job", Some("job"))]; + if let Some(k) = k { + tags.push(("k", k)); + } + source(name, false, timestamp, tags, DirectOrValue::Float64(value)) +} + +fn matrix_context(name: &str, k: Option>) -> PromPlannerContext { + direct_or_context( + name, + if k.is_some() { &["job", "k"] } else { &["job"] }, + "v", + ) +} + +async fn build_missing_le_or_normal_metric_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string()), + ], + &["pod", "instance"], + ) + .await +} + +fn assert_normal_metric_schema(plan: &LogicalPlan) { + let fields = plan.schema().fields(); + assert_eq!(fields.len(), 4, "{fields:?}"); + assert!( + fields.iter().any(|field| field.name() == "pod"), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| field.name() == "instance"), + "{fields:?}" + ); + assert!( + fields + .iter() + .any(|field| field.name() == greptime_timestamp()), + "{fields:?}" + ); + assert!( + fields.iter().any(|field| { + field.name() == greptime_value() && field.data_type() == &ArrowDataType::Float64 + }), + "{fields:?}" + ); +} + +async fn build_test_table_provider_with_distinct_tags( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn contains_histogram_fold(plan: &LogicalPlan) -> bool { + matches!(plan, LogicalPlan::Extension(Extension { node }) if node.as_any().is::()) + || plan.inputs().into_iter().any(contains_histogram_fold) +} + +async fn build_set_op_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("bucket_metric", &["job", "le"]), + ("normal_metric", &["job"]), + ("fallback_metric", &["instance"]), + ]) + .await +} + +async fn build_or_context_table_provider() -> DfTableSourceProvider { + build_test_table_provider_with_distinct_tags(&[ + ("normal_metric", &["job"]), + ("other_metric", &["instance"]), + ("non_hist_metric", &["instance"]), + ]) + .await +} + +async fn optimize_and_create_physical_plan( + state: &QueryEngineState, + plan: LogicalPlan, +) -> ( + LogicalPlan, + Arc, +) { + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(plan, &context).unwrap(); + let physical = state + .session_state() + .create_physical_plan(&optimized) + .await + .unwrap(); + (optimized, physical) +} + +async fn build_test_table_provider( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + for i in 0..num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_native_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + build_test_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_multi_histogram_table_provider(table_name: &str) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + "native_histogram_2".to_string(), + native_histogram_value_type().clone(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2, 3]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_mixed_native_histogram_table_provider( + table_name: &str, +) -> DfTableSourceProvider { + build_test_mixed_native_histogram_table_provider_with_marker(table_name, false).await +} + +async fn build_test_mixed_native_histogram_table_provider_with_marker( + table_name: &str, + temporality_marker: bool, +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if temporality_marker { + columns.push(ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + )); + } + let tag_count = columns.len(); + columns.extend([ + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![tag_count + 1, tag_count + 2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let mut arrays: Vec> = + vec![Arc::new(StringArray::from(vec!["float", "histogram"]))]; + if temporality_marker { + arrays.push(Arc::new(StringArray::from(vec![ + Some(GREPTIME_TEMPORALITY_DELTA), + Some(GREPTIME_TEMPORALITY_DELTA), + ]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])) as Arc, + build_histogram_array(&[None, Some(direct_or_histogram())]), + Arc::new(Float64Array::from(vec![Some(2.0), None])), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 1024, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +fn classic_and_native_histogram_table_provider( + native_tag: &str, + native_le: Option<&str>, + native_histogram: NativeHistogram, +) -> DfTableSourceProvider { + let table_name = "mixed_histogram"; + let catalog = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3, 4]) + .next_column_id(5) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![ + "classic", "classic", native_tag, "classic", "classic", native_tag, + ])), + Arc::new(StringArray::from(vec![ + Some("1"), + Some("+Inf"), + native_le, + Some("1"), + Some("+Inf"), + native_le, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, + ])), + build_histogram_array(&[ + None, + None, + Some(native_histogram.clone()), + None, + None, + Some(native_histogram), + ]), + Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(2.0), + Some(4.0), + None, + ])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 2_200, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 2_200, + table, + }) + .unwrap(); + + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_tsid( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, +) -> DfTableSourceProvider { + let table_specs = table_name_tuples + .iter() + .map(|(schema_name, table_name)| ((schema_name.clone(), table_name.clone()), num_field)) + .collect::>(); + build_test_table_provider_with_tsid_fields(&table_specs, num_tag).await +} + +async fn build_test_table_provider_with_tsid_fields( + table_specs: &[((String, String), usize)], + num_tag: usize, +) -> DfTableSourceProvider { + let table_specs = table_specs + .iter() + .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field)) + .collect::>(); + build_test_table_provider_with_tsid_tag_fields(&table_specs).await +} + +async fn build_test_table_provider_with_tsid_tag_fields( + table_specs: &[((String, String), usize, usize)], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + + let physical_table_name = "phy"; + let physical_table_id = 999u32; + let physical_num_tag = table_specs + .iter() + .map(|(_, num_tag, _)| *num_tag) + .max() + .unwrap_or(0); + let physical_num_field = table_specs + .iter() + .map(|(_, _, num_field)| *num_field) + .max() + .unwrap_or(0); + + // Register a metric engine physical table with internal columns. + { + let mut columns = vec![ + ColumnSchema::new( + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + ConcreteDataType::uint32_datatype(), + false, + ), + ColumnSchema::new( + DATA_SCHEMA_TSID_COLUMN_NAME.to_string(), + ConcreteDataType::uint64_datatype(), + false, + ), + ]; + for i in 0..physical_num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..physical_num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let primary_key_indices = (0..(2 + physical_num_tag)).collect::>(); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(primary_key_indices) + .value_indices( + (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(), + ) + .engine(METRIC_ENGINE_NAME.to_string()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(physical_table_id) + .name(physical_table_name) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: physical_table_name.to_string(), + table_id: physical_table_id, + table, + }) + .is_ok() + ); + } + + // Register metric engine logical tables without `__tsid`, referencing the physical table. + for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate() { + let mut columns = vec![]; + for i in 0..*num_tag { + columns.push(ColumnSchema::new( + format!("tag_{i}"), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + for i in 0..*num_field { + columns.push(ColumnSchema::new( + format!("field_{i}"), + ConcreteDataType::float64_datatype(), + true, + )); + } + + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + physical_table_name.to_string(), + ); + let table_id = 1024u32 + idx as u32; + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..*num_tag).collect()) + .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect()) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(table_id) + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_test_table_provider_with_fields( + table_name_tuples: &[(String, String)], + tags: &[&str], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (schema_name, table_name) in table_name_tuples { + let mut columns = vec![]; + let num_tag = tags.len(); + for tag in tags { + columns.push(ColumnSchema::new( + tag.to_string(), + ConcreteDataType::string_datatype(), + false, + )); + } + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.clone()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: schema_name.clone(), + table_name: table_name.clone(), + table_id: 1024, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +// { +// input: `abs(some_metric{foo!="bar"})`, +// expected: &Call{ +// Func: MustGetFunction("abs"), +// Args: Expressions{ +// &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchNotEqual, "foo", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// }, +// }, +// }, +async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) { + let prom_expr = parser::parse(&format!("{fn_name}(some_metric{{tag_0!=\"bar\"}})")).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: TEMPLATE(field_0) IS NOT NULL [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, TEMPLATE(some_metric.field_0) AS TEMPLATE(field_0), some_metric.tag_0 [timestamp:Timestamp(ms), TEMPLATE(field_0):Float64;N, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]" + ).replace("TEMPLATE", plan_name); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_abs() { + do_single_instant_function_call("abs", "abs").await; +} + +#[tokio::test] +#[should_panic] +async fn single_absent() { + do_single_instant_function_call("absent", "").await; +} + +#[tokio::test] +async fn single_ceil() { + do_single_instant_function_call("ceil", "ceil").await; +} + +#[tokio::test] +async fn single_exp() { + do_single_instant_function_call("exp", "exp").await; +} + +#[tokio::test] +async fn single_ln() { + do_single_instant_function_call("ln", "ln").await; +} + +#[tokio::test] +async fn single_log2() { + do_single_instant_function_call("log2", "log2").await; +} + +#[tokio::test] +async fn single_log10() { + do_single_instant_function_call("log10", "log10").await; +} + +#[tokio::test] +#[should_panic] +async fn single_scalar() { + do_single_instant_function_call("scalar", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sgn() { + do_single_instant_function_call("sgn", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort() { + do_single_instant_function_call("sort", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sort_desc() { + do_single_instant_function_call("sort_desc", "").await; +} + +#[tokio::test] +async fn single_sqrt() { + do_single_instant_function_call("sqrt", "sqrt").await; +} + +#[tokio::test] +async fn single_timestamp_plan_preserves_source_value() { + let eval_stmt = build_eval_stmt(r#"timestamp(some_metric{tag_0!="bar"})"#); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn single_acos() { + do_single_instant_function_call("acos", "acos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_acosh() { + do_single_instant_function_call("acosh", "").await; +} + +#[tokio::test] +async fn single_asin() { + do_single_instant_function_call("asin", "asin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_asinh() { + do_single_instant_function_call("asinh", "").await; +} + +#[tokio::test] +async fn single_atan() { + do_single_instant_function_call("atan", "atan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_atanh() { + do_single_instant_function_call("atanh", "").await; +} + +#[tokio::test] +async fn single_cos() { + do_single_instant_function_call("cos", "cos").await; +} + +#[tokio::test] +#[should_panic] +async fn single_cosh() { + do_single_instant_function_call("cosh", "").await; +} + +#[tokio::test] +async fn single_sin() { + do_single_instant_function_call("sin", "sin").await; +} + +#[tokio::test] +#[should_panic] +async fn single_sinh() { + do_single_instant_function_call("sinh", "").await; +} + +#[tokio::test] +async fn single_tan() { + do_single_instant_function_call("tan", "tan").await; +} + +#[tokio::test] +#[should_panic] +async fn single_tanh() { + do_single_instant_function_call("tanh", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_deg() { + do_single_instant_function_call("deg", "").await; +} + +#[tokio::test] +#[should_panic] +async fn single_rad() { + do_single_instant_function_call("rad", "").await; +} + +// { +// input: "avg by (foo)(some_metric)", +// expected: &AggregateExpr{ +// Op: AVG, +// Expr: &VectorSelector{ +// Name: "some_metric", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// PosRange: PositionRange{ +// Start: 13, +// End: 24, +// }, +// }, +// Grouping: []string{"foo"}, +// PosRange: PositionRange{ +// Start: 0, +// End: 25, +// }, +// }, +// }, +async fn do_aggregate_expr_plan(fn_name: &str, plan_name: &str) { + let prom_expr = parser::parse(&format!( + "{fn_name} by (tag_1)(some_metric{{tag_0!=\"bar\"}})", + )) + .unwrap(); + let mut eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // test group by + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_no_without = String::from( + "Sort: some_metric.tag_1 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_1, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_1:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!( + plan.display_indent_schema().to_string(), + expected_no_without + ); + + // test group without + if let PromExpr::Aggregate(AggregateExpr { modifier, .. }) = &mut eval_stmt.expr { + *modifier = Some(LabelModifier::Exclude(Labels { + labels: vec![String::from("tag_1")].into_iter().collect(), + })); + } + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 2, + 2, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected_without = String::from( + "Sort: some_metric.tag_0 ASC NULLS LAST, some_metric.timestamp ASC NULLS LAST [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n Aggregate: groupBy=[[some_metric.tag_0, some_metric.timestamp]], aggr=[[TEMPLATE(some_metric.field_0), TEMPLATE(some_metric.field_1)]] [tag_0:Utf8, timestamp:Timestamp(ms), TEMPLATE(some_metric.field_0):Float64;N, TEMPLATE(some_metric.field_1):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\"] [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.tag_1 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, tag_1:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N]" + ).replace("TEMPLATE", plan_name); + assert_eq!(plan.display_indent_schema().to_string(), expected_without); +} + +#[tokio::test] +async fn aggregate_sum() { + do_aggregate_expr_plan("sum", "sum").await; +} + +#[tokio::test] +async fn tsid_is_used_for_series_divide_when_available() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + assert!(plan_str.contains("__tsid ASC NULLS FIRST")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + let manipulate = find_instant_manipulate(&plan).unwrap(); + let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new( + &[], + Arc::new( + datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0] + .schema() + .as_arrow() + .clone(), + ), + None, + ) + .unwrap(), + )))); + assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); +} + +#[tokio::test] +async fn default_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + !plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn reject_binary_fill_modifiers() { + let state = build_query_engine_state(); + + for query in [ + "some_metric + fill(0) some_alt_metric", + "some_metric + fill_left(0) some_alt_metric", + "some_metric + fill_right(0) some_alt_metric", + "(some_metric + fill(0) some_alt_metric) + some_metric", + ] { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider(&[], 0, 0).await; + let err = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &state) + .await + .unwrap_err(); + + assert!( + matches!( + &err, + crate::promql::error::Error::UnsupportedExpr { name, .. } + if name == "PromQL fill modifiers" + ), + "{err}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_falls_back_when_tsid_is_projected_out() { + for query in [ + "timestamp(some_metric) / some_metric", + "some_metric / timestamp(some_metric)", + ] { + let eval_stmt = build_eval_stmt(query); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{query}: {plan_str}"); + assert!( + plan_str.contains("lhs.tag_0 = rhs.tag_0"), + "{query}: {plan_str}" + ); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn timestamp_binary_join_rejects_default_matching_on_mismatched_labels() { + let eval_stmt = build_eval_stmt("timestamp(left_host_job) / right_by_job"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "left_host_job".to_string()), + 2, + 1, + ), + ( + (DEFAULT_SCHEMA_NAME.to_string(), "right_by_job".to_string()), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + + assert!( + plan_str.contains("Boolean(false)") || plan_str.contains("false"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_is_preserved_for_nested_default_binary_joins() { + let eval_stmt = build_eval_stmt("(some_metric - some_alt_metric) / some_third_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_third_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_leaf_plan() { + let eval_stmt = build_eval_stmt("((some_metric - some_alt_metric) / some_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn repeated_tsid_binary_operand_reuses_shorter_field_side() { + let eval_stmt = + build_eval_stmt("((two_field_metric - one_field_metric) / one_field_metric) * 100"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1025)") + .count(), + 1, + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); +} + +#[tokio::test] +async fn binary_island_reuses_self_operand_without_join() { + let eval_stmt = build_eval_stmt("some_metric / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 0, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_reuses_leaf_across_two_branches() { + let eval_stmt = + build_eval_stmt("(some_metric + some_alt_metric) / (some_metric + third_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + (DEFAULT_SCHEMA_NAME.to_string(), "third_metric".to_string()), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str + .matches("Filter: phy.__table_id = UInt32(1024)") + .count(), + 1, + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_generated_alias_avoids_user_column_names() { + let eval_stmt = build_eval_stmt("(some_metric + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + &["prom_v0", "__prom_v0"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan.schema().field_names(); + assert!(field_names.iter().any(|name| name.ends_with(".prom_v0"))); + assert!(field_names.iter().any(|name| name.ends_with(".__prom_v0"))); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("SubqueryAlias: __prom_v0"), "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_clears_qualifier_for_nested_unary_projection() { + let eval_stmt = build_eval_stmt("-((some_metric + some_alt_metric) / some_metric)"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 1, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_distinct_matcher_leaves() { + let eval_stmt = build_eval_stmt( + "(some_metric{tag_0=\"foo\"} + some_alt_metric) / some_metric{tag_0=\"bar\"}", + ); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_keeps_offset_leaves_distinct() { + let eval_stmt = build_eval_stmt("(some_metric offset 5m + some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_group_modifier() { + let eval_stmt = + build_eval_stmt("(some_metric + ignoring(tag_0) group_left some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn binary_island_falls_back_for_comparison_filter() { + let eval_stmt = build_eval_stmt("(some_metric > some_alt_metric) / some_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!(plan_str.matches("__tsid =").count(), 2, "{plan_str}"); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 3, + "{plan_str}" + ); +} + +#[tokio::test] +async fn tsid_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("one_field_metric / two_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + let value_columns = field_names + .iter() + .filter(|name| { + *name != "tag_0" && *name != "timestamp" && *name != DATA_SCHEMA_TSID_COLUMN_NAME + }) + .count(); + assert_eq!(value_columns, 1, "{field_names:?}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_shorter_field_side() { + let eval_stmt = build_eval_stmt("two_field_metric > one_field_metric"); + + let table_provider = build_test_table_provider_with_tsid_fields( + &[ + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "two_field_metric".to_string(), + ), + 2, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "one_field_metric".to_string(), + ), + 1, + ), + ], + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let field_names = plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + assert!( + field_names.iter().any(|name| name == "field_0"), + "{field_names:?}" + ); + assert!( + !field_names.iter().any(|name| name == "field_1"), + "{field_names:?}" + ); +} + +#[tokio::test] +async fn label_matching_modifier_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_1 = some_alt_metric.tag_1"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn ignoring_absent_label_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn range_function_keeps_tsid_for_absent_ignoring_binary_join() { + let eval_stmt = build_eval_stmt("rate(some_metric[5m]) / ignoring(missing) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_full_label_set_keeps_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0, tag_1) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_partial_label_set_disables_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() { + let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid_tag_fields(&[ + ( + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + 2, + 1, + ), + ( + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + 1, + 1, + ), + ]) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("__tsid ="), "{plan_str}"); + assert!( + plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { + let eval_stmt = build_eval_stmt("some_metric > some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), + promql_annotations: None, + }; + let plan = planner + .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME), + "{plan_str}" + ); + assert!(planner.ctx.use_tsid, "{plan_str}"); +} + +#[tokio::test] +async fn comparison_bool_binary_join_uses_tsid_when_available() { + let eval_stmt = build_eval_stmt("some_metric > bool some_alt_metric"); + + let table_provider = build_test_table_provider_with_tsid( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "some_alt_metric".to_string(), + ), + ], + 2, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("some_metric.__tsid = some_alt_metric.__tsid"), + "{plan_str}" + ); + assert!(!plan_str.contains("tag_0 ="), "{plan_str}"); + assert!(!plan_str.contains("tag_1 ="), "{plan_str}"); +} + +#[tokio::test] +async fn scalar_count_count_range_keeps_full_window() { + let plan_str = build_optimized_tsid_plan( + "scalar(count(count(some_metric) by (tag_0)))", + 1, + 1, + 100_000, + 1, + ) + .await; + assert!(plan_str.contains("ScalarCalculate: tags=[]")); + assert!(plan_str.contains("PromInstantManipulate: range=[0..100000000]")); + assert!(!plan_str.contains("PromInstantManipulate: range=[99999000..99999000]")); +} + +#[tokio::test] +async fn scalar_count_count_rewrite_applies_inside_binary_expr_for_tsid_input() { + let plan_str = build_optimized_tsid_plan( + "sum(irate(some_metric[1h])) / scalar(count(count(some_metric) by (tag_0)))", + 2, + 1, + 10, + 300, + ) + .await; + assert!(plan_str.contains("Distinct:"), "{plan_str}"); +} + +#[tokio::test] +async fn nested_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(count(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(count(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_sum_count_rewrite_keeps_full_series_key_with_tsid_input() { + assert_nested_count_rewrite_applies( + "count(sum(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(sum(some_metric.field_0))]]" + ) + .await; +} + +#[tokio::test] +async fn nested_supported_inner_aggs_rewrite_apply_for_tsid_input() { + for (query, expected_outer_agg) in [ + ( + "count(avg(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(avg(some_metric.field_0))]]", + ), + ( + "count(min(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(min(some_metric.field_0))]]", + ), + ( + "count(max(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(max(some_metric.field_0))]]", + ), + ( + "count(stddev(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(stddev_pop(some_metric.field_0))]]", + ), + ( + "count(stdvar(some_metric) by (tag_0))", + "Aggregate: groupBy=[[some_metric.timestamp]], aggr=[[count(Int64(1)) AS count(var_pop(some_metric.field_0))]]", + ), + ] { + assert_nested_count_rewrite_applies(query, expected_outer_agg).await; + } +} + +#[tokio::test] +async fn nested_non_count_inner_aggs_rewrite_filter_null_values_for_tsid_input() { + let count_plan = + build_optimized_tsid_plan("count(count(some_metric) by (tag_0))", 2, 1, 100_000, 1).await; + assert!( + !count_plan.contains("some_metric.field_0 IS NOT NULL"), + "{count_plan}" + ); + + for query in [ + "count(sum(some_metric) by (tag_0))", + "count(avg(some_metric) by (tag_0))", + "count(min(some_metric) by (tag_0))", + "count(max(some_metric) by (tag_0))", + "count(stddev(some_metric) by (tag_0))", + "count(stdvar(some_metric) by (tag_0))", + ] { + let plan_str = build_optimized_tsid_plan(query, 2, 1, 100_000, 1).await; + assert!( + plan_str.contains("Filter: some_metric.field_0 IS NOT NULL"), + "{query}: {plan_str}" + ); + } +} + +#[tokio::test] +async fn nested_unsupported_or_non_direct_inner_aggs_do_not_rewrite() { + assert_nested_count_rewrite_missing("count(group(some_metric) by (tag_0))", 2, 1).await; + assert_nested_count_rewrite_missing("count(sum(irate(some_metric[1h])) by (tag_0))", 2, 300) + .await; +} + +#[tokio::test] +async fn physical_table_name_is_not_leaked_in_plan() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("TableScan: phy"), "{plan}"); + assert!(plan_str.contains("SubqueryAlias: some_metric")); + assert!(plan_str.contains("Filter: phy.__table_id = UInt32(1024)")); + assert!(!plan_str.contains("TableScan: some_metric")); +} + +#[tokio::test] +async fn sum_without_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum without (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn topk_without_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk without (tag_0) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn sum_by_does_not_group_by_tsid() { + let prom_expr = parser::parse("sum by (__tsid) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(!aggr_line.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn aggregate_over_binary_time_function_expr() { + for op in ["sum", "min", "max", "avg"] { + let prom_expr = parser::parse(&format!( + "{op} by (tag_0, tag_1, tag_2) (time() - some_metric)" + )) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 1, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + let aggr_line = plan_str + .lines() + .find(|line| line.contains("Aggregate: groupBy=")) + .unwrap(); + assert!(aggr_line.contains(op), "{plan_str}"); + assert!(aggr_line.contains("first_value"), "{plan_str}"); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| { field.name() == DATA_SCHEMA_TSID_COLUMN_NAME }) + ); + } +} + +#[tokio::test] +async fn topk_by_does_not_partition_by_tsid() { + let prom_expr = parser::parse("topk by (__tsid) (1, some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); + + let window_line = plan_str + .lines() + .find(|line| line.contains("WindowAggr: windowExpr=[[row_number()")) + .unwrap(); + let partition_by = window_line + .split("PARTITION BY [") + .nth(1) + .and_then(|s| s.split("] ORDER BY").next()) + .unwrap(); + assert!(!partition_by.contains(DATA_SCHEMA_TSID_COLUMN_NAME)); +} + +#[tokio::test] +async fn selector_matcher_on_tsid_does_not_use_internal_column() { + let prom_expr = parser::parse(r#"some_metric{__tsid="123"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + fn collect_filter_cols(plan: &LogicalPlan, out: &mut HashSet) { + if let LogicalPlan::Filter(filter) = plan { + datafusion_expr::utils::expr_to_columns(&filter.predicate, out).unwrap(); + } + for input in plan.inputs() { + collect_filter_cols(input, out); + } + } + + let mut filter_cols = HashSet::new(); + collect_filter_cols(&plan, &mut filter_cols); + assert!( + !filter_cols + .iter() + .any(|c| c.name == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn tsid_is_not_used_when_physical_table_is_missing() { + let prom_expr = parser::parse("some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let catalog_list = MemoryCatalogManager::with_default_setup(); + + // Register a metric engine logical table referencing a missing physical table. + let mut columns = vec![ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + columns.push( + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + let schema = Arc::new(Schema::new(columns)); + let mut options = table::requests::TableOptions::default(); + options + .extra_options + .insert(LOGICAL_TABLE_METADATA_KEY.to_string(), "phy".to_string()); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .engine(METRIC_ENGINE_NAME.to_string()) + .options(options) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .table_id(1024) + .name("some_metric") + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "some_metric".to_string(), + table_id: 1024, + table, + }) + .unwrap(); + + let table_provider = DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("PromSeriesDivide: tags=[\"tag_0\"]")); + assert!(!plan_str.contains("PromSeriesDivide: tags=[\"__tsid\"]")); +} + +#[tokio::test] +async fn tsid_is_carried_only_when_aggregate_preserves_label_set() { + let prom_expr = parser::parse("sum by (tag_0) (some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("first_value") && plan_str.contains("__tsid")); + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); + + // Merging aggregate: label set is reduced, tsid should not be carried. + let prom_expr = parser::parse("sum(some_metric)").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.contains("first_value")); +} + +#[tokio::test] +async fn or_operator_with_unknown_metric_does_not_require_tsid() { + let prom_expr = parser::parse("unknown_metric or some_metric").unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + assert!( + !plan + .schema() + .fields() + .iter() + .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME) + ); +} + +#[tokio::test] +async fn aggregate_avg() { + do_aggregate_expr_plan("avg", "avg").await; +} + +#[tokio::test] +#[should_panic] // output type doesn't match +async fn aggregate_count() { + do_aggregate_expr_plan("count", "count").await; +} + +#[tokio::test] +async fn aggregate_min() { + do_aggregate_expr_plan("min", "min").await; +} + +#[tokio::test] +async fn aggregate_max() { + do_aggregate_expr_plan("max", "max").await; +} + +#[tokio::test] +async fn aggregate_group() { + // Regression test for `group()` aggregator. + // PromQL: sum(group by (cluster)(kubernetes_build_info{service="kubernetes",job="apiserver"})) + // should be plannable, and `group()` should produce constant 1 for each group. + let prom_expr = parser::parse( + "sum(group by (cluster)(kubernetes_build_info{service=\"kubernetes\",job=\"apiserver\"}))", + ) + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "kubernetes_build_info".to_string(), + )], + &["cluster", "service", "job"], + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("max(Float64(1")); +} + +#[tokio::test] +async fn aggregate_stddev() { + do_aggregate_expr_plan("stddev", "stddev_pop").await; +} + +#[tokio::test] +async fn aggregate_stdvar() { + do_aggregate_expr_plan("stdvar", "var_pop").await; +} + +// TODO(ruihang): add range fn tests once exprs are ready. + +// { +// input: "some_metric{tag_0="foo"} + some_metric{tag_0="bar"}", +// expected: &BinaryExpr{ +// Op: ADD, +// LHS: &VectorSelector{ +// Name: "a", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchEqual, "tag_0", "foo"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// RHS: &VectorSelector{ +// Name: "sum", +// LabelMatchers: []*labels.Matcher{ +// MustLabelMatcher(labels.MatchxEqual, "tag_0", "bar"), +// MustLabelMatcher(labels.MatchEqual, model.MetricNameLabel, "some_metric"), +// }, +// }, +// VectorMatching: &VectorMatching{}, +// }, +// }, +#[tokio::test] +async fn binary_op_column_column() { + let prom_expr = + parser::parse(r#"some_metric{tag_0="foo"} + some_metric{tag_0="bar"}"#).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = String::from( + "Projection: rhs.tag_0, rhs.timestamp, CAST(lhs.field_0 AS Float64) + CAST(rhs.field_0 AS Float64) AS lhs.field_0 + rhs.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), lhs.field_0 + rhs.field_0:Float64;N]\ + \n Inner Join: lhs.tag_0 = rhs.tag_0, lhs.timestamp = rhs.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: lhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"foo\") AND some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: rhs [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.tag_0 = Utf8(\"foo\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +async fn indie_query_plan_compare>(query: &str, expected: T) { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); +} + +#[tokio::test] +async fn binary_op_literal_column() { + let query = r#"1 + some_metric{tag_0="bar"}"#; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, Float64(1) + CAST(some_metric.field_0 AS Float64) AS Float64(1) + field_0 [tag_0:Utf8, timestamp:Timestamp(ms), Float64(1) + field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 = Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn binary_op_literal_literal() { + let query = r#"1 + 1"#; + let expected = r#"EmptyMetric: range=[0..100000000], interval=[5000] [time:Timestamp(ms), value:Float64;N] + TableScan: dummy [time:Timestamp(ms), value:Float64;N]"#; + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_bool_grammar() { + let query = "some_metric != bool 1.2345"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 != Float64(1.2345) AS Float64) AS field_0 != Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 != Float64(1.2345):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn bool_with_additional_arithmetic() { + let query = "some_metric + (1 == bool 2)"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(some_metric.field_0 AS Float64) + CAST(Float64(1) = Float64(2) AS Float64) AS field_0 + Float64(1) = Float64(2) [tag_0:Utf8, timestamp:Timestamp(ms), field_0 + Float64(1) = Float64(2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn simple_unary() { + let query = "-some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, (- some_metric.field_0) AS (- field_0) [tag_0:Utf8, timestamp:Timestamp(ms), (- field_0):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn increase_aggr() { + let query = "increase(some_metric[5m])"; + let expected = String::from( + "Filter: prom_increase(timestamp_range,field_0,timestamp,Int64(300000)) IS NOT NULL [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_increase(timestamp_range, field_0, some_metric.timestamp, Int64(300000)) AS prom_increase(timestamp_range,field_0,timestamp,Int64(300000)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_increase(timestamp_range,field_0,timestamp,Int64(300000)):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +async fn native_histogram_plan(query: &str) -> String { + let table_provider = build_test_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent_schema().to_string() +} + +#[tokio::test] +async fn native_histogram_count_uses_native_udf() { + let plan = native_histogram_plan("histogram_count(some_metric)").await; + + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); +} + +#[tokio::test] +async fn timestamp_filters_native_histogram_stale_marker_before_projection() { + let mut stale = direct_or_histogram(); + stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let table = operator_metric_table( + "stale_histogram", + 2_100, + "a", + None, + DirectOrValue::NativeHistogram(stale), + ); + let catalog = MemoryCatalogManager::with_default_setup(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "stale_histogram".to_string(), + table_id: 2_100, + table, + }) + .unwrap(); + let provider = DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + provider, + &operator_eval_stmt("timestamp(stale_histogram)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { + let histograms = build_histogram_array(&[None]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new( + greptime_native_histogram(), + histograms.data_type().clone(), + true, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + histograms, + Arc::new(Float64Array::from(vec![f64::from_bits( + PROMETHEUS_STALE_NAN_BITS, + )])), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + input, + )), + }); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 1_000, + 1_000, + 5_000, + 1_000, + 0, + "timestamp".to_string(), + Vec::new(), + Some(greptime_native_histogram().to_string()), + input, + )), + }); + // Match timestamp()'s parent projection, which otherwise prunes the companion lane. + let plan = LogicalPlanBuilder::from(input) + .project([col("timestamp")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); +} + +#[tokio::test] +async fn native_histogram_rate_can_feed_count() { + let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; + + assert!(plan.contains("prom_native_histogram_rate"), "{plan}"); + assert!(plan.contains("prom_native_histogram_count"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_quantile_skips_classic_fold() { + let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); + assert!(plan.contains("some_metric.le"), "{plan}"); + // The phi literal is threaded into the native quantile UDF as its second argument. + assert!(plan.contains("Float64(0.9)"), "{plan}"); + // The empty-values filter drops NULL quantile results so the output is empty + // when all native histogram samples are dropped. + assert!(plan.contains("IS NOT NULL"), "{plan}"); +} + +#[tokio::test] +async fn mixed_native_histogram_quantile_uses_histogram_field() { + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + assert!( + plan.contains("prom_native_histogram_quantile(greptime_native_histogram"), + "{plan}" + ); + assert!(!plan.contains("EmptyRelation"), "{plan}"); +} + +#[tokio::test] +async fn mixed_histogram_helpers_execute_classic_and_native_samples() { + let state = build_query_engine_state(); + for (query, expected) in [ + ( + "histogram_quantile(0.5, mixed_histogram)", + vec![("classic", 1.0), ("native", 0.0)], + ), + ( + "histogram_fraction(-Inf, +Inf, mixed_histogram)", + vec![("classic", 1.0), ("native", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); + assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("tag") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) + }) + .collect::>(); + actual.sort_by_key(|(tag, _)| *tag); + assert_eq!(actual, expected, "{query}"); + } +} + +#[tokio::test] +async fn mixed_histogram_helpers_report_annotations() { + let state = build_query_engine_state(); + let mut native_histogram = direct_or_histogram(); + native_histogram.count = 2.0; + native_histogram.sum = f64::NAN; + for (native_tag, expected_rows, expected_warnings, expected_infos) in [ + ( + "classic", + 0, + vec!["vector contains a mix of classic and native histograms"], + vec![], + ), + ( + "native", + 2, + vec![], + vec!["input to histogram_quantile has NaN observations, result is skewed higher"], + ), + ] { + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + classic_and_native_histogram_table_provider(native_tag, None, native_histogram.clone()), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!(warnings, expected_warnings); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { + let state = build_query_engine_state(); + let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); + stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "classic", + Some("native"), + direct_or_histogram(), + ), + &stmt, + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert_eq!( + plan_text.matches("TableScan: mixed_histogram").count(), + 1, + "{plan_text}" + ); + + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let le = batch + .column_by_name(LE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| { + ( + timestamps.value(row), + (!le.is_null(row)).then(|| le.value(row).to_string()), + values.value(row), + ) + }) + }) + .collect::>(); + actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); + assert_eq!( + actual, + vec![ + (1_000, None, 1.0), + (1_000, Some("native".to_string()), 0.0), + (2_000, None, 1.0), + (2_000, Some("native".to_string()), 0.0), + ] + ); +} + +#[tokio::test] +async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { + let state = build_query_engine_state(); + for native_le in [None, Some("native")] { + for query in [ + "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", + "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "native", + native_le, + direct_or_histogram(), + ), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "native_le={native_le:?}, query={query}" + ); + } + } +} + +#[tokio::test] +async fn native_histogram_quantile_rejects_multi_field_input() { + let table_provider = build_test_multi_histogram_table_provider("some_metric").await; + let result = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_quantile(0.9, some_metric)"), + &build_query_engine_state(), + ) + .await; + + let err = result.expect_err("histogram_quantile on two native histogram fields must fail"); + assert!( + err.to_string() + .contains("Multi fields calculation is not supported in histogram_quantile"), + "{err}" + ); +} + +#[tokio::test] +async fn native_histogram_topk_uses_drop_udf() { + let plan = native_histogram_plan("topk(1, some_metric)").await; + + assert!(plan.contains("prom_native_histogram_drop_float"), "{plan}"); + assert!( + plan.contains("Filter: prom_native_histogram_drop_float") && plan.contains("IS NOT NULL"), + "{plan}" + ); +} + +#[tokio::test] +async fn mixed_or_topk_bottomk_ignore_native_histograms() { + for op in ["topk", "bottomk"] { + let collector = PromqlAnnotationCollector::default(); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), + "{plan:?}" + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!( + infos, + vec![format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + )] + ); + } +} + +#[tokio::test] +async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { + let plan = native_histogram_plan("scalar(some_metric)").await; + + assert!(plan.contains("ScalarCalculate"), "{plan}"); + assert!(plan.contains("Filter: Boolean(false)"), "{plan}"); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); +} + +#[tokio::test] +async fn native_histogram_value_sort_is_empty_but_label_sort_preserves_samples() { + for function in ["sort", "sort_desc"] { + let plan = native_histogram_plan(&format!("{function}(some_metric)")).await; + + assert!(plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + assert!( + !plan.contains(&format!("Sort: {}", greptime_native_histogram())), + "{plan}" + ); + assert!(!plan.contains("prom_native_histogram_drop"), "{plan}"); + } + + for (function, direction) in [("sort_by_label", "ASC"), ("sort_by_label_desc", "DESC")] { + let plan = native_histogram_plan(&format!("{function}(some_metric, \"tag_0\")")).await; + + assert!(plan.contains(&format!("tag_0 {direction}")), "{plan}"); + assert!(plan.contains(greptime_native_histogram()), "{plan}"); + assert!(!plan.contains("Float64(NULL) IS NOT NULL"), "{plan}"); + } +} + +#[tokio::test] +async fn unsupported_native_histogram_functions_use_drop_udf() { + for query in [ + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "round(some_metric)", + "abs(some_metric)", + ] { + let plan = native_histogram_plan(query).await; + + assert!( + plan.contains("prom_native_histogram_drop_float"), + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn native_histogram_absent_over_time_uses_native_udf() { + let plan = native_histogram_plan("absent_over_time(some_metric[5m])").await; + + assert!( + plan.contains("prom_native_histogram_absent_over_time"), + "{plan}" + ); +} + +#[tokio::test] +async fn native_histogram_all_function_arms_route_correctly() { + // Every native-histogram match arm in `create_function_expr` must route to the + // expected UDF when all field columns are native histograms. `holt_winters` shares + // the `double_exponential_smoothing` arm but is not registered in the promql + // parser (0.10), so it cannot be exercised through a query string. + let cases = [ + // Range functions routed to native histogram UDFs. + ( + "increase(some_metric[5m])", + "prom_native_histogram_increase", + ), + ("rate(some_metric[5m])", "prom_native_histogram_rate"), + ("delta(some_metric[5m])", "prom_native_histogram_delta"), + ("idelta(some_metric[5m])", "prom_native_histogram_idelta"), + ("irate(some_metric[5m])", "prom_native_histogram_irate"), + ("resets(some_metric[5m])", "prom_native_histogram_resets"), + ("changes(some_metric[5m])", "prom_native_histogram_changes"), + ( + "avg_over_time(some_metric[5m])", + "prom_native_histogram_avg_over_time", + ), + ( + "sum_over_time(some_metric[5m])", + "prom_native_histogram_sum_over_time", + ), + ( + "count_over_time(some_metric[5m])", + "prom_native_histogram_count_over_time", + ), + ( + "last_over_time(some_metric[5m])", + "prom_native_histogram_last_over_time", + ), + ( + "present_over_time(some_metric[5m])", + "prom_native_histogram_present_over_time", + ), + // Unsupported functions dropped with the float-null UDF. + ("deriv(some_metric[5m])", "prom_native_histogram_drop_float"), + ( + "min_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "max_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stddev_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "stdvar_over_time(some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "quantile_over_time(0.9, some_metric[5m])", + "prom_native_histogram_drop_float", + ), + ( + "predict_linear(some_metric[5m], 60)", + "prom_native_histogram_drop_float", + ), + ( + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + "prom_native_histogram_drop_float", + ), + ("round(some_metric)", "prom_native_histogram_drop_float"), + ("rad(some_metric)", "prom_native_histogram_drop_float"), + ("deg(some_metric)", "prom_native_histogram_drop_float"), + ("sgn(some_metric)", "prom_native_histogram_drop_float"), + // Instant helper functions routed to native histogram UDFs. + ( + "histogram_count(some_metric)", + "prom_native_histogram_count", + ), + ("histogram_sum(some_metric)", "prom_native_histogram_sum"), + ("histogram_avg(some_metric)", "prom_native_histogram_avg"), + ( + "histogram_stddev(some_metric)", + "prom_native_histogram_stddev", + ), + ( + "histogram_stdvar(some_metric)", + "prom_native_histogram_stdvar", + ), + ( + "histogram_fraction(-2 + 1, 2 / 2, some_metric)", + "prom_native_histogram_fraction", + ), + ]; + + for (query, expected_udf) in cases { + let plan = native_histogram_plan(query).await; + assert!(plan.contains(expected_udf), "{query}\n{plan}"); + if query.starts_with("histogram_fraction") { + assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); + } + } +} + +#[tokio::test] +async fn mixed_native_histogram_ranges_use_coordinated_udfs() { + let dual_output = [ + "increase(some_metric[5m])", + "rate(some_metric[5m])", + "delta(some_metric[5m])", + "idelta(some_metric[5m])", + "irate(some_metric[5m])", + "avg_over_time(some_metric[5m])", + "sum_over_time(some_metric[5m])", + "last_over_time(some_metric[5m])", + ]; + let float_output = [ + "resets(some_metric[5m])", + "changes(some_metric[5m])", + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "max_over_time(some_metric[5m])", + "count_over_time(some_metric[5m])", + "absent_over_time(some_metric[5m])", + "present_over_time(some_metric[5m])", + "stddev_over_time(some_metric[5m])", + "stdvar_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + ]; + + for query in dual_output.iter().chain(float_output.iter()) { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); + assert_eq!( + plan.contains("prom_mixed_range_histogram"), + dual_output.contains(query), + "{query}\n{plan}" + ); + } + + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + assert_eq!(plan, expected); +} + +#[tokio::test] +async fn mixed_native_histogram_rate_executes_real_ranges() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + Field::new( + greptime_native_histogram(), + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), + build_histogram_array(&[None, Some(direct_or_histogram()), None]), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let collector = PromqlAnnotationCollector::default(); + let mut planner = PromPlanner { + table_provider: build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await, + ctx: PromPlannerContext { + start: 3000, + end: 3000, + interval: 1000, + range: Some(3000), + time_index_column: Some("timestamp".to_string()), + field_columns: vec![ + greptime_native_histogram().to_string(), + greptime_value().to_string(), + ], + ..Default::default() + }, + promql_annotations: Some(collector.clone()), + }; + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 3000, + 3000, + 1000, + 0, + 3000, + "timestamp".to_string(), + planner.ctx.field_columns.clone(), + input, + ) + .unwrap(), + ), + }); + let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { + unreachable!() + }; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = Vec::new(); + collector.append_to(&mut warnings, &mut Vec::new()); + assert!( + warnings + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); +} + +#[tokio::test] +async fn native_histogram_mixed_field_table_behaves() { + // Exercise function planning after float and histogram samples have already been + // represented as alternative nullable fields. Histogram functions must select the + // histogram field without adding a NULL float field that would reject every row. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("histogram_count(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("prom_native_histogram_count"), + "{plan_str}" + ); + assert!(!plan_str.contains("Float64(NULL)"), "{plan_str}"); + assert!( + plan_str.contains("prom_native_histogram_count(greptime_native_histogram) IS NOT NULL"), + "{plan_str}" + ); + + // Value sorting keeps the float column and never sorts by the histogram column. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("sort(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("greptime_value ASC NULLS FIRST"), + "{plan_str}" + ); + assert!( + !plan_str.contains("greptime_native_histogram ASC"), + "{plan_str}" + ); + + // scalar() ignores histogram samples and evaluates only the float field. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt("scalar(some_metric)"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.contains("ScalarCalculate"), "{plan_str}"); + assert!( + plan_str.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); + + // Functions that preserve both alternative fields keep rows with either sample type. + let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(r#"label_replace(some_metric, "copied", "$1", "tag_0", "(.*)")"#), + &build_query_engine_state(), + ) + .await + .unwrap(); + let plan_str = plan.display_indent_schema().to_string(); + let filter = plan_str.lines().next().unwrap(); + assert!( + filter.starts_with("Filter: ") + && filter.contains("greptime_native_histogram IS NOT NULL") + && filter.contains(" OR ") + && filter.contains("greptime_value IS NOT NULL"), + "{plan_str}" + ); +} + +#[tokio::test] +async fn less_filter_on_value() { + let query = "some_metric < 1.2345"; + let expected = String::from( + "Filter: some_metric.field_0 < Float64(1.2345) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn count_over_time() { + let query = "count_over_time(some_metric[5m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +/// The outer `PromRangeManipulate` from a subquery must be preceded by +/// `Sort` + `PromSeriesDivide`. +#[tokio::test] +async fn count_over_time_subquery() { + let query = "count_over_time(some_metric[10m:1m])"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[-540000..100000000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-540999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn test_hash_join() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"http_server_requests_seconds_sum{uri="/accounts/login"} / ignoring(kubernetes_pod_name,kubernetes_namespace) http_server_requests_seconds_count{uri="/accounts/login"}"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_sum".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["uri", "kubernetes_namespace", "kubernetes_pod_name"], + ) + .await; + // Should be ok + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_count.greptime_timestamp, CAST(http_server_requests_seconds_sum.greptime_value AS Float64) / CAST(http_server_requests_seconds_count.greptime_value AS Float64) AS http_server_requests_seconds_sum.greptime_value / http_server_requests_seconds_count.greptime_value\ + \n Projection: http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.kubernetes_namespace, http_server_requests_seconds_sum.kubernetes_pod_name, http_server_requests_seconds_sum.greptime_timestamp, http_server_requests_seconds_sum.greptime_value, http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_sum.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_sum.uri, http_server_requests_seconds_sum.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n Inner Join: http_server_requests_seconds_sum.greptime_timestamp = http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_sum.uri = http_server_requests_seconds_count.uri\ + \n SubqueryAlias: http_server_requests_seconds_sum\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_sum.uri ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_sum.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_sum.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_sum.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_sum.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_sum.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_sum\ + \n SubqueryAlias: http_server_requests_seconds_count\ + \n Projection: http_server_requests_seconds_count.uri, http_server_requests_seconds_count.kubernetes_namespace, http_server_requests_seconds_count.kubernetes_pod_name, http_server_requests_seconds_count.greptime_timestamp, http_server_requests_seconds_count.greptime_value\ + \n Filter: prom_assert_unique_match_group(__promql_match_group_count, http_server_requests_seconds_count.uri)\ + \n WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [http_server_requests_seconds_count.uri, http_server_requests_seconds_count.greptime_timestamp] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS __promql_match_group_count]]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp]\ + \n PromSeriesDivide: tags=[\"uri\", \"kubernetes_namespace\", \"kubernetes_pod_name\"]\ + \n Sort: http_server_requests_seconds_count.uri ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_namespace ASC NULLS FIRST, http_server_requests_seconds_count.kubernetes_pod_name ASC NULLS FIRST, http_server_requests_seconds_count.greptime_timestamp ASC NULLS FIRST\ + \n Filter: http_server_requests_seconds_count.uri = Utf8(\"/accounts/login\") AND http_server_requests_seconds_count.greptime_timestamp >= TimestampMillisecond(-999, None) AND http_server_requests_seconds_count.greptime_timestamp <= TimestampMillisecond(100000000, None)\ + \n TableScan: http_server_requests_seconds_count"; + assert_eq!(plan.to_string(), expected); +} + +#[tokio::test] +async fn test_nested_histogram_quantile() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"label_replace(histogram_quantile(0.99, sum by(pod, le, path, code) (rate(greptime_servers_grpc_requests_elapsed_bucket{container="frontend"}[1m0s]))), "pod_new", "$1", "pod", "greptimedb-frontend-[0-9a-z]*-(.*)")"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "greptime_servers_grpc_requests_elapsed_bucket".to_string(), + )], + &["pod", "le", "path", "code", "container"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_histogram_quantile_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Arithmetic applied to a histogram_quantile() result. Regression for #8144: + // HistogramFold used to drop the input column qualifiers, so the binary-op + // projection failed to resolve the qualified tag column. + let case = r#"histogram_quantile(0.5, sum by (le, pod) (rate(http_request_duration_seconds_bucket[5m]))) + 0"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "http_request_duration_seconds_bucket".to_string(), + )], + &["pod", "le"], + ) + .await; + // Should plan without a "No field named ..." error. + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_and_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) and (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + r#"count (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} ) unless (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_used_bytes{namespace=~".+"} )) / (max by (persistentvolumeclaim,namespace) (kubelet_volume_stats_capacity_bytes{namespace=~".+"} )) >= (80 / 100)) or vector (0)"#, + ]; + + for case in cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_used_bytes".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "kubelet_volume_stats_capacity_bytes".to_string(), + ), + ], + &["namespace", "persistentvolumeclaim"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_nested_binary_op() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"sum(rate(nginx_ingress_controller_requests{job=~".*"}[2m])) - + ( + sum(rate(nginx_ingress_controller_requests{namespace=~".*"}[2m])) + or + vector(0) + )"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "nginx_ingress_controller_requests".to_string(), + )], + &["namespace", "job"], + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_parse_or_operator() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#" + sum(rate(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}[120s])) by (cluster_name,tenant_name) / + (sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) * 100) + or + 200 * sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) / + sum(sysstat{tenant_name=~"tenant1",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)"#; + + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) + + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0 + or + sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) / + (sum(delta(sysstat{tenant_name=~"sys",cluster_name=~"cluster1"}[2m])/120) by (cluster_name,tenant_name) *1000) >= 0"#; + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "sysstat".to_string())], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"(sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name) + + sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(background_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name)) or + (sum(foreground_waitevent_cnt{tenant_name=~"sys",cluster_name=~"cluster1"}) by (cluster_name,tenant_name))"#; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "background_waitevent_cnt".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "foreground_waitevent_cnt".to_string(), + ), + ], + &["tenant_name", "cluster_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let case = r#"avg(node_load1{cluster_name=~"cluster1"}) by (cluster_name,host_name) or max(container_cpu_load_average_10s{cluster_name=~"cluster1"}) by (cluster_name,host_name) * 100 / max(container_spec_cpu_quota{cluster_name=~"cluster1"}) by (cluster_name,host_name)"#; + let table_provider = build_test_table_provider_with_fields( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "node_load1".to_string()), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_cpu_load_average_10s".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "container_spec_cpu_quota".to_string(), + ), + ], + &["cluster_name", "host_name"], + ) + .await; + eval_stmt.expr = parser::parse(case).unwrap(); + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn value_matcher() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let cases = [ + // single equal matcher + ( + r#"some_metric{__field__="field_1"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two equal matchers + ( + r#"some_metric{__field__="field_1", __field__="field_0"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single not_eq matcher + ( + r#"some_metric{__field__!="field_1"}"#, + vec![ + "some_metric.field_0", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // two not_eq matchers + ( + r#"some_metric{__field__!="field_1", __field__!="field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (no conflict) + ( + r#"some_metric{__field__="field_1", __field__!="field_0"}"#, + vec![ + "some_metric.field_1", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // equal and not_eq matchers (conflict) + ( + r#"some_metric{__field__="field_2", __field__!="field_2"}"#, + vec![ + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex eq matcher + ( + r#"some_metric{__field__=~"field_1|field_2"}"#, + vec![ + "some_metric.field_1", + "some_metric.field_2", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + // single regex not_eq matcher + ( + r#"some_metric{__field__!~"field_1|field_2"}"#, + vec![ + "some_metric.field_0", + "some_metric.tag_0", + "some_metric.tag_1", + "some_metric.tag_2", + "some_metric.timestamp", + ], + ), + ]; + + for case in cases { + let prom_expr = parser::parse(case.0).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let mut fields = plan.schema().field_names(); + let mut expected = case.1.into_iter().map(String::from).collect::>(); + fields.sort(); + expected.sort(); + assert_eq!(fields, expected, "case: {:?}", case.0); + } + + let bad_cases = [ + r#"some_metric{__field__="nonexistent"}"#, + r#"some_metric{__field__!="nonexistent"}"#, + ]; + + for case in bad_cases { + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "case: {:?}", case); + } +} + +#[tokio::test] +async fn custom_schema() { + let query = "some_alt_metric{__schema__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__database__=\"greptime_private\"}"; + let expected = String::from( + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; + + let query = "some_alt_metric{__schema__=\"greptime_private\"} / some_metric"; + let expected = String::from( + "Projection: some_metric.tag_0, some_metric.timestamp, CAST(greptime_private.some_alt_metric.field_0 AS Float64) / CAST(some_metric.field_0 AS Float64) AS greptime_private.some_alt_metric.field_0 / some_metric.field_0 [tag_0:Utf8, timestamp:Timestamp(ms), greptime_private.some_alt_metric.field_0 / some_metric.field_0:Float64;N]\ + \n Inner Join: greptime_private.some_alt_metric.tag_0 = some_metric.tag_0, greptime_private.some_alt_metric.timestamp = some_metric.timestamp [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: greptime_private.some_alt_metric.tag_0 ASC NULLS FIRST, greptime_private.some_alt_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: greptime_private.some_alt_metric.timestamp >= TimestampMillisecond(-999, None) AND greptime_private.some_alt_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: greptime_private.some_alt_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n SubqueryAlias: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +#[tokio::test] +async fn only_equals_is_supported_for_special_matcher() { + let queries = &[ + "some_alt_metric{__schema__!=\"greptime_private\"}", + "some_alt_metric{__schema__=~\"lalala\"}", + "some_alt_metric{__database__!=\"greptime_private\"}", + "some_alt_metric{__database__=~\"lalala\"}", + ]; + + for query in queries { + let prom_expr = parser::parse(query).unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = build_test_table_provider( + &[ + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + ( + "greptime_private".to_string(), + "some_alt_metric".to_string(), + ), + ], + 1, + 1, + ) + .await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await; + assert!(plan.is_err(), "query: {:?}", query); + } +} + +#[tokio::test] +async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), + promql_annotations: None, + }; + planner.ctx.time_index_column = Some("timestamp".to_string()); + planner.ctx.start = 1_000; + planner.ctx.lookback_delta = 0; + let schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + false, + )])) + .unwrap(), + ); + for (end, interval, windows) in [ + (1_000, 1_000, 1), + (2_000, 1_000, 1), + (7_201_000, 7_200_000, 2), + ] { + planner.ctx.end = end; + planner.ctx.interval = interval; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!(filter.matches(">=").count(), windows, "{filter}"); + assert!( + filter.contains("TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + } + planner.ctx.end = i64::MAX; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert!( + filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + + // A lookback subtraction can underflow milliseconds while the upper bound remains + // representable. Keep that upper bound so LastRow cannot select a future sample. + let ms_schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )])) + .unwrap(), + ); + planner.ctx.start = i64::MIN + 100; + planner.ctx.end = planner.ctx.start; + planner.ctx.lookback_delta = 200; + let filter = planner + .build_time_index_filter(0, &ms_schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + format!( + "timestamp <= TimestampMillisecond({}, None)", + i64::MIN + 100 + ) + ); + + // The lower bound can also overflow while converting milliseconds to native nanoseconds. + // Its representable upper bound still has to reach the scan. + planner.ctx.start = 0; + planner.ctx.end = 0; + planner.ctx.lookback_delta = 300_000; + let filter = planner + .build_time_index_filter(9_223_372_036_854, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + "timestamp <= TimestampNanosecond(-9223372036854000000, None)" + ); +} + +#[tokio::test] +async fn test_non_ms_precision() { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "field".to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("metrics".to_string()) + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "metrics".to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + let plan = PromPlanner::stmt_to_plan( + DfTableSourceProvider::new( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("metrics{tag = \"1\"}").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert_eq!( + plan.display_indent_schema().to_string(), + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + ); + let plan = PromPlanner::stmt_to_plan( + DfTableSourceProvider::new( + catalog_list.clone(), + false, + QueryContext::arc(), + DummyDecoder::arc(), + true, + ), + &EvalStmt { + expr: parser::parse("avg_over_time(metrics{tag = \"1\"}[5s])").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert_eq!( + plan.display_indent_schema().to_string(), + "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + ); +} + +#[tokio::test] +async fn test_nonexistent_label() { + // template + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let case = r#"some_metric{nonexistent="hi"}"#; + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 3, + 3, + ) + .await; + // Should be ok + let _ = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_label_join() { + let prom_expr = + parser::parse("label_join(up{tag_0='api-server'}, 'foo', ',', 'tag_1', 'tag_2', 'tag_3')") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 4, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("api-server") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn test_label_replace() { + let prom_expr = + parser::parse("label_replace(up{tag_0=\"a:c\"}, \"foo\", \"$1\", \"tag_0\", \"(.*):.*\")") + .unwrap(); + let eval_stmt = EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + let table_provider = + build_test_table_provider(&[(DEFAULT_SCHEMA_NAME.to_string(), "up".to_string())], 1, 1) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + + let expected = r#" +Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + Filter: up.tag_0 = Utf8("a:c") AND up.timestamp >= TimestampMillisecond(-999, None) AND up.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] + TableScan: up [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]"#; + + let ret = plan.display_indent_schema().to_string(); + assert_eq!(format!("\n{ret}"), expected, "\n{}", ret); +} + +#[tokio::test] +async fn label_replace_aggregation_queries_plan_successfully() { + let aggregate = r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#; + let queries = [ + aggregate.to_string(), + format!("{aggregate} <= 10"), + format!("{aggregate} * 0.8"), + format!("0.8 * {aggregate}"), + format!("{aggregate} <= {aggregate} * 0.8"), + ]; + let state = build_query_engine_state(); + let mut failures = Vec::new(); + + for query in queries { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + if let Err(error) = + PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await + { + failures.push(format!("{query}: {error:?}")); + } + } + + assert!(failures.is_empty(), "{}", failures.join("\n")); +} + +#[tokio::test] +async fn test_matchers_to_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = + r#"sum(prometheus_tsdb_head_series{tag_1=~"(10.0.160.237:8080|10.0.160.237:9090)"})"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + 3, + 3, + ) + .await; + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.timestamp ASC NULLS LAST [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.timestamp]], aggr=[[sum(prometheus_tsdb_head_series.field_0), sum(prometheus_tsdb_head_series.field_1), sum(prometheus_tsdb_head_series.field_2)]] [timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.field_0):Float64;N, sum(prometheus_tsdb_head_series.field_1):Float64;N, sum(prometheus_tsdb_head_series.field_2):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\", \"tag_1\", \"tag_2\"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.tag_0 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_1 ASC NULLS FIRST, prometheus_tsdb_head_series.tag_2 ASC NULLS FIRST, prometheus_tsdb_head_series.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.tag_1 ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, field_1:Float64;N, field_2:Float64;N]"; + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_topk_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"topk(10, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Projection: sum(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [sum(prometheus_tsdb_head_series.greptime_value):Float64;N, ip:Utf8, greptime_timestamp:Timestamp(ms)]\ + \n Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Filter: row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= Float64(10) [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n WindowAggr: windowExpr=[[row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N, row_number() PARTITION BY [prometheus_tsdb_head_series.greptime_timestamp] ORDER BY [sum(prometheus_tsdb_head_series.greptime_value) DESC NULLS FIRST, prometheus_tsdb_head_series.ip DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW:UInt64]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_count_values_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_value_alias() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip)"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + eval_stmt = QueryLanguageParser::apply_alias_extension(eval_stmt, "my_series"); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = r#" +Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.series, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, series:Utf8;N, greptime_timestamp:Timestamp(ms)] + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value) AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Utf8;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prom_float_to_string(prometheus_tsdb_head_series.greptime_value)]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), prom_float_to_string(prometheus_tsdb_head_series.greptime_value):Utf8;N, count(prometheus_tsdb_head_series.greptime_value):Int64] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; + assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); +} + +#[tokio::test] +async fn test_quantile_expr() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10.0.160.237:8080|10.0.160.237:9090)"}) by (ip))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "http_server_requests_seconds_count".to_string(), + ), + ], + &["ip"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + let expected = "Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(Float64(0.3), sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(ms), quantile(Float64(0.3),sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + + assert_eq!(plan.display_indent_schema().to_string(), expected); +} + +#[tokio::test] +async fn test_or_not_exists_table_label() { + let state = build_query_engine_state(); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "normal_metric".to_string())], + &["job"], + ) + .await; + let raw = PromPlanner::stmt_to_plan( + provider, + &build_eval_stmt(r#"missing_metric or on(absent_label) normal_metric"#), + &state, + ) + .await + .unwrap(); + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_0@") + ); + let (optimized, batches) = execute(raw, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); +} + +#[tokio::test] +async fn test_histogram_quantile_missing_le_column() { + let mut eval_stmt = EvalStmt { + expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + + // Test case: histogram_quantile with a table that doesn't have 'le' column + let case = r#"histogram_quantile(0.99, sum by(pod,instance,le) (rate(non_existent_histogram_bucket{instance=~"xxx"}[1m])))"#; + + let prom_expr = parser::parse(case).unwrap(); + eval_stmt.expr = prom_expr; + + // Create a table provider with a table that doesn't have 'le' column + let table_provider = build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "non_existent_histogram_bucket".to_string(), + )], + &["pod", "instance"], // Note: no 'le' column + ) + .await; + + // Should return empty result instead of error + let result = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()).await; + + // This should succeed now (returning empty result) instead of failing with "Cannot find column le" + assert!( + result.is_ok(), + "Expected successful plan creation with empty result, but got error: {:?}", + result.err() + ); + + // Verify that the result is an EmptyRelation + let plan = result.unwrap(); + match plan { + LogicalPlan::EmptyRelation(_) => { + // This is what we expect + } + _ => panic!("Expected EmptyRelation, but got: {:?}", plan), + } +} + +#[tokio::test] +async fn test_direct_or_normalizes_missing_match_labels() { + type Case<'a> = ( + Option>, + Option>, + i64, + i64, + &'a [(f64, Option<&'a str>)], + ); + + let modifier = or_modifier("lhs or on(k) rhs"); + #[rustfmt::skip] + let cases: &[Case<'_>] = &[ + (None, None, 1, 1, &[(1.0, None)]), + (None, Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(Some("")), None, 1, 1, &[(1.0, Some(""))]), + (None, Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("l")), None, 1, 1, &[(1.0, Some("l")), (2.0, None)]), + (Some(None), Some(Some("")), 1, 1, &[(1.0, None)]), + (Some(None), Some(Some("r")), 1, 1, &[(1.0, None), (2.0, Some("r"))]), + (Some(Some("same")), Some(Some("same")), 1, 2, &[(1.0, Some("same")), (2.0, Some("same"))]), + ]; + for &(left, right, left_ts, right_ts, expected) in cases { + let (optimized, batches) = run( + &matrix_source("lhs", left, left_ts, 1.0), + &matrix_source("rhs", right, right_ts, 2.0), + matrix_context("lhs", left), + matrix_context("rhs", right), + &modifier, + ) + .await; + assert_no_internal_or_keys(optimized.schema()); + assert_eq!( + rows(&batches), + expected + .iter() + .map(|(value, label)| (*value, label.map(str::to_string))) + .collect::>() + ); + } +} + +#[tokio::test] +async fn test_direct_or_match_modifiers() { + for (modifier, left, right, expected) in [ + (None, "left", "right", 2), + (or_modifier("lhs or on(k) rhs"), "same", "same", 1), + (or_modifier("lhs or on() rhs"), "left", "right", 1), + (or_modifier("lhs or ignoring(k) rhs"), "left", "right", 1), + ] { + let (_, batches) = run( + &matrix_source("lhs", Some(Some(left)), 1, 1.0), + &matrix_source("rhs", Some(Some(right)), 1, 2.0), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &modifier, + ) + .await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected + ); + } +} + +#[tokio::test] +async fn test_direct_or_nested_projection_uses_left_context() { + let left = matrix_source("lhs", Some(Some("k")), 1, 1.0); + let right = matrix_source("rhs", Some(Some("k")), 1, 2.0); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert!(raw.schema().iter().any(|(qualifier, field)| { + qualifier.as_ref().is_some_and(|q| q.to_string() == "lhs") && field.name() == "v" + })); + let nested = LogicalPlanBuilder::from(raw) + .project(vec![ + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(DfExpr::Column(Column::new( + Some(TableReference::bare("lhs")), + "v", + ))), + op: Operator::Plus, + right: Box::new(lit(1.0)), + }) + .alias("v_plus"), + ]) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(nested, &build_query_engine_state()).await; + assert_eq!(values(&batches, "v_plus"), vec![2.0]); +} + +#[tokio::test] +async fn test_direct_or_skips_user_internal_key_name() { + const USER_TAG: &str = "__promql_or_match_0"; + let left = tagged_source( + "lhs", + false, + (USER_TAG, Some("left")), + DirectOrValue::Float64(1.0), + ); + let right = tagged_source( + "rhs", + false, + (USER_TAG, Some("right")), + DirectOrValue::Float64(2.0), + ); + let raw = plan_direct_or( + scan(&left), + scan(&right), + direct_or_context("lhs", &["job", USER_TAG], "v"), + direct_or_context("rhs", &["job", USER_TAG], "v"), + &or_modifier("lhs or on(missing_label) rhs"), + ) + .await; + assert!( + raw.display_indent_schema() + .to_string() + .contains("__promql_or_match_1@") + ); + let (_, batches) = execute(raw, &build_query_engine_state()).await; + assert!( + batches + .iter() + .all(|batch| batch.column_by_name(USER_TAG).is_some()) + ); +} + +#[tokio::test] +async fn test_direct_or_substrait_round_trip_with_normalized_key() { + let state = build_query_engine_state(); + let ctx = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog + .register_schema("public", Arc::new(MemorySchemaProvider::new())) + .unwrap(); + ctx.register_catalog("datafusion", catalog); + let left = matrix_source("lhs", Some(Some("")), 1, 1.0); + let right = matrix_source("rhs", None, 1, 2.0); + ctx.register_table( + TableReference::full("datafusion", "public", "lhs"), + table(&left), + ) + .unwrap(); + ctx.register_table( + TableReference::full("datafusion", "public", "rhs"), + table(&right), + ) + .unwrap(); + let raw = plan_direct_or( + ctx.table("datafusion.public.lhs") + .await + .unwrap() + .into_unoptimized_plan(), + ctx.table("datafusion.public.rhs") + .await + .unwrap() + .into_unoptimized_plan(), + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + let decoded = DFLogicalSubstraitConvertor + .decode( + DFLogicalSubstraitConvertor + .encode(&raw, DefaultSerializer) + .unwrap(), + ctx.state(), + ) + .await + .unwrap(); + let (optimized, batches) = execute(decoded, &state).await; + assert_no_internal_or_keys(optimized.schema()); + assert!(batches.iter().all(|batch| { + batch + .schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_or_match_")) + })); + assert_eq!(values(&batches, "v"), vec![1.0]); +} + +#[tokio::test] +async fn test_direct_or_numeric_value_types() { + let left = tagged_source("lhs", true, ("k", Some("lhs")), DirectOrValue::Int64(0)); + let right = tagged_source( + "rhs", + false, + ("k", Some("rhs")), + DirectOrValue::Float64(0.5), + ); + let (optimized, batches) = run( + &left, + &right, + direct_or_context("lhs", &["job", "k"], "v"), + direct_or_context("rhs", &["job", "k"], "v"), + &or_modifier("lhs or on(k) rhs"), + ) + .await; + assert_eq!( + optimized + .schema() + .field_with_name(None, "v") + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!(values(&batches, "v"), vec![0.5]); + let provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider: provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job"], "v"); + let right_context = direct_or_context("rhs", &["job"], "v"); + let error = planner + .or_operator( + scan(&job_source("lhs", DirectOrValue::Utf8("x"))), + scan(&job_source("rhs", DirectOrValue::Float64(1.0))), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on() rhs"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("OR value fields have incompatible types") + ); +} + +#[tokio::test] +async fn test_or_with_histogram_quantile_missing_le_column() { + let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); +} + +#[tokio::test] +async fn test_or_with_right_empty_histogram_restores_left_context() { + let eval_stmt = build_eval_stmt( + r#"abs(sum by(instance) (normal_metric) or histogram_quantile(0.99, sum by(pod) (non_existent_histogram_bucket)))"#, + ); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_or_with_both_empty_histograms() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, sum by(pod) (left_histogram_bucket)) or histogram_quantile(0.99, sum by(instance) (right_histogram_bucket))"#, + ); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + match plan { + LogicalPlan::EmptyRelation(relation) => { + assert!(!relation.produce_one_row); + assert!(!relation.schema.fields().is_empty()); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64) + ); + assert!( + relation + .schema + .fields() + .iter() + .any(|field| field.name() == "pod") + ); + assert!( + !relation + .schema + .fields() + .iter() + .any(|field| field.name() == "instance") + ); + } + _ => panic!("Expected EmptyRelation, but got: {plan:?}"), + } +} + +#[tokio::test] +async fn test_nested_or_with_both_empty_histograms() { + for case in [ + r#"abs(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket))"#, + r#"(histogram_quantile(0.99, left_histogram_bucket) or histogram_quantile(0.99, right_histogram_bucket)) + 1"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_test_table_provider_with_fields( + &[ + ( + DEFAULT_SCHEMA_NAME.to_string(), + "left_histogram_bucket".to_string(), + ), + ( + DEFAULT_SCHEMA_NAME.to_string(), + "right_histogram_bucket".to_string(), + ), + ], + &["pod", "instance"], + ) + .await; + + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn test_or_with_empty_histogram_modifiers() { + for case in [ + r#"histogram_quantile(0.99, non_existent_histogram_bucket) or on(pod) normal_metric"#, + r#"normal_metric or ignoring(instance) histogram_quantile(0.99, non_existent_histogram_bucket)"#, + ] { + let eval_stmt = build_eval_stmt(case); + let table_provider = build_missing_le_or_normal_metric_table_provider().await; + + let plan = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap(); + assert_normal_metric_schema(&plan); + } +} + +#[tokio::test] +async fn test_unless_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric unless on(job) normal_metric) or fallback_metric"#, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &state, + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); + let (optimized, physical) = optimize_and_create_physical_plan(&state, plan).await; + assert!(contains_histogram_fold(&optimized), "{optimized:?}"); + let batches = datafusion::physical_plan::collect(physical, state.session_state().task_ctx()) + .await + .unwrap(); + assert!(batches.iter().all(|batch| batch.num_rows() == 0)); +} + +#[tokio::test] +async fn test_and_preserves_left_context_for_histogram() { + let eval_stmt = build_eval_stmt( + r#"histogram_quantile(0.99, bucket_metric and on(job) normal_metric) or fallback_metric"#, + ); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(contains_histogram_fold(&plan), "{plan:?}"); +} + +#[tokio::test] +async fn test_and_preserves_left_context_when_le_is_missing() { + let eval_stmt = + build_eval_stmt(r#"histogram_quantile(0.99, normal_metric and on(job) bucket_metric)"#); + let plan = PromPlanner::stmt_to_plan( + build_set_op_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!(matches!(&plan, LogicalPlan::EmptyRelation(_)), "{plan:?}"); + assert!(!plan.schema().fields().is_empty()); + assert!(!contains_histogram_fold(&plan), "{plan:?}"); +} + +async fn build_matching_filter_plan(query: &str) -> String { + let table_provider = build_test_table_provider_with_distinct_tags(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +/// [`build_test_table_provider_with_distinct_tags`] plus a `status` string column: a value +/// field that is neither a primary key nor a value column of the metric, so +/// `count by(status) (...)` still reports it among the aggregation's tag columns. +async fn build_test_table_provider_with_string_field( + table_tags: &[(&str, &[&str])], +) -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + for (table_name, tags) in table_tags { + let mut columns = tags + .iter() + .map(|tag| { + ColumnSchema::new( + (*tag).to_string(), + ConcreteDataType::string_datatype(), + false, + ) + }) + .collect::>(); + columns.push( + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ); + columns.push(ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + )); + columns.push(ColumnSchema::new( + "status".to_string(), + ConcreteDataType::string_datatype(), + true, + )); + let table_meta = TableMetaBuilder::empty() + .schema(Arc::new(Schema::new(columns))) + .primary_key_indices((0..tags.len()).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name((*table_name).to_string()) + .meta(table_meta) + .build() + .unwrap(); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: (*table_name).to_string(), + table_id: 1024, + table: EmptyTable::from_table_info(&table_info), + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +async fn build_matching_filter_plan_with_string_field(query: &str) -> String { + let table_provider = build_test_table_provider_with_string_field(&[ + ("metric_a", &["host", "device"]), + ("metric_b", &["host", "device"]), + ]) + .await; + let plan = PromPlanner::stmt_to_plan( + table_provider, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap(); + plan.display_indent().to_string() +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_both_operands() { + for query in [ + r#"metric_a / metric_b{host="foo"}"#, + r#"metric_a / on(host, device) metric_b{host="foo"}"#, + r#"count_over_time(metric_a[1m]) / on(host) count_over_time(metric_b{host="foo"}[1m])"#, + r#"metric_a / ignoring(device) metric_b{host="foo"}"#, + r#"sum by(host) (metric_a) / on(host) sum by(host) (metric_b{host="foo"})"#, + r#"metric_a / on(host) avg without(device) (metric_b{host="foo"})"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_scalar_ranking_and_grouped_operands() { + for query in [ + r#"(8 * metric_a{host="foo"}) / on(host) metric_b"#, + r#"topk(1, metric_a{host="foo"}) / on(host, device) metric_b"#, + r#"(8 * metric_a{host="foo"}) / on(host) group_left topk by(host)(1, max by(host)(metric_b))"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 2, + "{query}\n{plan}" + ); + } + // A global ranking one-side must see every host, so the matcher stays put. + let query = r#"metric_a{host="foo"} / on(host) group_left topk(1, max by(host)(metric_b))"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!( + plan.matches(r#"host = Utf8("foo")"#).count(), + 1, + "{query}\n{plan}" + ); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_selecting_aggregations() { + // `topk` ranks its input, so filtering before it changes the candidate set. + let query = r#"topk(1, metric_a) / on(host, device) metric_b{host="foo"}"#; + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_value_field_matcher_stays_on_its_own_operand() { + let value = greptime_value(); + for query in [ + format!(r#"metric_a / metric_b{{{value}="2"}}"#), + format!(r#"metric_a / on(host, device, {value}) metric_b{{{value}="2"}}"#), + ] { + let plan = build_matching_filter_plan(&query).await; + assert_eq!(plan.matches(r#"Utf8("2")"#).count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn binary_value_field_matcher_is_not_copied_across_aggregations() { + // `status` varies between the samples of one series, so filtering the other operand by it + // would drop the newest sample before sample selection (#9242). + for query in [ + r#"count by(status) (metric_a) / on(status) count by(status) (metric_b{status="ready"})"#, + r#"(8 * count by(status)(metric_a{__field__="status"})) / on(status) topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + r#"count by(status)(metric_a{__field__="status"}) / on(status) group_left topk by(status)(1, count by(status)(metric_b{__field__="status",status="ready"}))"#, + ] { + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!( + plan.matches(r#"Utf8("ready")"#).count(), + 1, + "{query}\n{plan}" + ); + } +} + +#[tokio::test] +async fn binary_matching_label_filter_reaches_aggregations_grouping_by_tags() { + // `count`, not `sum`: the string value field is not summable. + let query = r#"count by(host) (metric_a) / on(host) count by(host) (metric_b{host="foo"})"#; + let plan = build_matching_filter_plan_with_string_field(query).await; + assert_eq!(plan.matches(r#"Utf8("foo")"#).count(), 2, "{query}\n{plan}"); +} + +#[tokio::test] +async fn binary_matching_label_filter_skips_unproven_expressions() { + for query in [ + r#"metric_a > on(host, device) metric_b{host="foo"}"#, + r#"metric_a / on(host) group_left metric_b{host="foo"}"#, + r#"metric_a / on(host) label_replace(metric_b{host="foo"},"extra","e","host",".*")"#, + r#"metric_a / on(device) metric_b{host="foo"}"#, + ] { + let plan = build_matching_filter_plan(query).await; + assert_eq!(plan.matches("foo").count(), 1, "{query}\n{plan}"); + } +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_output() { + let case = r#"(normal_metric or other_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let state = build_query_engine_state(); + let plan = + PromPlanner::stmt_to_plan(build_or_context_table_provider().await, &eval_stmt, &state) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); + let (_optimized, _physical) = optimize_and_create_physical_plan(&state, plan).await; +} + +#[tokio::test] +async fn test_or_context_uses_left_qualified_empty_histogram_output() { + let case = r#"(abs(histogram_quantile(0.99, non_hist_metric)) or normal_metric) + 1"#; + let eval_stmt = build_eval_stmt(case); + let plan = PromPlanner::stmt_to_plan( + build_or_context_table_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .any(|field| field.data_type() == &ArrowDataType::Float64), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_direct_or_preserves_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (planner, plan) = mixed_direct_or(histogram_on_left).await; + + let float_field = &planner.ctx.field_columns[0]; + let histogram_field = &planner.ctx.field_columns[1]; + assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); + assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); + assert_eq!( + plan.schema() + .field_with_name(None, float_field) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + plan.schema() + .field_with_name(None, histogram_field) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + + let (optimized, batches) = execute(plan, &build_query_engine_state()).await; + assert_no_internal_or_keys(optimized.schema()); + let mut sample_kinds = batches + .iter() + .flat_map(|batch| { + let values = batch.column_by_name(float_field).unwrap(); + let histograms = batch.column_by_name(histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (values.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + sample_kinds.sort_unstable(); + assert_eq!(sample_kinds, vec![(false, true), (true, false)]); + } +} + +#[tokio::test] +async fn malformed_classic_bucket_does_not_drop_native_histogram() { + let state = build_query_engine_state(); + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &value_field), vec![0.0]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert!(infos.is_empty()); +} + +#[tokio::test] +async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!( + plan_text.contains("prom_native_histogram_mul_scalar"), + "{plan_text}" + ); + assert!( + plan_text.contains("prom_native_histogram_scalar_mul"), + "{plan_text}" + ); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert!(values(&batches, &float_field).is_empty()); + let mut sums = histograms(&batches, &histogram_field) + .into_iter() + .map(|histogram| histogram.sum) + .collect::>(); + sums.sort_by(f64::total_cmp); + assert_eq!(sums, vec![2.0, 3.0]); +} + +#[tokio::test] +async fn test_mixed_binary_operator_reports_only_dropped_samples() { + for (query, expected_rows, expected_infos) in [ + ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + ] { + let state = build_query_engine_state(); + let annotations = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + Some(annotations.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows, + "{query}" + ); + let mut warnings = vec![]; + let mut infos = vec![]; + annotations.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty(), "{query}: {warnings:?}"); + assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); + } +} + +#[tokio::test] +async fn test_histogram_only_min_drops_empty_aggregate_group() { + // `min` over native-histogram-only input drops every sample in the group, so the + // NULL-valued aggregate row must be filtered out. Otherwise an outer expression + // like `group()` resurrects the group Prometheus considers unseen. + let state = build_query_engine_state(); + for query in ["min(lh)", "group(min(lh))"] { + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "{query}" + ); + } +} + +#[tokio::test] +async fn test_mixed_min_drops_histogram_only_group() { + // With alternative float/histogram fields, `min by (tag)` keeps float-only groups + // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of + // emitting a NULL-valued row for them. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &float_field), vec![2.0]); +} + +#[tokio::test] +async fn test_mixed_or_can_feed_another_or() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let mut float_values = values(&batches, &float_field); + float_values.sort_by(f64::total_cmp); + assert_eq!(float_values, vec![2.0, 7.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_mixed_fields_align_with_single_float_vector() { + let (planner, mixed) = mixed_direct_or(false).await; + let scale = tagged_source( + "scale", + false, + ("k", Some("float")), + DirectOrValue::Float64(2.0), + ); + let scale = scan(&scale); + let scale_fields = vec!["v".to_string()]; + let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { + unreachable!() + }; + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + mixed.schema(), + scale.schema(), + &planner.ctx.field_columns, + &scale_fields, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(_, right)| *right == &scale_fields[0]) + ); + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + scale.schema(), + mixed.schema(), + &scale_fields, + &planner.ctx.field_columns, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(left, _)| *left == &scale_fields[0]) + ); +} + +#[tokio::test] +async fn test_non_bool_comparison_filters_mixed_sample_lanes() { + let (planner, input) = mixed_direct_or(false).await; + let input_schema = input.schema().clone(); + let plan = planner + .filter_on_field_column(input, |field| { + if PromPlanner::field_column_is_native_histogram(&input_schema, field) { + Ok(lit(false)) + } else { + Ok(col(field).gt(lit(0.0))) + } + }) + .unwrap(); + let float_field = planner.ctx.field_columns[0].clone(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_field), vec![1.25]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); +} + +#[tokio::test] +async fn test_mixed_left_and_unless_preserve_sample_lanes() { + for (expression, expected_sample_kind) in [ + ("lhs and on(k) mask", (false, true)), + ("lhs unless on(k) mask", (true, false)), + ] { + let (mut planner, left) = mixed_direct_or(false).await; + let left_context = planner.ctx.clone(); + let float_field = left_context.field_columns[0].clone(); + let histogram_field = left_context.field_columns[1].clone(); + let mask = tagged_source( + "mask", + false, + ("k", Some("histogram")), + DirectOrValue::Float64(1.0), + ); + let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { + unreachable!() + }; + let plan = planner + .set_op_on_non_field_columns( + left, + scan(&mask), + left_context, + direct_or_context("mask", &["job", "k"], "v"), + binary.op, + &binary.modifier, + ) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let sample_kinds = batches + .iter() + .flat_map(|batch| { + let floats = batch.column_by_name(&float_field).unwrap(); + let histograms = batch.column_by_name(&histogram_field).unwrap(); + (0..batch.num_rows()).map(|row| (floats.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); + } +} + +#[tokio::test] +async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("some_metric * scalar(vector(2))"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let schema = plan.schema(); + assert_eq!( + schema + .field_with_unqualified_name(greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + schema + .field_with_unqualified_name(greptime_native_histogram()) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + assert!( + plan.display_indent_schema() + .to_string() + .contains("prom_native_histogram_mul_scalar"), + "{plan:?}" + ); +} + +#[tokio::test] +async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![0.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); +} + +#[tokio::test] +async fn test_unary_negates_mixed_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (mut planner, input) = mixed_direct_or(histogram_on_left).await; + let plan = planner.negate_field_columns(input).unwrap(); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let float_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap(); + let histogram_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, float_field), vec![-1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(histogram_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, -1.0); + assert_eq!(histogram.sum, -1.0); + assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); + } +} + +#[tokio::test] +async fn test_native_histogram_sum_and_avg_execute_real_batches() { + for op_name in ["sum", "avg"] { + for incompatible in [false, true] { + let mut second = direct_or_histogram(); + if incompatible { + second.schema = CUSTOM_BUCKETS_SCHEMA; + second.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = + mixed_aggregate_input(vec![direct_or_histogram(), second]).await; + planner.promql_annotations = Some(collector.clone()); + let histogram_column = planner.ctx.field_columns[1].clone(); + planner.ctx.field_columns = vec![histogram_column.clone()]; + let input = LogicalPlanBuilder::from(input) + .project([col("ts"), col(&histogram_column)]) + .unwrap() + .build() + .unwrap(); + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse(&format!("{op_name}(mixed)")).unwrap() + else { + unreachable!() + }; + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter(planner.create_empty_values_filter_expr(false).unwrap()) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(infos.is_empty()); + if incompatible { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + assert!(warnings.iter().any(|warning| { + warning + == &format!( + "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" + ) + })); + } else { + let histograms = histograms(&batches, &histogram_column); + assert_eq!(histograms.len(), 1); + let expected = if op_name == "sum" { 2.0 } else { 1.0 }; + assert_eq!(histograms[0].count, expected); + assert_eq!(histograms[0].sum, expected); + assert!(warnings.is_empty()); + } + } + } +} + +#[tokio::test] +async fn test_canonical_mixed_count_group_and_count_values_execute() { + let state = build_query_engine_state(); + for (query, expected) in [ + ("count(some_metric)", vec![2.0]), + ("group(some_metric)", vec![1.0]), + (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_sample_count")), + "{query}: {plan:?}" + ); + let value_fields = plan + .schema() + .fields() + .iter() + .filter(|field| { + matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) || field.data_type() == &native_histogram_value_type().as_arrow_type() + }) + .collect::>(); + assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); + assert_ne!( + value_fields[0].data_type(), + &native_histogram_value_type().as_arrow_type(), + "{query}: {plan:?}" + ); + let value_column = value_fields[0].name().clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = numeric_values(&batches, &value_column); + actual.sort_by(f64::total_cmp); + assert_eq!(actual, expected, "{query}"); + + if query.starts_with("count_values") { + let mut sample_labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("sample") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + }) + .collect::>(); + sample_labels.sort(); + let mut expected_labels = vec!["2".to_string(), direct_or_histogram().promql_string()]; + expected_labels.sort(); + assert_eq!(sample_labels, expected_labels); + } + } +} + +#[tokio::test] +async fn test_mixed_or_sum_aggregates_each_sample_type() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts"), col("k")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([ + col(&float_column), + col(&histogram_column), + col("ts"), + col("k"), + ]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(&histogram_column)? + .as_any() + .downcast_ref::()?; + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, 1.0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!( + warnings, + vec!["sum: dropped aggregation result containing both float and native histogram samples"] + ); +} + +#[tokio::test] +async fn test_mixed_or_sum_drops_incompatible_mixed_group() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + let mut custom = direct_or_histogram(); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.iter().any(|warning| { + warning + == "sum: dropped aggregation result containing both float and native histogram samples" + })); +} + +#[tokio::test] +async fn test_mixed_or_min_records_only_present_histograms() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = parser::parse("min(lhs)").unwrap() + else { + unreachable!() + }; + let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; + + for (histograms, expected_infos) in [ + (vec![], vec![]), + (vec![direct_or_histogram()], vec![expected_info]), + ] { + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(histograms).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_ignored_histogram_filter_expr(op, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!(infos, expected_infos); + } +} + +#[tokio::test] +async fn test_mixed_or_value_aliases_do_not_replace_labels() { + let left = source( + "lhs", + false, + 1, + vec![("job", Some("job")), ("k", Some("float"))], + DirectOrValue::Float64(1.0), + ); + let right = source( + "rhs", + false, + 1, + vec![ + ("job", Some("job")), + ("k", Some("histogram")), + (greptime_value(), Some("value-label")), + ], + DirectOrValue::NativeHistogram(direct_or_histogram()), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left = LogicalPlanBuilder::from(scan(&left)) + .project(vec![ + col("ts"), + col("job"), + col("k"), + col("v").alias(greptime_value()), + ]) + .unwrap() + .build() + .unwrap(); + let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); + let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); + let plan = planner + .or_operator( + left, + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + + assert_eq!( + plan.schema() + .field_with_name(None, greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Utf8 + ); + assert!( + planner + .ctx + .field_columns + .iter() + .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) + ); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(labels, vec!["value-label"]); +} + +#[tokio::test] +async fn test_mixed_or_routes_float_histogram_and_label_functions() { + for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { + unreachable!() + }; + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let values = batches + .iter() + .flat_map(|batch| { + batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .into_iter() + .flatten() + }) + .collect::>(); + assert_eq!(values, vec![expected], "{function}"); + } + + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = + parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() + else { + unreachable!() + }; + let args = planner.create_function_args(&call.args.args).unwrap(); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, args.literals, input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(sample_count, 2); +} + +/// Table provider with a single metric `cv_metric` holding three series at `ts=1000`: +/// `k="a"` and `k="c"` carry the value 1.0, `k="b"` carries 2.0. +async fn build_count_values_table_provider() -> DfTableSourceProvider { + build_count_values_table_provider_with_values(&[1.0, 2.0, 1.0]).await +} + +/// Table provider with a single metric `cv_metric` holding one series per given value at +/// `ts=1000` (`k` names the series, the sample value is the given value). +async fn build_count_values_table_provider_with_values(values: &[f64]) -> DfTableSourceProvider { + build_count_values_table_provider_with_value_array(Arc::new(Float64Array::from( + values.to_vec(), + ))) + .await +} + +/// Like [`build_count_values_table_provider_with_values`], but with a caller provided value +/// column, so tests can cover value columns that are not `Float64` (e.g. `BIGINT`). +async fn build_count_values_table_provider_with_value_array( + values: ArrayRef, +) -> DfTableSourceProvider { + let value_data_type = ConcreteDataType::from_arrow_type(values.data_type()); + let catalog_list = MemoryCatalogManager::with_default_setup(); + let columns = vec![ + ColumnSchema::new("k".to_string(), ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value().to_string(), value_data_type, true), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(3_001) + .name("cv_metric") + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from( + (0..values.len()) + .map(|index| format!("k{index}")) + .collect::>(), + )), + Arc::new(TimestampMillisecondArray::from(vec![1_000; values.len()])), + values.clone(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + "cv_metric", + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 3_001, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "cv_metric".to_string(), + table_id: 3_001, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +/// Collects `(label, value)` pairs of a `count_values` result, where `label` is the +/// PromQL label generated by `count_values` and `value` is the aggregated sample value. +/// +/// The generated label holds the original sample value in PromQL's textual form, so it +/// is asserted as a string: comparing it as a number would not catch formatting bugs +/// (`1.0` instead of `1`, scientific notation, ...). +fn count_values_rows<'a>(batches: &'a [RecordBatch], label: &str) -> Vec<(&'a str, f64)> { + let mut rows = batches + .iter() + .flat_map(|batch| { + // The aggregated value is the only numeric column that is not the generated label. + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| { + field.name() != label + && matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) + }) + .expect("no aggregated value column"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + let values = datafusion::arrow::compute::cast( + batch.column(value_index), + &ArrowDataType::Float64, + ) + .unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + labels + .iter() + .zip(values.iter()) + .map(|(label, value)| (label.unwrap(), value.unwrap())) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| left.0.cmp(right.0).then(left.1.total_cmp(&right.1))); + rows +} + +/// Asserts that a `count_values` result holds one sample per label set and evaluation +/// timestamp: Prometheus groups by the generated label, so a timestamp must never repeat +/// a label set (that would mean the samples were still grouped by the overwritten input +/// label). +fn assert_unique_label_set_per_timestamp(batches: &[RecordBatch], label: &str) { + let mut seen = HashMap::>::new(); + for batch in batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let labels = batch + .column_by_name(label) + .expect("no generated label column") + .as_any() + .downcast_ref::() + .expect("the generated label must be a string column"); + for (timestamp, label) in timestamps.iter().zip(labels.iter()) { + let timestamp = timestamp.unwrap(); + let label = label.unwrap(); + assert!( + seen.entry(timestamp).or_default().insert(label.to_string()), + "duplicated label set `{label}` at timestamp {timestamp}" + ); + } + } +} + +#[tokio::test] +async fn test_count_values_generated_label_survives_enclosing_expr() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + for (case, label) in [ + (r#"count_values("v", prometheus_tsdb_head_series)"#, "v"), + ( + r#"abs(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"round(count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + (r#"count_values("v", prometheus_tsdb_head_series) + 1"#, "v"), + ( + r#"topk(1, count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"sum by (v) (count_values("v", prometheus_tsdb_head_series))"#, + "v", + ), + ( + r#"label_replace(count_values("v", prometheus_tsdb_head_series), "vcopy", "$1", "v", "(.*)")"#, + "v", + ), + ( + r#"count_values("v", prometheus_tsdb_head_series) by (ip) + 1"#, + "v", + ), + // The generated label overwrites an input label with the same name. + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip)"#, + "ip", + ), + ( + r#"count_values("ip", prometheus_tsdb_head_series) by (ip) + 1"#, + "ip", + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_table_provider_with_fields( + &[( + DEFAULT_SCHEMA_NAME.to_string(), + "prometheus_tsdb_head_series".to_string(), + )], + &["ip"], + ) + .await, + &build_eval_stmt(case), + &build_query_engine_state(), + ) + .await + .unwrap(); + + let label_columns = plan + .schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(); + assert_eq!( + label_columns, + 1, + "{case}: the `{label}` label must survive: {}", + plan.display_indent() + ); + } +} + +#[tokio::test] +async fn test_count_values_generated_label_in_enclosing_expr_execute() { + // https://github.com/GreptimeTeam/greptimedb/issues/9181 + let state = build_query_engine_state(); + // (query, generated label, expected `(label value, aggregated value)` pairs) + for (query, label, expected) in [ + ( + r#"count_values("v", cv_metric)"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"abs(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"round(count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"count_values("v", cv_metric) + 1"#, + "v", + vec![("1", 3.0), ("2", 2.0)], + ), + ( + r#"sum by (v) (count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + ( + r#"topk(10, count_values("v", cv_metric))"#, + "v", + vec![("1", 2.0), ("2", 1.0)], + ), + // The generated label overwrites the input label with the same name, and the + // samples are grouped by the generated label only: `{k="1"}` holds the two + // samples of value `1.0` instead of one row per (overwritten label, value). + ( + r#"count_values("k", cv_metric) by (k)"#, + "k", + vec![("1", 2.0), ("2", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider().await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap_or_else(|err| panic!("{query}: {err}")); + + assert_eq!( + plan.schema() + .fields() + .iter() + .filter(|field| field.name() == label) + .count(), + 1, + "{query}: {}", + plan.display_indent() + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(count_values_rows(&batches, label), expected, "{query}"); + assert_unique_label_set_per_timestamp(&batches, label); + } +} + +#[tokio::test] +async fn test_count_values_label_is_prometheus_formatted_value() { + // PromQL materializes the generated label with `strconv.FormatFloat(value, 'f', -1, 64)`: + // the shortest decimal form of the sample value without an exponent. The label is a + // label, so it must be a string column holding exactly that text: arrow's + // `Float64 -> Utf8` cast would render `1`/`200`/`1e21` as `1.0`/`200.0`/`1e21`. + let state = build_query_engine_state(); + // Samples are grouped by that formatted text, exactly like Prometheus groups by the + // generated label: `-0.0` and `0.0` are two series (`-0` and `0`), while values that + // round to the same text share one group. `0.0` formats as "0". + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_values(&[ + -0.0, 0.0, 1.0, 0.5, 200.0, 1e21, 1e-7, 2.5, + ]) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_unique_label_set_per_timestamp(&batches, "v"); + let mut labels = count_values_rows(&batches, "v") + .into_iter() + .map(|(label, _)| label) + .collect::>(); + labels.sort(); + assert_eq!( + labels, + vec![ + "-0", + "0", + "0.0000001", + "0.5", + "1", + "1000000000000000000000", + "2.5", + "200", + ] + ); +} + +#[tokio::test] +async fn test_count_values_groups_by_formatted_value_for_bigint_input() { + // The grouping key of `count_values` is the formatted sample value, not the raw input + // value. Two `BIGINT` values that differ below the `Float64` precision (`2^53` and + // `2^53 + 1`) cast and format to the same label, so they must share one group and one + // count, exactly like Prometheus, which groups by the generated label text. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + build_count_values_table_provider_with_value_array(Arc::new(Int64Array::from(vec![ + 9_007_199_254_740_992_i64, + 9_007_199_254_740_993_i64, + ]))) + .await, + &operator_eval_stmt(r#"count_values("v", cv_metric)"#), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + // One timestamp must never carry the same label set twice. + assert_unique_label_set_per_timestamp(&batches, "v"); + assert_eq!( + count_values_rows(&batches, "v"), + vec![("9007199254740992", 2.0)] + ); +} From 43eaea7a9af5a3906b6d918872e02b6536a664ef Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 28 Sep 2026 06:08:14 +0000 Subject: [PATCH 050/133] fix(ci): teach check-builder-rust-version.sh to handle stable channels (#9369) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ci): teach check-builder-rust-version.sh to handle stable channels The script extracted a YYYY-MM-DD date from rust-toolchain.toml to compare against the rustc build date inside the dev-builder image — a nightly-era design. With channel = "1.96.1" there is no date in the file, so every release build failed with 'Error: No rust toolchain version found in rust-toolchain.toml'. Extract the channel token instead and branch on it: - stable channel (X.Y[.Z]): require the builder image's rustc to exactly match the pinned version - nightly-YYYY-MM-DD: keep the legacy date-difference check Verified against a mocked docker for all four paths (stable match / mismatch, nightly fresh / stale). Part of #9289. Signed-off-by: Ning Sun * chore(toolchain): finish stable-migration cleanup in docs and query-regression pin - README/AGENTS: the toolchain is now stable Rust pinned by rust-toolchain.toml, not nightly - query-regression: align the benchmark toolchain pin with the workspace (nightly-2026-03-21 = 1.96.0-nightly -> stable 1.96.1), including the exact-version assertions (cargo 356927216, rustc 31fca3adb, both 2026-06-26) and the runner image default The query-regression runner image must be rebuilt and QUERY_REGRESSION_ECS_IMAGE_ID bumped together with these pins (per .github/runner-scale-sets/query-regression/README.md) before the next regression run. Part of #9289. Signed-off-by: Ning Sun * ci(query-regression): derive the Rust toolchain pin from rust-toolchain.toml Replace the hard-coded RUSTUP_TOOLCHAIN value and the hard-coded version strings in the runner Verify assertions with a pin resolved from rust-toolchain.toml: - the always-running test-tooling job exports the channel parsed from rust-toolchain.toml as a job output - query-regression sets RUSTUP_TOOLCHAIN from that output - the Verify step escapes the pin into the cargo/rustc/active-toolchain regexes at runtime; the exact commit hash and date are asserted generically since a stable version identifies the release Removing the redundant require_eq (workflow yaml vs runner env) since both now flow from the single source of truth. When rust-toolchain.toml is bumped, the run fails with a clear signal until the runner image is rebuilt with the new toolchain, keeping the existing lockstep contract. Part of #9289. Signed-off-by: Ning Sun * ci(query-regression): derive the runner image toolchain from rust-toolchain.toml Remove the hard-coded 'ARG RUST_TOOLCHAIN=1.96.1' from the query-regression runner Dockerfile. The pin is now parsed from a COPY'd rust-toolchain.toml at build time (the bootstrap script builds with the repo root as context, so the file is in the build context): - rustup-init installs the parsed channel as the default toolchain - the baked ENV RUSTUP_TOOLCHAIN is dropped: the rustup default makes bare cargo/rustc resolve correctly without it, and the workflow supplies RUSTUP_TOOLCHAIN explicitly at run time - the build-time self-verification asserts the active toolchain against the same parsed pin With this, rust-toolchain.toml is the single source of truth for the benchmark toolchain end to end: the image bakes whatever the toml says at build time and the workflow asserts against the toml at run time. A toolchain bump now only requires rebuilding the image. The changed mechanics were verified natively with the real rustup-init 1.29.0 and the real 1.96.1 toolchain (registry pulls are unavailable in this sandbox): parsing, default-toolchain installation without RUSTUP_TOOLCHAIN env, bare cargo/rustc resolution, and the active-toolchain assertion for both '(default)' and '(overridden)' output forms. Part of #9289. Signed-off-by: Ning Sun * ci(query-regression): rebuild the runner image automatically on toolchain changes Mirror the dev-builder automation for the query-regression ECS runner image: a new rebuild-query-regression-runner-image.yaml workflow runs whenever rust-toolchain.toml or the query-regression runner directory changes on main (or via manual dispatch). It drives the existing build-ecs-image.py ops tool, then completes the documented lockstep updates in order: bump RUNNER_IMAGE_EPOCH in query-regression.yml and push the commit to main, and only then point the QUERY_REGRESSION_ECS_IMAGE_ID repo variable at the new image, so the next regression run picks up image and epoch together. Also fix build-ecs-image.py to stage rust-toolchain.toml into the temporary docker build context: the AMI path builds the embedded Dockerfile from an empty /tmp/image-context, which would break on the Dockerfile's COPY of rust-toolchain.toml introduced earlier. The user-data now base64-stages the toml next to the Dockerfile before docker build. Verified: py_compile, render_user_data round-trip (mkdir -> stage -> docker build ordering), and the RUNNER_IMAGE_EPOCH bump sed against the real workflow file. Requires a new ALIYUN_ECS_BASE_IMAGE_ID repo variable (Ubuntu 24.04 public image id in the region); all other secrets/vars are shared with the provisioning job. Part of #9289. Signed-off-by: Ning Sun * ci: fold the query-regression runner rebuild into release-dev-builder-images.yaml Merge the standalone rebuild workflow into the existing builder-image release workflow, as one entry point for all builder artifacts: - push paths extended with .github/runner-scale-sets/query-regression/** - new 'release_query_regression_runner_image' dispatch input - a 'changes' job diffs the pushed range (github.event.before..sha, with an everything-changed fallback for dispatch or unknown bases) so each expensive rebuild only fires for its own paths: rust-toolchain.toml gates both, docker/dev-builder/** gates the dev-builder images, the query-regression runner directory gates the ECS image rebuild - the rebuild job itself is unchanged from the standalone workflow (build-ecs-image.py, then RUNNER_IMAGE_EPOCH commit to main, then the QUERY_REGRESSION_ECS_IMAGE_ID variable update) The dev-builder jobs, their ECR/CN/tag-update dependents, and the runner rebuild now share one workflow; the changes filter preserves the previous on-push behavior for the dev-builder images while the runner rebuild keeps its own trigger. The epoch-bump commit only touches query-regression.yml, which is outside the trigger paths, so no re-trigger loop. Part of #9289. Signed-off-by: Ning Sun * fix(ci): auto-resolve the ECS base image for the runner rebuild The automated rebuild failed with 'Missing required configuration: --base-image-id' because the ALIYUN_ECS_BASE_IMAGE_ID repo variable does not exist yet (it was flagged as a one-time setup item). Remove the setup dependency instead: build-ecs-image.py now defaults --base-image-id to the latest public Ubuntu 24.04 x86_64 system image in the region (DescribeImages with image_owner_alias=system), so no manual variable is required. The runner Dockerfile pins every tool version itself, so base-image drift is low-risk; --base-image-id or the ALIYUN_ECS_BASE_IMAGE_ID variable still pin a specific base image deterministically, and the workflow only passes the flag when the variable is set. Part of #9289. Signed-off-by: Ning Sun * docs(query-regression): clarify what an image rebuild requires A routine runner-image rebuild needs no manual file updates: the rebuild job updates QUERY_REGRESSION_ECS_IMAGE_ID and RUNNER_IMAGE_EPOCH; the toolchain derives from rust-toolchain.toml; uv, sccache, otelgen, rustup, and the runner base are pinned by digest/sha/commit in the Dockerfile. Only an apt package revision bump (mold, protoc, python3) between rebuilds requires bumping the corresponding Verify pins, and that failure is loud with the observed version. Part of #9289. Signed-off-by: Ning Sun * fix(ci): correct SDK field names in the base-image resolver DescribeImagesRequest takes 'ostype' (not 'os_type') and the image items expose 'osname'/'osname_en' (not 'os_name') in the pinned alibabacloud_ecs20140526 SDK range, so the auto-resolution added in 5a9fd2c769 crashed with a TypeError before describing anything. Fix the request fields, move the architecture filter server-side, and paginate (page_size=100 until a short page) instead of relying on a single default-sized response. Match Ubuntu 24.04 on the localized osname or the English osname_en. Verified against the real SDK models (uv run --with 'alibabacloud_ecs20140526>=4.1.0,<6'): a two-page fake client picks the newest Ubuntu 24.04 via osname_en and rejects 22.04/Windows decoys. Part of #9289. Signed-off-by: Ning Sun * fix(ci): correct the repo-root path in build-ecs-image.py ASSETS_DIR.parent.parent lands on runner-scale-sets, not the repo root -- the toml lookup failed with FileNotFoundError. The root is four levels above ecs-image; express it as an explicit REPO_ROOT constant (ASSETS_DIR.parents[3]). Verified every path main() reads against the real checkout layout (Dockerfile, rust-toolchain.toml, start-runner.sh, the systemd unit, plus REPO_ROOT sanity against Cargo.toml/.git), re-checked the user-data toml staging round-trip, and re-ran the base-image resolver regression test. Part of #9289. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .../query-regression/Dockerfile | 17 +- .../query-regression/README.md | 34 +++- .../ecs-image/build-ecs-image.py | 75 ++++++++- .github/workflows/query-regression.yml | 23 ++- .../workflows/release-dev-builder-images.yaml | 146 +++++++++++++++++- AGENTS.md | 2 +- README.md | 2 +- scripts/check-builder-rust-version.sh | 48 ++++-- 8 files changed, 313 insertions(+), 34 deletions(-) diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index d3e89e81769..9b6a8d5ebb6 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -22,11 +22,14 @@ ARG SCCACHE_SHA256=aec995a83ad3dff3d14b6314e08858b7b73d35ca85a5bcf3d3a9ec07dee35 ARG RUSTUP_INIT_VERSION=1.29.0 ARG RUSTUP_INIT_TARGET=x86_64-unknown-linux-gnu ARG RUSTUP_INIT_SHA256=4acc9acc76d5079515b46346a485974457b5a79893cfb01112423c89aeb5aa10 -ARG RUST_TOOLCHAIN=nightly-2026-03-21 + +# Single source of truth for the Rust toolchain: parsed from +# rust-toolchain.toml at build time, so the image always bakes the +# workspace toolchain without a second hard-coded pin here. +COPY rust-toolchain.toml /opt/rust-toolchain.toml ENV RUSTUP_HOME=/opt/rustup \ CARGO_HOME=/opt/cargo \ - RUSTUP_TOOLCHAIN=${RUST_TOOLCHAIN} \ RUSTUP_AUTO_INSTALL=0 \ PATH=/opt/cargo/bin:${PATH} @@ -70,12 +73,14 @@ RUN curl --fail --location --silent --show-error \ && install --mode=0755 /tmp/sccache/sccache /usr/local/bin/sccache \ && rm -rf /tmp/sccache.tar.gz /tmp/sccache -RUN curl --fail --location --silent --show-error \ +RUN rust_toolchain="$(grep -E '^channel' /opt/rust-toolchain.toml | cut -d'"' -f2)" \ + && test -n "${rust_toolchain}" \ + && curl --fail --location --silent --show-error \ --output /tmp/rustup-init \ "https://static.rust-lang.org/rustup/archive/${RUSTUP_INIT_VERSION}/${RUSTUP_INIT_TARGET}/rustup-init" \ && echo "${RUSTUP_INIT_SHA256} /tmp/rustup-init" | sha256sum --check --status - \ && chmod 0755 /tmp/rustup-init \ - && /tmp/rustup-init -y --profile minimal --default-toolchain "${RUST_TOOLCHAIN}" --no-modify-path \ + && /tmp/rustup-init -y --profile minimal --default-toolchain "${rust_toolchain}" --no-modify-path \ && rm -f /tmp/rustup-init \ && chown -R root:root /opt/rustup /opt/cargo \ && chmod -R go-w /opt/rustup /opt/cargo \ @@ -86,6 +91,8 @@ USER runner RUN set -eu \ && test "$(id -u)" = "1001" \ + && rust_toolchain="$(grep -E '^channel' /opt/rust-toolchain.toml | cut -d'"' -f2)" \ + && test -n "${rust_toolchain}" \ && temporary_cargo_home="$(mktemp --directory)" \ && temporary_proto_dir="" \ && cleanup() { rm -rf "${temporary_cargo_home}" "${temporary_proto_dir}"; } \ @@ -97,7 +104,7 @@ RUN set -eu \ && rustc --version \ && active_toolchain="$(rustup show active-toolchain)" \ && printf 'Active toolchain: %s\n' "${active_toolchain}" \ - && case "${active_toolchain}" in "${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu"|"${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ + && case "${active_toolchain}" in "${rust_toolchain}-x86_64-unknown-linux-gnu"|"${rust_toolchain}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ && test ! -w /opt/rustup \ && test ! -w /opt/cargo/bin \ && test -r /usr/include/google/protobuf/any.proto \ diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index fa8d436d2c2..70ed4a88345 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -86,15 +86,30 @@ sweep. ### Building and updating the ECS image The runner `Dockerfile` in the parent directory stays the single source of the -tool contract. Build a new ECS image from it: +tool contract. The image is rebuilt **automatically** by the +`rebuild-query-regression-runner-image` job in +`.github/workflows/release-dev-builder-images.yaml` whenever +`rust-toolchain.toml` or anything under this directory changes on main (or via +manual dispatch): it runs the ops tool below, then bumps +`RUNNER_IMAGE_EPOCH` in `query-regression.yml` and points the +`QUERY_REGRESSION_ECS_IMAGE_ID` repo variable at the new image, so the next +regression run picks up image and epoch together. + +The manual fallback (also what the workflow runs): ```bash ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ --region-id --vswitch-id --security-group-id \ - --base-image-id + [--base-image-id ] ``` +`--base-image-id` is optional: the script defaults to the latest public +Ubuntu 24.04 image in the region (the Dockerfile pins every tool version +itself, so base drift is low-risk); pass it — or set the +`ALIYUN_ECS_BASE_IMAGE_ID` repo variable consumed by the automated job — +to pin a specific base image. + The script boots a temporary builder instance, `docker build`s the runner image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and `/home/runner` (actions-runner) onto the host, installs the ephemeral-runner @@ -158,10 +173,17 @@ overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, `Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and -the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` -come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump -the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the -image is rebuilt. Rustup, Cargo, and Rustc +the image-baked Rust toolchain matching `rust-toolchain.toml` (the image +parses the pin from the toml at build time, and the workflow asserts it +dynamically at run time — there is no separately pinned toolchain version). +`mold` and `python3` +come from apt at image-build time (not Ubuntu 24.04's default 3.12); if the +Ubuntu archive ships a newer package revision between rebuilds, the Verify +step fails with the observed version — bump those pins in `query-regression.yml` +when that happens. Everything else (toolchain, uv, sccache, otelgen, rustup, +the runner base) is pinned by digest/sha/commit in the Dockerfile, and the +image id + `RUNNER_IMAGE_EPOCH` updates are automated by the rebuild job, so +a routine rebuild needs no manual file updates. Rustup, Cargo, and Rustc must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. Protobuf well-known includes, including `google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py index 7031917f910..78b01e721c3 100644 --- a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -52,6 +52,8 @@ import time from pathlib import Path ASSETS_DIR = Path(__file__).resolve().parent +# Repo root: ecs-image -> query-regression -> runner-scale-sets -> .github -> root. +REPO_ROOT = ASSETS_DIR.parents[3] DONE_MARKER = "QREG_IMAGE_BUILD_DONE" FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" POLL_INTERVAL_SECONDS = 15 @@ -92,8 +94,9 @@ APT_PACKAGES = [ DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" -def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: +def render_user_data(dockerfile: str, rust_toolchain_toml: str, start_runner: str, unit: str) -> str: dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() + rust_toolchain_b64 = base64.b64encode(rust_toolchain_toml.encode()).decode() start_runner_b64 = base64.b64encode(start_runner.encode()).decode() unit_b64 = base64.b64encode(unit.encode()).decode() packages = " ".join(APT_PACKAGES) @@ -122,6 +125,11 @@ base64 -d > /tmp/Dockerfile <<'EOF' {dockerfile_b64} EOF mkdir -p /tmp/image-context +# The Dockerfile COPYies rust-toolchain.toml (single source of truth for +# the baked toolchain); stage it into the build context. +base64 -d > /tmp/image-context/rust-toolchain.toml <<'EOF' +{rust_toolchain_b64} +EOF docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context # Materialize the tool contract onto the host filesystem. @@ -229,6 +237,62 @@ def call_api_with_retry(fn, description: str, attempts: int = 5): raise RuntimeError("unreachable: retry loop exited without returning") +def resolve_base_image_id(client, region_id: str) -> str: + """Resolve the latest public Ubuntu 24.04 x86_64 system image. + + Used as the default for --base-image-id: the runner Dockerfile pins + every tool version itself, so a current stock Ubuntu 24.04 base is all + the builder needs. Pass --base-image-id (or ALIYUN_ECS_BASE_IMAGE_ID) + to pin a specific base image deterministically. + """ + from alibabacloud_ecs20140526 import models as ecs_models + + def _is_ubuntu_2404(image) -> bool: + # osname is localized (e.g. "Ubuntu 24.04 64ä½"), osname_en the + # English form; accept either. + for os_name in (image.osname_en, image.osname): + if os_name and "Ubuntu" in os_name and "24.04" in os_name: + return True + return False + + page_size = 100 + images: list = [] + page_number = 1 + while True: + page = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest( + region_id=region_id, + image_owner_alias="system", + ostype="linux", + architecture="x86_64", + page_number=page_number, + page_size=page_size, + ) + ), + f"DescribeImages(system base, page {page_number})", + ).body.images.image or [] + images.extend(page) + if len(page) < page_size: + break + page_number += 1 + + candidates = [image for image in images if _is_ubuntu_2404(image)] + if not candidates: + raise SystemExit( + "No public Ubuntu 24.04 x86_64 system image found in region " + f"{region_id}; pass --base-image-id explicitly" + ) + candidates.sort(key=lambda image: image.creation_time or "", reverse=True) + picked = candidates[0] + print( + f"Resolved base image: {picked.image_id} ({picked.osname_en or picked.osname}, " + f"created {picked.creation_time}) from {len(candidates)} candidates", + flush=True, + ) + return picked.image_id + + def read_console_output(client, region_id: str, instance_id: str) -> str: """Fetch the instance's serial console output; no in-guest agent needed.""" from alibabacloud_ecs20140526 import models as ecs_models @@ -253,13 +317,19 @@ def main() -> int: parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") args = parser.parse_args() - for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + for name in ("region_id", "vswitch_id", "security_group_id"): if not getattr(args, name): raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") from alibabacloud_ecs20140526 import models as ecs_models client = make_ecs_client(args.region_id) + # --base-image-id is optional: default to the latest public Ubuntu 24.04 + # image in the region (the Dockerfile pins every tool version itself, so + # base drift is low-risk; pass --base-image-id or set + # ALIYUN_ECS_BASE_IMAGE_ID to pin deterministically). + if not args.base_image_id: + args.base_image_id = resolve_base_image_id(client, args.region_id) image_name = args.image_name or time.strftime( "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() ) @@ -267,6 +337,7 @@ def main() -> int: user_data = base64.b64encode( render_user_data( (ASSETS_DIR.parent / "Dockerfile").read_text(), + (REPO_ROOT / "rust-toolchain.toml").read_text(), (ASSETS_DIR / "start-runner.sh").read_text(), (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), ).encode() diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 8c9a93789cf..8ce77827de7 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -120,12 +120,23 @@ jobs: # Ordinary PRs also run the same tests from checks.yml. runs-on: ubuntu-latest timeout-minutes: 10 + outputs: + # The Rust toolchain pin, resolved from rust-toolchain.toml so the + # benchmark always runs the workspace toolchain without a second + # hard-coded copy in this workflow. + rust_toolchain: ${{ steps.rust-toolchain.outputs.pin }} steps: - name: Checkout uses: actions/checkout@v4 with: persist-credentials: false + - name: Resolve Rust toolchain pin + id: rust-toolchain + run: | + pin="$(grep -E '^channel' rust-toolchain.toml | cut -d'"' -f2)" + echo "pin=${pin}" >> "$GITHUB_OUTPUT" + - name: Test query regression tooling run: | python3 tests/perf/test_query_regression_runner_compaction_toctou.py @@ -190,7 +201,7 @@ jobs: CARGO_HOME: /home/runner/.cargo UV_CACHE_DIR: /home/runner/.cargo/uv-cache RUSTUP_HOME: /opt/rustup - RUSTUP_TOOLCHAIN: nightly-2026-03-21 + RUSTUP_TOOLCHAIN: ${{ needs.test-tooling.outputs.rust_toolchain }} RUSTUP_AUTO_INSTALL: "0" CARGO_TARGET_DIR: /home/runner/query-regression-target QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta @@ -504,14 +515,16 @@ jobs: require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + # The toolchain pin flows from rust-toolchain.toml (resolved in the + # test-tooling job); escape it for the regex assertions below. + pin_regex="${RUSTUP_TOOLCHAIN//./\\.}" require_match cargo "$(capture cargo --version)" \ - '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + "^cargo[[:space:]]${pin_regex}[[:space:]]\([0-9a-f]+[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$" require_match rustc "$(capture rustc --version)" \ - '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + "^rustc[[:space:]]${pin_regex}[[:space:]]\([0-9a-f]+[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$" require_match active_toolchain "$(capture rustup show active-toolchain)" \ - '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' + "^${pin_regex}-x86_64-unknown-linux-gnu([[:space:]]|$)" require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" - require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" require "readable /opt/rustup" test -r /opt/rustup require "executable /opt/rustup" test -x /opt/rustup diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index 41325f61ab9..f9a55cb9985 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -1,4 +1,4 @@ -name: Release dev-builder images +name: Release builder images on: push: @@ -7,6 +7,7 @@ on: paths: - rust-toolchain.toml - 'docker/dev-builder/**' + - '.github/runner-scale-sets/query-regression/**' workflow_dispatch: # Allows you to run this workflow manually. inputs: release_dev_builder_ubuntu_image: @@ -29,6 +30,11 @@ on: description: Release dev-builder-riscv64 image required: false default: false + release_query_regression_runner_image: + type: boolean + description: Rebuild the query-regression ECS runner image (and update QUERY_REGRESSION_ECS_IMAGE_ID + RUNNER_IMAGE_EPOCH) + required: false + default: false update_dev_builder_image_tag: type: boolean description: Update the DEV_BUILDER_IMAGE_TAG in Makefile and create a PR @@ -36,12 +42,55 @@ on: default: false jobs: + changes: + name: Detect changed components + runs-on: ubuntu-latest + outputs: + dev-builder: ${{ steps.filter.outputs.dev-builder }} + query-regression-runner: ${{ steps.filter.outputs.query-regression-runner }} + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + - name: Filter changes + id: filter + shell: bash + run: | + set -euo pipefail + base="${{ github.event.before }}" + head="${{ github.sha }}" + if [[ "${{ github.event_name }}" != "push" ]] || [[ -z "${base}" ]] || ! git cat-file -e "${base}^{commit}" 2>/dev/null; then + # Dispatch, or an unknown base (new branch/force-push): treat + # everything as changed so gating falls back to job inputs. + echo "dev-builder=true" >> "$GITHUB_OUTPUT" + echo "query-regression-runner=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + files="$(git diff --name-only "${base}" "${head}")" + dev_builder=false + query_regression_runner=false + if grep -qx 'rust-toolchain.toml' <<<"${files}"; then + dev_builder=true + query_regression_runner=true + fi + if grep -q '^docker/dev-builder/' <<<"${files}"; then + dev_builder=true + fi + if grep -q '^\.github/runner-scale-sets/query-regression/' <<<"${files}"; then + query_regression_runner=true + fi + echo "dev-builder=${dev_builder}" >> "$GITHUB_OUTPUT" + echo "query-regression-runner=${query_regression_runner}" >> "$GITHUB_OUTPUT" + release-dev-builder-images: name: Release dev builder images # The jobs are triggered by the following events: # 1. Manually triggered workflow_dispatch event # 2. Push event when the PR that modifies the `rust-toolchain.toml` or `docker/dev-builder/**` is merged to main - if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} + needs: [changes] + if: ${{ (github.event_name == 'push' && needs.changes.outputs.dev-builder == 'true') || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} runs-on: ubuntu-latest outputs: version: ${{ steps.set-version.outputs.version }} @@ -281,3 +330,96 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | ./.github/scripts/update-dev-builder-version.sh ${{ needs.release-dev-builder-images.outputs.version }} + + rebuild-query-regression-runner-image: + # Rebuilds the query-regression ECS runner image via the ops tool in + # .github/runner-scale-sets/query-regression/ecs-image/: boots a + # temporary pay-as-you-go ECS builder, snapshots a new custom image, + # then completes the documented lockstep updates in order: + # 1. bumps RUNNER_IMAGE_EPOCH in query-regression.yml (target-cache + # invalidation) and pushes the commit to main, and only then + # 2. points the repo variable QUERY_REGRESSION_ECS_IMAGE_ID at the + # new image, so the next regression run picks up image and epoch + # together. + # + # Required repository configuration (same names the provisioning job + # uses): vars ALIYUN_ECS_REGION_ID, ALIYUN_ECS_VSWITCH_ID, + # ALIYUN_ECS_SECURITY_GROUP_ID, optionally ALIYUN_ECS_RESOURCE_GROUP_ID + # and ALIYUN_ECS_BASE_IMAGE_ID (deterministic base-image pin; otherwise + # the rebuild auto-resolves the latest public Ubuntu 24.04 image). + # Secrets: ALICLOUD_ECS_ACCESS_KEY_ID, ALICLOUD_ECS_ACCESS_KEY_SECRET, + # GH_PERSONAL_ACCESS_TOKEN (repo push + actions-variable write). + name: Rebuild query-regression runner image + needs: [changes] + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && ((github.event_name == 'push' && needs.changes.outputs.query-regression-runner == 'true') || inputs.release_query_regression_runner_image) }} + runs-on: ubuntu-latest + timeout-minutes: 150 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + # Keep the PAT credential for the epoch-bump push below. + persist-credentials: true + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Rebuild the runner image + id: rebuild + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + run: | + set -euo pipefail + args=( + --region-id "${{ vars.ALIYUN_ECS_REGION_ID }}" + --vswitch-id "${{ vars.ALIYUN_ECS_VSWITCH_ID }}" + --security-group-id "${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }}" + --base-image-id "${{ vars.ALIYUN_ECS_BASE_IMAGE_ID }}" + ) + if [[ -n "${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }}" ]]; then + args+=(--resource-group-id "${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }}") + fi + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + "${args[@]}" 2>&1 | tee rebuild.log + image_id="$(sed -n 's/^Custom image ready: \([^ ]*\) .*/\1/p' rebuild.log | tail -n1)" + if [[ -z "${image_id}" ]]; then + echo "Error: no image id found in build output" >&2 + exit 1 + fi + echo "image_id=${image_id}" >> "$GITHUB_OUTPUT" + echo "## Rebuilt runner image" >> "$GITHUB_STEP_SUMMARY" + echo "New image: \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" + + - name: Bump RUNNER_IMAGE_EPOCH and update the image id variable + env: + GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + run: | + set -euo pipefail + image_id="${{ steps.rebuild.outputs.image_id }}" + old_epoch="$(grep -oE 'readonly RUNNER_IMAGE_EPOCH="[0-9]+"' \ + .github/workflows/query-regression.yml | grep -oE '[0-9]+' | head -n1)" + if [[ -z "${old_epoch}" ]]; then + echo "Error: RUNNER_IMAGE_EPOCH not found in query-regression.yml" >&2 + exit 1 + fi + new_epoch=$((old_epoch + 1)) + sed -i "s/readonly RUNNER_IMAGE_EPOCH=\"${old_epoch}\"/readonly RUNNER_IMAGE_EPOCH=\"${new_epoch}\"/" \ + .github/workflows/query-regression.yml + + git config user.name "greptimedb-ci" + git config user.email "greptimedb-ci@users.noreply.github.com" + git add .github/workflows/query-regression.yml + git commit -m "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch} for image ${image_id} [skip ci]" + git pull --rebase origin main + git push origin HEAD:main + + # Only point the variable at the new image after the epoch commit + # landed, so the next run picks up image and epoch together. + gh api -X PATCH "repos/${{ github.repository }}/actions/variables/QUERY_REGRESSION_ECS_IMAGE_ID" \ + -f value="${image_id}" + + echo "## Lockstep updates" >> "$GITHUB_STEP_SUMMARY" + echo "- QUERY_REGRESSION_ECS_IMAGE_ID → \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" + echo "- RUNNER_IMAGE_EPOCH → \`${new_epoch}\` (committed to main)" >> "$GITHUB_STEP_SUMMARY" diff --git a/AGENTS.md b/AGENTS.md index ef56c326805..20b2689a401 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,7 +25,7 @@ provides sub-second querying at PB scale with high cost efficiency. | Lint | `make clippy` (= `cargo clippy --workspace --all-targets --all-features -- -D warnings`) | | Type check | `make check` | -Toolchain: Rust nightly, Protobuf compiler (>= 3.15), C/C++ build essentials. +Toolchain: Rust (stable, pinned by `rust-toolchain.toml`), Protobuf compiler (>= 3.15), C/C++ build essentials. Install the test runner with `cargo install cargo-nextest --locked`. ## Repo map diff --git a/README.md b/README.md index 1c8ed985e1d..96058cc79b7 100644 --- a/README.md +++ b/README.md @@ -168,7 +168,7 @@ Read more in the [full Install Guide](https://docs.greptime.com/getting-started/ ## Build From Source **Prerequisites:** -* [Rust toolchain](https://www.rust-lang.org/tools/install) — nightly, pinned by [`rust-toolchain.toml`](https://github.com/GreptimeTeam/greptimedb/blob/main/rust-toolchain.toml) +* [Rust toolchain](https://www.rust-lang.org/tools/install) — stable, pinned by [`rust-toolchain.toml`](https://github.com/GreptimeTeam/greptimedb/blob/main/rust-toolchain.toml) * [Protobuf compiler](https://grpc.io/docs/protoc-installation/) (>= 3.15) * C/C++ building essentials: `gcc` / `g++` / `autoconf` and the glibc dev package (`libc6-dev` on Ubuntu, `glibc-devel` on Fedora) * Python toolchain (optional, only for some test scripts) diff --git a/scripts/check-builder-rust-version.sh b/scripts/check-builder-rust-version.sh index fcfbd7d189d..3ca7acf9f84 100755 --- a/scripts/check-builder-rust-version.sh +++ b/scripts/check-builder-rust-version.sh @@ -16,26 +16,50 @@ function check_rust_toolchain_version() { DEV_BUILDER_UBUNTU_IMAGE="$DEV_BUILDER_UBUNTU_REGISTRY/$DEV_BUILDER_UBUNTU_NAMESPACE/$DEV_BUILDER_UBUNTU_NAME:$DEV_BUILDER_IMAGE_TAG" - CURRENT_VERSION=$(grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}' "$RUST_TOOLCHAIN_VERSION_FILE") - if [ -z "$CURRENT_VERSION" ]; then - echo "Error: No rust toolchain version found in $RUST_TOOLCHAIN_VERSION_FILE" + # The pinned channel, e.g. "1.96.1" or "nightly-2026-03-21". + RUST_TOOLCHAIN=$(grep -E '^channel' "$RUST_TOOLCHAIN_VERSION_FILE" | cut -d'"' -f2) + if [ -z "$RUST_TOOLCHAIN" ]; then + echo "Error: No rust toolchain channel found in $RUST_TOOLCHAIN_VERSION_FILE" exit 1 fi - RUST_TOOLCHAIN_VERSION_IN_BUILDER=$(docker run "$DEV_BUILDER_UBUNTU_IMAGE" rustc --version | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') - if [ -z "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" ]; then + RUSTC_VERSION_OUTPUT=$(docker run "$DEV_BUILDER_UBUNTU_IMAGE" rustc --version) + RUSTC_VERSION=$(echo "$RUSTC_VERSION_OUTPUT" | awk '{print $2}') + if [ -z "$RUSTC_VERSION" ]; then echo "Error: No rustc version found in $DEV_BUILDER_UBUNTU_IMAGE" exit 1 fi - # Compare the version and the difference should be less than 1 day. - current_rust_toolchain_seconds=$(date -d "$CURRENT_VERSION" +%s) - rust_toolchain_in_dev_builder_ubuntu_seconds=$(date -d "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" +%s) - date_diff=$(( (current_rust_toolchain_seconds - rust_toolchain_in_dev_builder_ubuntu_seconds) / 86400 )) + if [[ "$RUST_TOOLCHAIN" =~ ^[0-9]+\.[0-9]+ ]]; then + # Stable channel: the builder image must ship exactly the pinned version. + if [ "$RUSTC_VERSION" != "$RUST_TOOLCHAIN" ]; then + echo "Error: The rust toolchain '$RUSTC_VERSION' in builder '$DEV_BUILDER_UBUNTU_IMAGE' doesn't match the pinned '$RUST_TOOLCHAIN', please rebuild the dev-builder image" + exit 1 + fi + else + # Nightly channel: keep the legacy date-based check. The difference + # between the pinned nightly date and the rustc build date shipped in + # the builder should be less than 1 day. + CURRENT_VERSION=$(echo "$RUST_TOOLCHAIN" | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + if [ -z "$CURRENT_VERSION" ]; then + echo "Error: Unsupported toolchain channel '$RUST_TOOLCHAIN' in $RUST_TOOLCHAIN_VERSION_FILE" + exit 1 + fi - if [ $date_diff -gt 1 ]; then - echo "Error: The rust toolchain '$RUST_TOOLCHAIN_VERSION_IN_BUILDER' in builder '$DEV_BUILDER_UBUNTU_IMAGE' maybe outdated, please update it to '$CURRENT_VERSION'" - exit 1 + RUST_TOOLCHAIN_VERSION_IN_BUILDER=$(echo "$RUSTC_VERSION_OUTPUT" | grep -Eo '[0-9]{4}-[0-9]{2}-[0-9]{2}') + if [ -z "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" ]; then + echo "Error: No rustc version found in $DEV_BUILDER_UBUNTU_IMAGE" + exit 1 + fi + + current_rust_toolchain_seconds=$(date -d "$CURRENT_VERSION" +%s) + rust_toolchain_in_dev_builder_ubuntu_seconds=$(date -d "$RUST_TOOLCHAIN_VERSION_IN_BUILDER" +%s) + date_diff=$(( (current_rust_toolchain_seconds - rust_toolchain_in_dev_builder_ubuntu_seconds) / 86400 )) + + if [ $date_diff -gt 1 ]; then + echo "Error: The rust toolchain '$RUST_TOOLCHAIN_VERSION_IN_BUILDER' in builder '$DEV_BUILDER_UBUNTU_IMAGE' maybe outdated, please update it to '$CURRENT_VERSION'" + exit 1 + fi fi } From 0a03e64d8cc5fd593e4bde56dfe6a494556d616d Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 06:22:15 +0000 Subject: [PATCH 051/133] perf(index): batch bloom filter searches across row groups (#9361) * perf(index): search bloom filters of all row groups in one read Signed-off-by: Dennis Zhuang * perf(index): bound the filter bytes of each batched bloom search Signed-off-by: Dennis Zhuang * perf(index): count batch filter bytes exactly as the search reads them Signed-off-by: Dennis Zhuang * chore(index): allow single-range groups in the budget test Signed-off-by: Dennis Zhuang * test(index): cover filters shared across segments in the budget test Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/index/src/bloom_filter/applier.rs | 308 ++++++++++++++++++ .../src/sst/index/bloom_filter/applier.rs | 21 +- .../src/sst/index/fulltext_index/applier.rs | 26 +- 3 files changed, 323 insertions(+), 32 deletions(-) diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index db219c9e611..ebd44cf84e6 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -23,6 +23,11 @@ use crate::Bytes; use crate::bloom_filter::error::Result; use crate::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; +/// Filter bytes one batch of [`BloomFilterApplier::search_groups`] reads. A single row +/// group larger than this is still read as one batch, so this is not a memory limit; a +/// batch holds both its raw bytes and the decoded filters. +const MAX_BATCH_FILTER_BYTES: u64 = 8 * 1024 * 1024; + /// `InListPredicate` contains a list of acceptable values. A value needs to match at least /// one of the elements (logical OR semantic) for the predicate to be satisfied. #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -43,6 +48,93 @@ impl BloomFilterApplier { Ok(Self { reader, meta }) } + /// Runs [`Self::search`] over several groups of ranges (e.g. row groups) and keeps in + /// each group only its matching ranges. + /// + /// Consecutive groups are searched together, one read per batch, as long as the + /// filters of a batch stay within [`MAX_BATCH_FILTER_BYTES`]. This saves a round trip + /// per group on object storage without holding the filters of a whole file at once. + /// Groups must be ordered and their ranges sorted and disjoint, as for `search`. + pub async fn search_groups( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + self.search_groups_in_batches(predicates, groups, metrics, MAX_BATCH_FILTER_BYTES) + .await + } + + async fn search_groups_in_batches( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + mut metrics: Option<&mut BloomFilterReadMetrics>, + max_batch_bytes: u64, + ) -> Result<()> { + let mut start = 0; + while start < groups.len() { + // Mirrors `load_bloom_filters`: segments map to filter locations in order and + // only consecutive equal locations share a read, so this is the exact number of + // bytes the batch requests. + let mut last_loc = None; + let mut batch_bytes = 0; + let mut end = start; + while end < groups.len() { + let mut group_last = last_loc; + let mut bytes = 0; + for seg in self.row_ranges_to_segments(groups[end]) { + let loc = self.meta.segment_loc_indices[seg]; + if group_last != Some(loc) { + bytes += self.meta.bloom_filter_locs[loc as usize].size; + group_last = Some(loc); + } + } + // A group whose filters alone exceed the budget still forms its own batch. + if end > start && batch_bytes + bytes > max_batch_bytes { + break; + } + last_loc = group_last; + batch_bytes += bytes; + end += 1; + } + self.search_batch(predicates, &mut groups[start..end], metrics.as_deref_mut()) + .await?; + start = end; + } + Ok(()) + } + + /// Searches `groups` with one `search`, i.e. one read for all their filters. + async fn search_batch( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + let all = groups + .iter() + .flat_map(|g| g.iter().cloned()) + .collect::>(); + if all.is_empty() { + return Ok(()); + } + // Each matched range lies within one input range, so it belongs to one group. + let mut matched = self + .search(predicates, &all, metrics) + .await? + .into_iter() + .peekable(); + for group in groups.iter_mut() { + for range in std::mem::take(*group) { + while let Some(m) = matched.next_if(|m| m.start < range.end) { + group.push(m); + } + } + } + Ok(()) + } + /// Searches ranges of rows that match all the given predicates in the search ranges. /// Each predicate represents an OR condition of probes, and all predicates must match (AND semantics). /// The logic is: (probe1 OR probe2 OR ...) AND (probe3 OR probe4 OR ...) @@ -217,6 +309,222 @@ mod tests { use crate::bloom_filter::reader::BloomFilterReaderImpl; use crate::external_provider::MockExternalTempFileProvider; + #[tokio::test] + async fn test_search_groups_matches_per_group_search() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Row i holds "v{i / 3}", so values straddle segment boundaries. + for i in 0..40 { + creator + .push_row_elems([format!("v{}", i / 3).into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + + // Row groups of 10 rows, some already narrowed by other predicates. + let groups = vec![ + vec![0..10], + vec![10..13, 15..20], + vec![], + vec![30..33, 37..40], + ]; + for values in [ + vec!["v1"], + vec!["v4", "v5"], + vec!["v3", "v10", "v12"], + vec!["x"], + ] { + let predicates = vec![InListPredicate { + list: values.iter().map(|v| v.as_bytes().to_vec()).collect(), + }]; + let mut applier = + BloomFilterApplier::new(Box::new(BloomFilterReaderImpl::new(bytes.clone()))) + .await + .unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + let mut expected = Vec::new(); + for group in &groups { + expected.push(if group.is_empty() { + vec![] + } else { + applier.search(&predicates, group, None).await.unwrap() + }); + } + // Budgets of zero (a batch per group), one filter and unlimited (one batch). + for budget in [0, filter_size, u64::MAX] { + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(actual, expected, "values: {values:?}, budget: {budget}"); + } + } + } + + /// Records the bytes of every `read_vec`, i.e. of every batch. + struct RecordingReader { + inner: BloomFilterReaderImpl>, + reads: Arc>>, + } + + #[async_trait::async_trait] + impl BloomFilterReader for RecordingReader { + async fn range_read( + &self, + offset: u64, + size: u32, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + self.inner.range_read(offset, size, metrics).await + } + + async fn read_vec( + &self, + ranges: &[Range], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + let bytes = ranges.iter().map(|r| r.end - r.start).sum(); + self.reads.lock().unwrap().push(bytes); + self.inner.read_vec(ranges, metrics).await + } + + async fn metadata( + &self, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + self.inner.metadata(metrics).await + } + } + + #[tokio::test] + #[allow(clippy::single_range_in_vec_init)] + async fn test_search_groups_respects_budget() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Distinct values everywhere: 10 segments with filters of the same size. + for i in 0..40 { + creator + .push_row_elems([format!("v{i}").into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + let predicates = vec![InListPredicate { + list: BTreeSet::from([b"v1".to_vec()]), + }]; + + let reads = Arc::new(std::sync::Mutex::new(Vec::new())); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(bytes), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + + // Row groups of 6 rows: most share a boundary segment with their neighbor, which + // a batch reads only once. + let groups = (0..40) + .step_by(6) + .map(|s| vec![s..(s + 6).min(40)]) + .collect::>(); + for (budget, expected_reads) in [ + (u64::MAX, vec![10 * filter_size]), + // Batches close before exceeding the budget; shared boundary filters count once. + ( + 5 * filter_size, + vec![5 * filter_size, 5 * filter_size, filter_size], + ), + // Smaller than any row group: one row group per batch, still above the budget. + ( + filter_size, + vec![ + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + filter_size, + ], + ), + ] { + reads.lock().unwrap().clear(); + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(*reads.lock().unwrap(), expected_reads, "budget: {budget}"); + } + + // Segments 0..3 hold the same value and share one filter, which a batch counts + // and reads once even across row groups. + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + for i in 0..24 { + let value = if i < 12 { + "a".to_string() + } else { + format!("v{i}") + }; + creator.push_row_elems([value.into_bytes()]).await.unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(writer.into_inner()), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + assert_eq!(applier.meta.bloom_filter_locs.len(), 4); + let filter_size = applier.meta.bloom_filter_locs[0].size; + assert!( + applier + .meta + .bloom_filter_locs + .iter() + .all(|l| l.size == filter_size) + ); + reads.lock().unwrap().clear(); + let mut groups = (0..24) + .step_by(6) + .map(|s| vec![s..s + 6]) + .collect::>(); + let mut refs = groups.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, 2 * filter_size) + .await + .unwrap(); + // Row groups 0 and 1 need only the shared filter; 2 and 3 need two each. + assert_eq!( + *reads.lock().unwrap(), + vec![filter_size, 2 * filter_size, 2 * filter_size] + ); + } + #[tokio::test] #[allow(clippy::single_range_in_vec_init)] async fn test_appliter() { diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index e2d25e4bf82..1faae466c89 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -422,23 +422,14 @@ impl BloomFilterIndexApplier { reader: R, predicates: &[InListPredicate], output: &mut [(usize, Vec>)], - mut metrics: Option<&mut BloomFilterIndexApplyMetrics>, + metrics: Option<&mut BloomFilterIndexApplyMetrics>, ) -> std::result::Result<(), index::bloom_filter::error::Error> { let mut applier = BloomFilterApplier::new(Box::new(reader)).await?; - - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - let read_metrics = metrics.as_deref_mut().map(|m| &mut m.read_metrics); - *row_group_output = applier - .search(predicates, row_group_output, read_metrics) - .await?; - } - - Ok(()) + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + let read_metrics = metrics.map(|m| &mut m.read_metrics); + applier + .search_groups(predicates, &mut row_groups, read_metrics) + .await } /// Returns compatible bloom filter predicates with the given SST metadata. diff --git a/src/mito2/src/sst/index/fulltext_index/applier.rs b/src/mito2/src/sst/index/fulltext_index/applier.rs index 5ccf4e57a4b..4efb2a5ccb9 100644 --- a/src/mito2/src/sst/index/fulltext_index/applier.rs +++ b/src/mito2/src/sst/index/fulltext_index/applier.rs @@ -464,23 +464,15 @@ impl FulltextIndexApplier { let mut applier = BloomFilterApplier::new(reader) .await .context(ApplyBloomFilterIndexSnafu)?; - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - *row_group_output = applier - .search( - &predicates, - row_group_output, - metrics - .as_deref_mut() - .map(|m| &mut m.bloom_filter_read_metrics), - ) - .await - .context(ApplyBloomFilterIndexSnafu)?; - } + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + applier + .search_groups( + &predicates, + &mut row_groups, + metrics.map(|m| &mut m.bloom_filter_read_metrics), + ) + .await + .context(ApplyBloomFilterIndexSnafu)?; Ok(true) } From 55b7e08a1a2aa01dfe57efb7cad0131eabd7e459 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 28 Sep 2026 07:15:36 +0000 Subject: [PATCH 052/133] feat: allow customized time index unit for metric engine table (#9236) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: allow customized time index unit for metric engine table * test: provide query tests * refactor: revert unnecessary change * refactor: share timestamp unit conversions in api helper Address review feedback on the time index unit changeset: - Add shared timestamp_unit/timestamp_datatype helpers to api::helper (the only crate that sees both proto ColumnDataType and TimeUnit due to layering; common-time and datatypes have no greptime-proto dep). This removes the ColumnDataType -> TimeUnit match duplicated between operator's insert path and the OTLP logs path. - Collapse the two TimeUnit <-> ValueData matches in convert_timestamp_value_data by reusing api::helper::to_grpc_value for the construction side. - Note that convert_rows_time_unit rewrites the schema before the values, so an overflow mid-batch leaves the request half-converted; harmless because the error aborts the whole insert request. Signed-off-by: Ning Sun * fix: align time units per destination table and floor remote-read timestamps Address review feedback on PR #9236: - Align each metric insert request to the unit of the table it actually targets: an existing logical table keeps its own unit (it may be bound to a different physical table than the one selected by the request), and only new tables use the selected physical table's unit. The previous blanket conversion rewrote valid millisecond samples to the selected physical table's unit and the engine rejected them. Regression test: writing an existing millisecond logical table and a new table in one request that selects a microsecond physical table. - Remote read now floors narrowing timestamp conversions towards negative infinity (div_euclid), consistent with Timestamp::convert_to on the ingestion path; arrow's cast truncates towards zero and returned -1ms for a stored -1001us. Widening (second -> millisecond) keeps the exact arrow cast. Regression test: a negative, non-aligned timestamp round-trips as -2ms. Signed-off-by: Ning Sun * perf: fold time unit alignment into existing table lookups Address review feedback on PR #9236: - The per-destination unit alignment no longer runs its own pass of table lookups: create_or_alter_tables_on_demand gains an align_time_index_unit parameter (metric engine path only) and converts each request inside the lookups it already performs — existing tables to their own unit, new tables to the selected physical table's. Default ingest paths now issue zero additional catalog lookups compared to main; the separate alignment pass remains only in the opt-in logical batcher pre-gate, next to the eligibility check that already looks up the same tables. - convert_rows_time_unit indexes the time index position directly (validate_column_count_match guarantees row widths) instead of Optional get_mut; the gate-side alignment validates widths itself. Signed-off-by: Ning Sun * perf: resolve the batcher time index guard once per write target All batches of one remote write request share the same write target (catalog, schema, physical table), so the batcher time index guard now resolves each distinct target once instead of once per batch. Signed-off-by: Ning Sun * fix: address review comments * fix: address review issue --------- Signed-off-by: Ning Sun --- src/api/src/helper.rs | 37 ++ src/frontend/src/instance/otlp.rs | 25 +- src/metric-engine/src/engine/put.rs | 89 ++++- src/metric-engine/src/test_util.rs | 82 +++- src/operator/src/insert.rs | 478 +++++++++++++++++++++- src/servers/src/batcher/logical_table.rs | 98 ++++- src/servers/src/http/prom_store.rs | 182 +++++++-- src/servers/src/otlp/logs.rs | 12 +- src/servers/src/otlp/metrics.rs | 88 ++++- src/servers/src/prom_store.rs | 53 ++- tests-integration/src/otlp.rs | 161 ++++++++ tests-integration/src/prom_store.rs | 438 +++++++++++++++++++- tests-integration/tests/http.rs | 482 ++++++++++++++++++++++- 13 files changed, 2128 insertions(+), 97 deletions(-) diff --git a/src/api/src/helper.rs b/src/api/src/helper.rs index 40f9bdfb6d6..c23ae2587d5 100644 --- a/src/api/src/helper.rs +++ b/src/api/src/helper.rs @@ -87,6 +87,28 @@ impl ColumnDataTypeWrapper { } } +/// Returns the time unit if `datatype` is a timestamp type. +pub fn timestamp_unit(datatype: ColumnDataType) -> Option { + match datatype { + ColumnDataType::TimestampSecond => Some(TimeUnit::Second), + ColumnDataType::TimestampMillisecond => Some(TimeUnit::Millisecond), + ColumnDataType::TimestampMicrosecond => Some(TimeUnit::Microsecond), + ColumnDataType::TimestampNanosecond => Some(TimeUnit::Nanosecond), + _ => None, + } +} + +/// Returns the timestamp [ColumnDataType] for the given time unit. +/// This is the inverse of [timestamp_unit]. +pub fn timestamp_datatype(unit: TimeUnit) -> ColumnDataType { + match unit { + TimeUnit::Second => ColumnDataType::TimestampSecond, + TimeUnit::Millisecond => ColumnDataType::TimestampMillisecond, + TimeUnit::Microsecond => ColumnDataType::TimestampMicrosecond, + TimeUnit::Nanosecond => ColumnDataType::TimestampNanosecond, + } +} + impl From for ConcreteDataType { fn from(datatype_wrapper: ColumnDataTypeWrapper) -> Self { match datatype_wrapper.datatype { @@ -1215,6 +1237,21 @@ mod tests { use super::*; use crate::v1::Column; + #[test] + fn test_timestamp_unit_roundtrip() { + for unit in [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ] { + assert_eq!(timestamp_unit(timestamp_datatype(unit)), Some(unit)); + } + // Non-timestamp types have no time unit. + assert_eq!(timestamp_unit(ColumnDataType::String), None); + assert_eq!(timestamp_unit(ColumnDataType::Datetime), None); + } + #[test] fn test_values_with_capacity() { let values = values_with_capacity(ColumnDataType::Int8, 2); diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 7a0f706c58d..e82aa343870 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -124,7 +124,7 @@ impl OpenTelemetryProtocolHandler for Instance { metric_ctx.resource_info = self.otlp_resource_info; let otlp::metrics::MetricsConversion { - requests, + mut requests, rows, semantic_index, resource_info, @@ -173,15 +173,30 @@ impl OpenTelemetryProtocolHandler for Instance { 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 + let batcher = if batcher.is_some() { + // Align the requests' time index unit with the physical table's + // before the bulk eligibility check: the OTLP encoder keeps + // nanosecond precision on the metric engine path, while the bulk + // path assumes the physical table's unit (today millisecond). + // Without this, nanosecond requests would silently skip the + // batcher, and a non-millisecond physical table must not enter + // the bulk path either (the eligibility check rejects its unit). + self.inserter + .align_metric_row_inserts_time_unit(&ctx, &physical_table, &mut requests) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)?; + if self .inserter .can_batch_metric_rows(&requests, &ctx, &physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu)? - { - batcher + { + batcher + } else { + None + } } else { None }; diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 56d4af2284f..fdf9b02a5ff 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -791,7 +791,9 @@ mod tests { use common_function::utils::partition_expr_version; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_recordbatch::RecordBatches; - use datatypes::arrow::array::{Float64Array, TimestampMillisecondArray}; + use datatypes::arrow::array::{ + Float64Array, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::value::Value as PartitionValue; @@ -1116,6 +1118,91 @@ mod tests { expr.as_json_str().unwrap() } + #[tokio::test] + async fn test_put_and_scan_microsecond_physical_region() { + let env = TestEnv::new().await; + let engine = env.metric(); + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + env.create_physical_region_with_ts_type( + physical_region_id, + &TestEnv::default_table_dir(), + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + ) + .await; + + // A logical region with a matching microsecond time index is accepted. + let region_create_request = test_util::create_logical_region_request_with_ts_type( + &["job"], + physical_region_id, + &table_dir("test", logical_region_id.table_id()), + ConcreteDataType::timestamp_microsecond_datatype(), + ); + engine + .handle_request( + logical_region_id, + RegionRequest::Create(region_create_request), + ) + .await + .unwrap(); + + // Writing microsecond rows works. + let affected_rows = engine + .handle_request( + logical_region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags_and_ts_datatype( + &["job"], + ColumnDataType::TimestampMicrosecond, + ), + rows: test_util::build_rows_with_ts_datatype( + 1, + 2, + ColumnDataType::TimestampMicrosecond, + ), + }, + hint: None, + partition_expr_version: None, + skip_wal: false, + }), + ) + .await + .unwrap(); + assert_eq!(affected_rows.affected_rows, 2); + + // The scan returns timestamps in the physical region's unit. + let stream = engine + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut rows = Vec::new(); + for batch in batches.iter() { + let batch = batch.df_record_batch(); + let timestamps = batch + .column(batch.schema().index_of(greptime_timestamp()).unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(batch.schema().index_of(greptime_value()).unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend( + timestamps + .values() + .iter() + .copied() + .zip(values.values().iter().copied()), + ); + } + rows.sort_unstable_by_key(|(timestamp, _)| *timestamp); + assert_eq!(rows, vec![(0, 0.0), (1, 1.0)]); + } + async fn create_logical_region_with_tags( env: &TestEnv, physical_region_id: RegionId, diff --git a/src/metric-engine/src/test_util.rs b/src/metric-engine/src/test_util.rs index 1f177370988..0671351ddb3 100644 --- a/src/metric-engine/src/test_util.rs +++ b/src/metric-engine/src/test_util.rs @@ -158,6 +158,24 @@ impl TestEnv { physical_region_id: RegionId, table_dir: &str, options: Vec<(String, String)>, + ) { + self.create_physical_region_with_ts_type( + physical_region_id, + table_dir, + options, + ConcreteDataType::timestamp_millisecond_datatype(), + ) + .await + } + + /// Create regions in [MetricEngine] with specific `physical_region_id` + /// and time index column type. + pub async fn create_physical_region_with_ts_type( + &self, + physical_region_id: RegionId, + table_dir: &str, + options: Vec<(String, String)>, + ts_type: ConcreteDataType, ) { let region_create_request = RegionCreateRequest { engine: METRIC_ENGINE_NAME.to_string(), @@ -165,11 +183,7 @@ impl TestEnv { ColumnMetadata { column_id: 0, semantic_type: SemanticType::Timestamp, - column_schema: ColumnSchema::new( - greptime_timestamp(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), + column_schema: ColumnSchema::new(greptime_timestamp(), ts_type, false), }, ColumnMetadata { column_id: 1, @@ -330,16 +344,28 @@ pub fn create_logical_region_request( tags: &[&str], physical_region_id: RegionId, table_dir: &str, +) -> RegionCreateRequest { + create_logical_region_request_with_ts_type( + tags, + physical_region_id, + table_dir, + ConcreteDataType::timestamp_millisecond_datatype(), + ) +} + +/// Generate a [RegionCreateRequest] for logical region with the given time +/// index column type. +pub fn create_logical_region_request_with_ts_type( + tags: &[&str], + physical_region_id: RegionId, + table_dir: &str, + ts_type: ConcreteDataType, ) -> RegionCreateRequest { let mut column_metadatas = vec![ ColumnMetadata { column_id: 0, semantic_type: SemanticType::Timestamp, - column_schema: ColumnSchema::new( - greptime_timestamp(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ), + column_schema: ColumnSchema::new(greptime_timestamp(), ts_type, false), }, ColumnMetadata { column_id: 1, @@ -407,10 +433,20 @@ pub fn alter_logical_region_request(tags: &[&str]) -> RegionAlterRequest { /// /// The result will also contains default timestamp and value column at beginning. pub fn row_schema_with_tags(tags: &[&str]) -> Vec { + row_schema_with_tags_and_ts_datatype(tags, ColumnDataType::TimestampMillisecond) +} + +/// Generate a row schema with given tag columns and time index datatype. +/// +/// The result will also contains default timestamp and value column at beginning. +pub fn row_schema_with_tags_and_ts_datatype( + tags: &[&str], + ts_datatype: ColumnDataType, +) -> Vec { let mut schema = vec![ PbColumnSchema { column_name: greptime_timestamp().to_string(), - datatype: ColumnDataType::TimestampMillisecond as i32, + datatype: ts_datatype as i32, semantic_type: SemanticType::Timestamp as _, datatype_extension: None, options: None, @@ -440,11 +476,33 @@ pub fn row_schema_with_tags(tags: &[&str]) -> Vec { /// The schema is generated by [row_schema_with_tags]. `num_tags` doesn't need to be precise, /// it's used to determine the column id for new columns. pub fn build_rows(num_tags: usize, num_rows: usize) -> Vec { + build_rows_with_ts_datatype(num_tags, num_rows, ColumnDataType::TimestampMillisecond) +} + +/// Build [Row]s whose time index values use the given timestamp datatype. +/// +/// The schema is generated by [row_schema_with_tags_and_ts_datatype]. +/// `num_tags` doesn't need to be precise, it's used to determine the column id +/// for new columns. +pub fn build_rows_with_ts_datatype( + num_tags: usize, + num_rows: usize, + ts_datatype: ColumnDataType, +) -> Vec { + let ts_value = |i: usize| -> ValueData { + match ts_datatype { + ColumnDataType::TimestampSecond => ValueData::TimestampSecondValue(i as _), + ColumnDataType::TimestampMillisecond => ValueData::TimestampMillisecondValue(i as _), + ColumnDataType::TimestampMicrosecond => ValueData::TimestampMicrosecondValue(i as _), + ColumnDataType::TimestampNanosecond => ValueData::TimestampNanosecondValue(i as _), + _ => unreachable!("not a timestamp datatype"), + } + }; let mut rows = vec![]; for i in 0..num_rows { let mut values = vec![ Value { - value_data: Some(ValueData::TimestampMillisecondValue(i as _)), + value_data: Some(ts_value(i)), }, Value { value_data: Some(ValueData::F64Value(i as f64)), diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 49215f89713..24acdeaba28 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -22,6 +22,7 @@ use api::v1::region::{ InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests, RegionRequestHeader, }; +use api::v1::value::ValueData; use api::v1::{ AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests, RowInsertRequest, RowInsertRequests, Rows, SemanticType, @@ -46,6 +47,8 @@ use common_query::native_histogram::{is_native_histogram_value_type, native_hist use common_query::prelude::{greptime_timestamp, greptime_value}; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; +use common_time::Timestamp; +use common_time::timestamp::TimeUnit; use datatypes::schema::SkippingIndexOptions; use futures_util::future; use meter_core::data::MeterRecord; @@ -429,6 +432,7 @@ impl Inserter { statement_executor, accommodate_existing_schema, is_single_value, + None, ) .await?; @@ -565,10 +569,21 @@ impl Inserter { validate_column_count_match(&requests)?; // check and create physical table - self.create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) + let physical_table_ref = self + .create_physical_table_on_demand(&ctx, physical_table.clone(), statement_executor) .await?; - // check and create logical tables + // check and create logical tables; `create_or_alter_tables_on_demand` + // aligns each request's time index unit with the unit of the table it + // targets, inside its existing table lookups: existing tables keep + // their own unit (which matches the physical table they are bound + // to), and new tables use the selected physical table's unit (from + // `physical_table_ref`). Ingestion endpoints encode timestamps in a + // fixed unit (prometheus remote write always uses millisecond; OTLP + // keeps nanosecond precision on the metric engine path), and the + // metric engine requires each logical table's requests to match its + // time index unit. Narrowing conversions truncate the sub-unit part + // (floor), following `Timestamp::convert_to`. let CreateAlterTableResult { instant_table_ids, table_infos, @@ -580,6 +595,7 @@ impl Inserter { statement_executor, true, true, + table_time_index_unit(&physical_table_ref), ) .await?; let name_to_info = table_infos @@ -913,6 +929,7 @@ impl Inserter { statement_executor, false, false, + None, ) .await?; Ok(()) @@ -930,6 +947,13 @@ impl Inserter { /// custom schema, and then inserts data with endpoints that have default schema setting, like prometheus /// remote write. This will modify the `RowInsertRequests` in place. /// `is_single_value` indicates whether the default schema only contains single value column so we can accommodate it. + /// + /// `align_time_index_unit` is the selected physical metric table's time + /// index unit; passing `Some` (metric engine path only) rewrites each + /// request's time index column, inside this function's existing table + /// lookups (no extra catalog access): existing destination tables are + /// converted to their own unit, and new tables to the given unit. + #[allow(clippy::too_many_arguments)] async fn create_or_alter_tables_on_demand( &self, requests: &mut RowInsertRequests, @@ -938,6 +962,7 @@ impl Inserter { statement_executor: &StatementExecutor, accommodate_existing_schema: bool, is_single_value: bool, + align_time_index_unit: Option, ) -> Result { let _timer = crate::metrics::CREATE_ALTER_ON_DEMAND .with_label_values(&[auto_create_table_type.as_str()]) @@ -958,7 +983,7 @@ impl Inserter { && !has_auto_create_exempt_table { let mut instant_table_ids = HashSet::new(); - for req in &requests.inserts { + for req in &mut requests.inserts { let table = match self.get_table(catalog, &schema, &req.table_name).await? { Some(table) => table, // System-defined table: created canonically by the system, @@ -978,6 +1003,15 @@ impl Inserter { .fail(); } }; + // Metric path: an existing destination table keeps its own + // time index unit (it may be bound to another physical + // table than the one selected by this request). + if align_time_index_unit.is_some() + && let Some(rows) = req.rows.as_mut() + && let Some(target_unit) = table_time_index_unit(&table) + { + convert_rows_time_unit(rows, target_unit)?; + } let table_info = table.table_info(); if matches!(auto_create_table_type, AutoCreateTableType::Trace { .. }) { validate_trace_table_model(&table_info, ctx)?; @@ -1013,6 +1047,15 @@ impl Inserter { if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); } + // Metric path: an existing destination table keeps its + // own time index unit (it may be bound to another + // physical table than the one selected by this request). + if align_time_index_unit.is_some() + && let Some(rows) = req.rows.as_mut() + && let Some(target_unit) = table_time_index_unit(&table) + { + convert_rows_time_unit(rows, target_unit)?; + } if auto_create_allowed && let Some(alter_expr) = self.get_alter_table_expr_on_demand( req, @@ -1058,6 +1101,14 @@ impl Inserter { .fail(); } None => { + // Metric path: a new table uses the selected physical + // table's unit; convert before the create expression is + // derived from the request schema. + if let Some(physical_unit) = align_time_index_unit + && let Some(rows) = req.rows.as_mut() + { + convert_rows_time_unit(rows, physical_unit)?; + } let semantic_index = per_table_semantics .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) .as_ref(); @@ -1256,17 +1307,16 @@ impl Inserter { ctx: &QueryContextRef, physical_table: String, statement_executor: &StatementExecutor, - ) -> Result<()> { + ) -> Result { let catalog_name = ctx.current_catalog(); let schema_name = ctx.current_schema(); // check if exist - if self + if let Some(table) = self .get_table(catalog_name, &schema_name, &physical_table) .await? - .is_some() { - return Ok(()); + return Ok(table); } // Gate here too, otherwise a disabled switch would still leak the physical table. @@ -1318,7 +1368,7 @@ impl Inserter { .await; match res { - Ok(_) => Ok(()), + Ok(table) => Ok(table), Err(err) => { error!(err; "Failed to create table {table_reference}"); Err(err) @@ -1326,6 +1376,75 @@ impl Inserter { } } + /// Aligns each request's time index unit with the unit of the table the + /// request targets, so that write-path gates (e.g. the logical batcher + /// eligibility check) observe the destination table's unit instead of the + /// ingestion endpoint's encoding unit (Prometheus remote write always + /// uses millisecond; OTLP keeps nanosecond precision on the metric engine + /// path). Existing destination tables keep their own unit — they may be + /// bound to a different physical table than `physical_table` — while new + /// tables use `physical_table`'s unit. A missing physical table is + /// treated as the millisecond unit of the auto-created default; the + /// actual creation (if any) happens later in + /// [`Inserter::handle_metric_row_inserts`]. + /// + /// This is only needed on paths that choose a write path before + /// `handle_metric_row_inserts` (whose own table lookups perform the + /// alignment again, as a no-op after this). + pub async fn align_metric_row_inserts_time_unit( + &self, + ctx: &QueryContextRef, + physical_table: &str, + requests: &mut RowInsertRequests, + ) -> Result<()> { + // The unit conversion indexes rows by the time index position, which + // requires well-formed requests. + validate_column_count_match(requests)?; + let physical_unit = match self + .get_table(ctx.current_catalog(), &ctx.current_schema(), physical_table) + .await? + { + Some(table) => table_time_index_unit(&table), + // A missing physical table is auto-created with the millisecond + // unit later. + None => Some(TimeUnit::Millisecond), + }; + self.align_metric_rows_per_destination(ctx, physical_unit, requests) + .await + } + + /// Converts each request's time index column to the unit of the table the + /// request targets: the existing table's own unit when the table exists, + /// and `physical_unit` (of the request's selected physical table) for new + /// tables. + async fn align_metric_rows_per_destination( + &self, + ctx: &QueryContextRef, + physical_unit: Option, + requests: &mut RowInsertRequests, + ) -> Result<()> { + for request in &mut requests.inserts { + let Some(rows) = request.rows.as_mut() else { + continue; + }; + let target_unit = match self + .get_table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + ) + .await? + { + Some(table) => table_time_index_unit(&table), + None => physical_unit, + }; + if let Some(target_unit) = target_unit { + convert_rows_time_unit(rows, target_unit)?; + } + } + Ok(()) + } + async fn get_table( &self, catalog: &str, @@ -1592,6 +1711,84 @@ fn request_is_native_histogram(request_schema: &[ColumnSchema]) -> bool { ) } +/// Returns the table's time index unit, if any. A metric table without a +/// timestamp column is left alone by the unit alignment: the metric engine +/// rejects it anyway. +fn table_time_index_unit(table: &TableRef) -> Option { + table + .table_info() + .meta + .schema + .timestamp_column() + .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) +} + +fn convert_rows_time_unit(rows: &mut Rows, target_unit: TimeUnit) -> Result<()> { + let Some(ts_index) = rows + .schema + .iter() + .position(|col| col.semantic_type == SemanticType::Timestamp as i32) + else { + return Ok(()); + }; + let Some(source_unit) = ColumnDataType::try_from(rows.schema[ts_index].datatype) + .ok() + .and_then(api::helper::timestamp_unit) + else { + return Ok(()); + }; + if source_unit == target_unit { + return Ok(()); + } + + rows.schema[ts_index].datatype = api::helper::timestamp_datatype(target_unit) as i32; + // Timestamp columns never carry a datatype extension. + rows.schema[ts_index].datatype_extension = None; + + // Note: the schema is rewritten before the rows are converted, so an + // overflow error mid-batch leaves this request half-converted. That is + // harmless: the error aborts the whole insert request. + // + // `validate_column_count_match` guarantees every row carries exactly one + // value per schema column, so the time index position is directly in + // bounds; no per-value search is needed. + for row in &mut rows.rows { + debug_assert_eq!(row.values.len(), rows.schema.len()); + let value = &mut row.values[ts_index]; + let Some(value_data) = value.value_data.take() else { + continue; + }; + value.value_data = + convert_timestamp_value_data(value_data, source_unit, target_unit, ts_index)?; + } + Ok(()) +} + +fn convert_timestamp_value_data( + value_data: ValueData, + source_unit: TimeUnit, + target_unit: TimeUnit, + column_index: usize, +) -> Result> { + let timestamp = match value_data { + ValueData::TimestampSecondValue(v) => Timestamp::new_second(v), + ValueData::TimestampMillisecondValue(v) => Timestamp::new_millisecond(v), + ValueData::TimestampMicrosecondValue(v) => Timestamp::new_microsecond(v), + ValueData::TimestampNanosecondValue(v) => Timestamp::new_nanosecond(v), + // Null or non-timestamp value; nothing to convert. + other => return Ok(Some(other)), + }; + let converted = timestamp + .convert_to(target_unit) + .with_context(|| InvalidInsertRequestSnafu { + reason: format!( + "timestamp column {column_index} value {} in unit {source_unit:?} overflows when converting to unit {target_unit:?}", + timestamp.value() + ), + })?; + Ok(api::helper::to_grpc_value(datatypes::value::Value::Timestamp(converted)).value_data) +} + fn table_is_native_histogram(table: &TableRef) -> bool { let mut fields = table.field_columns(); let Some(col) = fields.next() else { @@ -1909,7 +2106,7 @@ mod tests { use api::helper::ColumnDataTypeWrapper; use api::v1::helper::{field_column_schema, time_index_column_schema}; - use api::v1::{RowInsertRequest, Rows, Value}; + use api::v1::{Row, RowInsertRequest, Rows, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_meta::cache::new_table_flownode_set_cache; use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler; @@ -1976,6 +2173,269 @@ mod tests { )) } + fn make_metric_physical_table_ref_with_time_unit(unit: TimeUnit) -> TableRef { + let schema = datatypes::schema::SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new( + greptime_timestamp(), + ConcreteDataType::timestamp_datatype(unit), + false, + ) + .with_time_index(true), + ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), + ]) + .unwrap() + .build() + .unwrap(); + let meta = TableMetaBuilder::empty() + .schema(Arc::new(schema)) + .primary_key_indices(vec![]) + .value_indices(vec![1]) + .engine("metric") + .next_column_id(0) + .options(Default::default()) + .created_on(Default::default()) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(1) + .table_version(0) + .name("greptime_physical_table") + .schema_name(DEFAULT_SCHEMA_NAME) + .catalog_name(DEFAULT_CATALOG_NAME) + .desc(None) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(table::Table::new( + info, + table::metadata::FilterPushDownType::Unsupported, + Arc::new(DummyDataSource), + )) + } + + fn ms_row_insert_request(timestamp_ms: i64) -> RowInsertRequest { + ms_row_insert_request_named("my_metric", timestamp_ms) + } + + fn ms_row_insert_request_named(table: &str, timestamp_ms: i64) -> RowInsertRequest { + RowInsertRequest { + table_name: table.to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + ), + field_column_schema(greptime_value(), ColumnDataType::Float64), + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(timestamp_ms)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }], + }), + } + } + + /// Converts the time index column of each request to `target_unit`. + fn convert_row_insert_requests_time_unit( + requests: &mut RowInsertRequests, + target_unit: TimeUnit, + ) -> Result<()> { + for request in &mut requests.inserts { + let Some(rows) = request.rows.as_mut() else { + continue; + }; + convert_rows_time_unit(rows, target_unit)?; + } + Ok(()) + } + + #[test] + fn test_convert_row_insert_requests_time_unit_noop_when_matching() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(123)], + }; + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Millisecond).unwrap(); + let rows = requests.inserts[0].rows.as_ref().unwrap(); + assert_eq!( + rows.schema[0].datatype, + ColumnDataType::TimestampMillisecond as i32 + ); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(123)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_widens_losslessly() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(123)], + }; + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Microsecond).unwrap(); + let rows = requests.inserts[0].rows.as_ref().unwrap(); + assert_eq!( + rows.schema[0].datatype, + ColumnDataType::TimestampMicrosecond as i32 + ); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(123_000)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_truncates_on_narrowing() { + // 123_456_789 ns floors to 123_456 us and 123 ms; negative values + // floor towards negative infinity, matching `Timestamp::convert_to`. + let requests = |value_ns: i64| RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "my_metric".to_string(), + rows: Some(Rows { + schema: vec![ + time_index_column_schema( + greptime_timestamp(), + ColumnDataType::TimestampNanosecond, + ), + field_column_schema(greptime_value(), ColumnDataType::Float64), + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampNanosecondValue(value_ns)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }], + }), + }], + }; + + let mut reqs = requests(123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(123_456)) + )); + + let mut reqs = requests(123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Millisecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(123)) + )); + + let mut reqs = requests(-123_456_789); + convert_row_insert_requests_time_unit(&mut reqs, TimeUnit::Microsecond).unwrap(); + let rows = reqs.inserts[0].rows.as_ref().unwrap(); + assert!(matches!( + rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(-123_457)) + )); + } + + #[test] + fn test_convert_row_insert_requests_time_unit_overflow() { + let mut requests = RowInsertRequests { + inserts: vec![ms_row_insert_request(i64::MAX)], + }; + let err = + convert_row_insert_requests_time_unit(&mut requests, TimeUnit::Nanosecond).unwrap_err(); + assert!(err.to_string().contains("overflows"), "{err}"); + } + + #[test] + fn test_table_time_index_unit() { + assert_eq!( + table_time_index_unit(&make_metric_physical_table_ref_with_time_unit( + TimeUnit::Microsecond + )), + Some(TimeUnit::Microsecond) + ); + assert_eq!( + table_time_index_unit(&make_metric_physical_table_ref_with_time_unit( + TimeUnit::Millisecond + )), + Some(TimeUnit::Millisecond) + ); + } + + #[tokio::test] + async fn test_align_metric_rows_per_destination() { + use catalog::RegisterTableRequest; + use catalog::memory::MemoryCatalogManager; + + // An existing millisecond logical table `existing`, plus a + // microsecond physical table selected by the request: the existing + // table's request keeps the millisecond unit (it is bound to another + // physical table), while the new table's request is converted to the + // selected physical table's microsecond unit. + let mut inserter = batcher_test_inserter().await; + let existing = make_metric_physical_table_ref_with_time_unit(TimeUnit::Millisecond); + let phy_us = make_metric_physical_table_ref_with_time_unit(TimeUnit::Microsecond); + let catalog = MemoryCatalogManager::with_default_setup(); + for (table_name, table_id, table) in [("existing", 1, existing), ("phy_us", 2, phy_us)] { + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id, + table, + }) + .unwrap(); + } + inserter.catalog_manager = catalog; + + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + let mut requests = RowInsertRequests { + inserts: vec![ + ms_row_insert_request_named("existing", 123), + ms_row_insert_request_named("fresh", 123), + ], + }; + inserter + .align_metric_row_inserts_time_unit(&ctx, "phy_us", &mut requests) + .await + .unwrap(); + + let existing_rows = requests.inserts[0].rows.as_ref().unwrap(); + assert_eq!( + existing_rows.schema[0].datatype, + ColumnDataType::TimestampMillisecond as i32 + ); + assert!(matches!( + existing_rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMillisecondValue(123)) + )); + + let fresh_rows = requests.inserts[1].rows.as_ref().unwrap(); + assert_eq!( + fresh_rows.schema[0].datatype, + ColumnDataType::TimestampMicrosecond as i32 + ); + assert!(matches!( + fresh_rows.rows[0].values[0].value_data, + Some(ValueData::TimestampMicrosecondValue(123_000)) + )); + } + #[tokio::test] async fn test_accommodate_existing_schema_and_reject_kind_changes() { let ts_name = "my_ts"; diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index bd7a26da613..fcd4e1ea38f 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::collections::HashSet; use std::future::{Future, ready}; use std::num::NonZeroUsize; use std::sync::Arc; @@ -35,6 +36,7 @@ use common_batcher::worker_registry::WorkerRegistry; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; +use common_time::timestamp::TimeUnit; use meter_core::data::MeterRecord; use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; @@ -67,15 +69,15 @@ const PHYSICAL_TABLE_KEY: &str = "physical_table"; const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; #[derive(Debug, Clone, Hash, Eq, PartialEq)] -struct BatchKey { - catalog: String, - schema: String, - physical_table: String, - skip_wal: bool, +pub(crate) struct BatchKey { + pub(crate) catalog: String, + pub(crate) schema: String, + pub(crate) physical_table: String, + pub(crate) skip_wal: bool, } // Requests can share a batch only when their write target and WAL policy match. -fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { +pub(crate) fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { let physical_table = ctx .extension(PHYSICAL_TABLE_KEY) .unwrap_or(GREPTIME_PHYSICAL_TABLE) @@ -171,6 +173,90 @@ impl LogicalTablePendingRowsBatcher { .map(|(rows, ())| rows) } + /// Returns whether the bulk path can accept `batches`: + /// - the physical metric table selected by each context must have a + /// millisecond time index (or not exist yet — it is auto-created with + /// the millisecond unit): new tables are created on it and the bulk + /// encode produces millisecond batches; + /// - every existing destination table must itself have a millisecond time + /// index: it may be bound to another physical table with another unit, + /// and the bulk encoder would build millisecond arrays against its + /// schema. + /// + /// Incompatible requests must stay on the ordinary insert path, which + /// converts the requests to each destination table's unit. Write targets + /// and destination tables are resolved once per distinct name. + pub(crate) async fn accepts_bulk_time_indexes( + &self, + batches: impl Iterator, + ) -> bool { + let mut checked_targets = HashSet::new(); + let mut checked_tables = HashSet::new(); + for (ctx, requests) in batches { + let key = batch_key_from_ctx(ctx); + if checked_targets.insert((key.catalog, key.schema, key.physical_table)) + && !self.accepts_physical_table_time_index(ctx).await + { + return false; + } + for request in &requests.inserts { + if !checked_tables.insert((ctx.current_schema(), request.table_name.clone())) { + continue; + } + let Ok(Some(table)) = self + .catalog_manager + .table( + ctx.current_catalog(), + &ctx.current_schema(), + &request.table_name, + None, + ) + .await + else { + // New table: governed by the selected physical table's + // unit, checked above. + continue; + }; + if table + .table_info() + .meta + .schema + .timestamp_column() + .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) + .is_none_or(|unit| unit != TimeUnit::Millisecond) + { + return false; + } + } + } + true + } + + /// Returns whether the physical metric table resolved from `ctx` can use + /// the bulk path. The bulk encode produces millisecond timestamp batches + /// only, so a physical table with another time index unit (e.g. created + /// as `TIMESTAMP(6)`) must stay on the ordinary insert path, which + /// converts the requests to the physical table's unit. A missing physical + /// table is accepted: it is auto-created with the millisecond unit. + async fn accepts_physical_table_time_index(&self, ctx: &QueryContextRef) -> bool { + let key = batch_key_from_ctx(ctx); + let Ok(Some(table)) = self + .catalog_manager + .table(&key.catalog, &key.schema, &key.physical_table, None) + .await + else { + return true; + }; + table + .table_info() + .meta + .schema + .timestamp_column() + .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) + .map(|unit| unit == TimeUnit::Millisecond) + .unwrap_or(true) + } + /// Submits with request-level accounting after schema preparation and before /// queue admission. Acknowledgement follows the global batching policy. pub async fn submit_with( diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 2c4b6927451..44abb18a5be 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -397,27 +397,45 @@ async fn write_prometheus_rows_with_progress( mut batches: Vec, ) -> std::result::Result { if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { + // Preflight before the bulk eligibility decision: pre_write hooks + // may redirect the contexts (e.g. to a per-tenant schema), so + // eligibility must be evaluated against the prepared batches — + // and the fallback below reuses them without re-running hooks or + // admission. preflight_prometheus_rows(&prom_store_handler, &mut batches) .await .map_err(|error| PromWriteError { error, rows_written: 0, })?; - let mut rows_written = 0; - for (temp_ctx, reqs) in batches { - let rows = batcher - .submit(reqs, temp_ctx) - .await - .map_err(|error| PromWriteError { - error, - rows_written, - })?; - rows_written += rows; + // The bulk encode produces millisecond batches only; write targets + // and existing destination tables with another time index unit stay + // on the ordinary insert path, which converts the requests to each + // destination table's unit. + if batcher.accepts_bulk_time_indexes(batches.iter()).await { + let mut rows_written = 0; + for (temp_ctx, reqs) in batches { + let rows = + batcher + .submit(reqs, temp_ctx) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + rows_written += rows; + } + return Ok(PromWriteOutcome { + write_cost: 0, + rows_written, + }); } - return Ok(PromWriteOutcome { - write_cost: 0, - rows_written, - }); + return write_prepared_prometheus_rows_with_progress( + prom_store_handler, + batches, + prom_store_with_metric_engine, + ) + .await; } let row_counts = batches @@ -456,6 +474,33 @@ async fn write_prometheus_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path; the pre_write hooks and admission have already run. +async fn write_prepared_prometheus_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut rows_written = 0; + for (ctx, request) in batches { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteError { + error, + rows_written, + })?; + write_cost += output.meta.cost; + rows_written += rows; + } + Ok(PromWriteOutcome { + write_cost, + rows_written, + }) +} + async fn write_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, pending_rows_batcher: Option>, @@ -483,20 +528,44 @@ async fn write_prometheus_v2_rows_with_progress( }); } + let sample_batch_count = sample_batches.len(); + let mut batches = sample_batches; + batches.extend(histogram_batches); + if prom_store_with_metric_engine && let Some(batcher) = pending_rows_batcher { - return write_batched_prometheus_v2_rows_with_progress( + // Same ordering as the v1 path: preflight (which may redirect the + // contexts) before the bulk eligibility decision, and the fallback + // reuses the prepared batches without re-running hooks or admission. + preflight_prometheus_rows(&prom_store_handler, &mut batches) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written: 0, + histograms_written: 0, + })?; + // The bulk encode produces millisecond batches only; write targets + // and existing destination tables with another time index unit stay + // on the ordinary insert path, which converts the requests to each + // destination table's unit. + if batcher.accepts_bulk_time_indexes(batches.iter()).await { + return write_batched_prometheus_v2_rows_with_progress( + prom_store_handler, + batcher.as_ref(), + prom_store_with_metric_engine, + sample_batch_count, + batches, + ) + .await; + } + return write_prepared_prometheus_v2_rows_with_progress( prom_store_handler, - batcher.as_ref(), + batches, + sample_batch_count, prom_store_with_metric_engine, - sample_batches, - histogram_batches, ) .await; } - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); let row_counts = batches .iter() .map(|(_, request)| prom_write_row_count(request)) @@ -544,24 +613,49 @@ async fn write_prometheus_v2_rows_with_progress( }) } +/// Writes already-preflighted batches through the ordinary (prepared) write +/// path with v2 partial-progress accounting; the pre_write hooks and +/// admission have already run. +async fn write_prepared_prometheus_v2_rows_with_progress( + prom_store_handler: PromStoreProtocolHandlerRef, + batches: Vec, + sample_batch_count: usize, + prom_store_with_metric_engine: bool, +) -> std::result::Result { + let mut write_cost = 0; + let mut samples_written = 0; + let mut histograms_written = 0; + for (index, (ctx, request)) in batches.into_iter().enumerate() { + let rows = prom_write_row_count(&request); + let output = prom_store_handler + .write_prepared(request, ctx, prom_store_with_metric_engine) + .await + .map_err(|error| PromWriteV2Error { + error, + samples_written, + histograms_written, + })?; + write_cost += output.meta.cost; + if index < sample_batch_count { + samples_written += rows; + } else { + histograms_written += rows; + } + } + Ok(PromWriteV2Outcome { + write_cost, + samples_written, + histograms_written, + }) +} + async fn write_batched_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, batcher: &B, prom_store_with_metric_engine: bool, - sample_batches: Vec, - histogram_batches: Vec, + sample_batch_count: usize, + batches: Vec, ) -> std::result::Result { - let sample_batch_count = sample_batches.len(); - let mut batches = sample_batches; - batches.extend(histogram_batches); - preflight_prometheus_rows(&prom_store_handler, &mut batches) - .await - .map_err(|error| PromWriteV2Error { - error, - samples_written: 0, - histograms_written: 0, - })?; - let mut samples_written = 0; let mut histograms_written = 0; let mut write_cost = 0; @@ -862,14 +956,18 @@ mod tests { events: events.clone(), }; - let Ok(outcome) = write_batched_prometheus_v2_rows_with_progress( - handler, - &batcher, - true, - vec![test_prom_write_batch("sample")], - vec![test_prom_write_batch("histogram")], - ) - .await + let mut batches = vec![ + test_prom_write_batch("sample"), + test_prom_write_batch("histogram"), + ]; + // The caller preflights before choosing the bulk path; the batched + // writer consumes the prepared batches. + preflight_prometheus_rows(&handler, &mut batches) + .await + .unwrap(); + let Ok(outcome) = + write_batched_prometheus_v2_rows_with_progress(handler, &batcher, true, 1, batches) + .await else { panic!("mixed remote write should succeed") }; diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index f32816fbe36..53a2189044d 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -15,7 +15,7 @@ use std::collections::BTreeMap; use ahash::{HashMap, HashMapExt}; -use api::helper::ColumnDataTypeWrapper; +use api::helper::{ColumnDataTypeWrapper, timestamp_unit}; use api::v1::column_data_type_extension::TypeExt; use api::v1::column_def::options_from_column_schema; use api::v1::value::ValueData; @@ -898,16 +898,6 @@ fn is_timestamp_type(datatype: ColumnDataType) -> bool { timestamp_unit(datatype).is_some() } -fn timestamp_unit(datatype: ColumnDataType) -> Option { - match datatype { - ColumnDataType::TimestampSecond => Some(TimeUnit::Second), - ColumnDataType::TimestampMillisecond => Some(TimeUnit::Millisecond), - ColumnDataType::TimestampMicrosecond => Some(TimeUnit::Microsecond), - ColumnDataType::TimestampNanosecond => Some(TimeUnit::Nanosecond), - _ => None, - } -} - fn parse_export_logs_service_request_to_rows( request: ExportLogsServiceRequest, select_info: Box, diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index e518d732479..fd496db5a98 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -998,9 +998,15 @@ fn write_timestamp( table: &mut TableData, row: &mut Vec, time_nano: i64, - legacy_mode: bool, + metric_ctx: &OtlpMetricCtx, ) -> Result<()> { - if legacy_mode { + // Keep the full nanosecond precision whenever the request is headed for + // the metric engine: `Inserter::handle_metric_row_inserts` converts the + // timestamps to the physical table's time index unit (which may be + // micro/nanosecond). Only the non-metric prometheus-compatible path is + // fixed to milliseconds, to keep auto-created mito tables on the + // millisecond time index. + if metric_ctx.is_legacy || metric_ctx.with_metric_engine { row_writer::write_ts_to_nanos( table, greptime_timestamp(), @@ -1102,7 +1108,7 @@ fn write_tags_and_timestamp( )?; } - write_timestamp(table, row, timestamp_nanos, metric_ctx.is_legacy)?; + write_timestamp(table, row, timestamp_nanos, metric_ctx)?; Ok(()) } @@ -1541,6 +1547,7 @@ fn encode_summary( #[cfg(test)] mod tests { + use api::v1::ColumnDataType; use common_query::prelude::set_default_prefix; use otel_arrow_rust::proto::opentelemetry::common::v1::AnyValue; use otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value as Val; @@ -2625,6 +2632,81 @@ mod tests { } } + #[test] + fn test_metric_engine_path_keeps_nanosecond_precision() { + let time_unix_nano = 1_704_067_200_123_456_789u64; + let request = metrics_request(vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }]); + + // The metric engine path keeps the full nanosecond precision here; + // `Inserter::handle_metric_row_inserts` converts the timestamps to + // the physical table's time index unit afterwards. + let mut metric_ctx = OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + }; + let MetricsConversion { requests, .. } = + to_grpc_insert_requests(request, &mut metric_ctx).unwrap(); + + let rows = requests.inserts[0].rows.as_ref().unwrap(); + let ts_index = rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + rows.schema[ts_index].datatype, + ColumnDataType::TimestampNanosecond as i32 + ); + assert!(matches!( + rows.rows[0].values[ts_index].value_data, + Some(ValueData::TimestampNanosecondValue( + 1_704_067_200_123_456_789 + )) + )); + + // The non-metric prometheus-compatible path stays millisecond so + // auto-created mito tables keep the millisecond time index. + let mut compat_ctx = OtlpMetricCtx::default(); + let request = metrics_request(vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }]); + let MetricsConversion { requests, .. } = + to_grpc_insert_requests(request, &mut compat_ctx).unwrap(); + + let rows = requests.inserts[0].rows.as_ref().unwrap(); + let ts_index = rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap(); + assert_eq!( + rows.schema[ts_index].datatype, + ColumnDataType::TimestampMillisecond as i32 + ); + assert!(matches!( + rows.rows[0].values[ts_index].value_data, + Some(ValueData::TimestampMillisecondValue(1_704_067_200_123)) + )); + } + #[test] fn test_exponential_histogram_rejection_metrics() { // Other conversion tests update the same process-global counters. diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 9182f64ad00..e202070d98d 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -18,14 +18,19 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; +use std::sync::Arc; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{Label, Query, Sample, TimeSeries, WriteRequest}; use api::v1::RowInsertRequests; use arrow::array::{ - Array, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray, + Array, ArrayRef, AsArray, DictionaryArray, LargeStringArray, StringArray, StringViewArray, +}; +use arrow::compute::kernels::cast as casts; +use arrow::datatypes::{ + DataType, Float64Type, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, UInt32Type, }; -use arrow::datatypes::{Float64Type, TimestampMillisecondType, UInt32Type}; use common_grpc::precision::Precision; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; @@ -398,6 +403,50 @@ fn append_recordbatch_to_timeseries( .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { msg: format!("missing timestamp column '{timestamp_column_name}' in query result"), })?; + // The Prometheus remote read wire format carries millisecond timestamps, + // while the table's time index can use any time unit (e.g. a metric + // physical table created with TIMESTAMP(6)). Narrowing floors towards + // negative infinity, consistent with `Timestamp::convert_to` on the + // ingestion path; arrow's cast would truncate towards zero, so a + // pre-epoch sample (e.g. -1001us) would round to -1ms instead of -2ms + // and disagree with what ingesting the same instant into a millisecond + // table would have stored. + let ts_column: ArrayRef = match ts_column.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, _) => ts_column.clone(), + // Second -> millisecond is a widening (exact) conversion. + DataType::Timestamp(TimeUnit::Second, _) => casts::cast( + ts_column, + &DataType::Timestamp(TimeUnit::Millisecond, None), + ) + .map_err(|e| { + error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!( + "failed to cast timestamp column '{timestamp_column_name}' of datatype {:?} to millisecond: {e}", + ts_column.data_type() + ), + } + .build() + })?, + DataType::Timestamp(TimeUnit::Microsecond, _) => Arc::new( + ts_column + .as_primitive::() + .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000)), + ), + DataType::Timestamp(TimeUnit::Nanosecond, _) => Arc::new( + ts_column + .as_primitive::() + .unary::<_, TimestampMillisecondType>(|v| v.div_euclid(1_000_000)), + ), + _ => { + return error::InvalidPromRemoteReadQueryResultSnafu { + msg: format!( + "Expect timestamp column of datatype Timestamp(Millisecond), actual {:?}", + ts_column.data_type() + ), + } + .fail(); + } + }; let ts_column = ts_column .as_primitive_opt::() .with_context(|| error::InvalidPromRemoteReadQueryResultSnafu { diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 67a4e512f7d..59ce38919f7 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -19,6 +19,10 @@ mod test { use client::{DEFAULT_CATALOG_NAME, OutputData}; use common_recordbatch::RecordBatches; use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{ + DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, + }; use frontend::instance::Instance; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; @@ -37,6 +41,7 @@ mod test { use servers::query_handler::OpenTelemetryProtocolHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; use crate::standalone::GreptimeDbStandaloneBuilder; use crate::tests; @@ -1164,6 +1169,162 @@ WITH( ); } + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_microsecond_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_us_physical") + .build() + .await; + + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_704_067_200_123_456, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_microsecond_physical_table_on_distributed() { + let instance = tests::create_distributed_instance("test_otlp_us_physical_dist").await; + + test_otlp_metrics_into_non_millisecond_physical_table( + &instance.frontend(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_704_067_200_123_456, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_seconds_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_s_physical") + .build() + .await; + + // Narrowing truncates: 1704067200123456789ns -> 1704067200s. + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(0)", + TimeUnit::Second, + 1_704_067_200, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + pub async fn test_otlp_metrics_into_nanoseconds_physical_table_on_standalone() { + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_ns_physical") + .build() + .await; + + // Same unit: the full nanosecond precision is kept verbatim. + test_otlp_metrics_into_non_millisecond_physical_table( + standalone.fe_instance(), + "TIMESTAMP(9)", + TimeUnit::Nanosecond, + 1_704_067_200_123_456_789, + ) + .await; + } + + /// Regression test for : + /// a physical metric table pre-created with a non-millisecond time index + /// accepts OTLP ingestion; the nanosecond samples are converted to the + /// physical table's unit (kept verbatim for a nanosecond table, + /// truncating the sub-unit part for micro/second tables). + async fn test_otlp_metrics_into_non_millisecond_physical_table( + instance: &Arc, + sql_ts_type: &str, + expected_unit: TimeUnit, + expected_value: i64, + ) { + let db = "otlp_non_ms_physical"; + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, db); + // Route the request to the metric engine, like the OTLP HTTP handler + // does with the default `prom_store.with_metric_engine = true`. + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "my_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + attributes: vec![keyvalue("host", "h1")], + time_unix_nano: 1_704_067_200_123_456_789, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + instance.metrics(request, ctx.clone()).await.unwrap(); + + let mut output = instance + .do_query("SELECT greptime_timestamp FROM my_gauge", ctx.clone()) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = RecordBatches::try_collect(stream).await.unwrap().take(); + assert_eq!(batches[0].num_rows(), 1); + // The logical table's time index keeps the physical table's unit. + let ts_column = batches[0].column(0); + assert_eq!( + ts_column.data_type(), + &DataType::Timestamp(expected_unit, None), + "unexpected time index type" + ); + let stored = match expected_unit { + TimeUnit::Second => ts_column.as_primitive::().value(0), + TimeUnit::Millisecond => ts_column + .as_primitive::() + .value(0), + TimeUnit::Microsecond => ts_column + .as_primitive::() + .value(0), + TimeUnit::Nanosecond => ts_column.as_primitive::().value(0), + }; + assert_eq!(stored, expected_value); + } + fn build_request() -> ExportMetricsServiceRequest { let data_points = vec![ NumberDataPoint { diff --git a/tests-integration/src/prom_store.rs b/tests-integration/src/prom_store.rs index 0649b2dd778..af901120454 100644 --- a/tests-integration/src/prom_store.rs +++ b/tests-integration/src/prom_store.rs @@ -18,9 +18,19 @@ mod tests { use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ - Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, WriteRequest, + Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, TimeSeries, WriteRequest, }; + use api::v1::value::ValueData; + use api::v1::{ + ColumnDataType, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, Value, + }; + use client::OutputData; use common_catalog::consts::DEFAULT_CATALOG_NAME; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{ + DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, + }; use frontend::instance::Instance; use prost::Message; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; @@ -83,6 +93,432 @@ mod tests { .await; } + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_microsecond_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_us_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_000_000, + 2_000_000, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_distributed_prom_store_remote_rw_microsecond_physical_table() { + common_telemetry::init_default_ut_logging(); + let distributed = + tests::create_distributed_instance("test_prom_store_remote_rw_us_physical_table").await; + test_prom_store_remote_rw_non_millisecond_physical_table( + &distributed.frontend(), + "TIMESTAMP(6)", + TimeUnit::Microsecond, + 1_000_000, + 2_000_000, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_seconds_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_s_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + // Narrowing truncates: 1000ms/2000ms -> 1s/2s. + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(0)", + TimeUnit::Second, + 1, + 2, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_remote_rw_nanoseconds_physical_table() { + common_telemetry::init_default_ut_logging(); + let standalone = + GreptimeDbStandaloneBuilder::new("test_prom_store_remote_rw_ns_physical_table") + .build() + .await; + let instance = standalone.fe_instance(); + + // Lossless widening: 1000ms/2000ms -> 1e9ns/2e9ns. + test_prom_store_remote_rw_non_millisecond_physical_table( + instance, + "TIMESTAMP(9)", + TimeUnit::Nanosecond, + 1_000_000_000, + 2_000_000_000, + ) + .await; + } + + /// Regression test for : + /// prometheus remote write/read against a physical metric table + /// pre-created with a non-millisecond time index. Millisecond samples are + /// converted to the physical table's unit on write (lossless widening, + /// truncating narrowing) and narrowed back to milliseconds on remote read. + async fn test_prom_store_remote_rw_non_millisecond_physical_table( + instance: &Arc, + sql_ts_type: &str, + expected_unit: TimeUnit, + expected_first: i64, + expected_second: i64, + ) { + let db = "prometheus_non_ms"; + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, db)); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + let write_request = WriteRequest { + timeseries: vec![prom_store::mock_timeseries()[0].clone()], + ..Default::default() + }; + let (row_inserts, _) = to_grpc_row_insert_requests(&write_request).unwrap(); + instance + .write(row_inserts, ctx.clone(), true) + .await + .unwrap(); + + let read_request = ReadRequest { + queries: vec![Query { + start_timestamp_ms: 1000, + end_timestamp_ms: 2000, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "metric1".to_string(), + r#type: 0, + }], + ..Default::default() + }], + ..Default::default() + }; + let resp = instance.read(read_request, ctx.clone()).await.unwrap(); + let body = prom_store::snappy_decompress(&resp.body).unwrap(); + let read_response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(1, read_response.results.len()); + assert_eq!(1, read_response.results[0].timeseries.len()); + let timeseries = &read_response.results[0].timeseries[0]; + assert_eq!( + timeseries.samples, + vec![ + Sample { + value: 1.0, + timestamp: 1000, + }, + Sample { + value: 2.0, + timestamp: 2000, + }, + ] + ); + + // The stored timestamps keep the physical table's time index unit. + let mut output = instance + .do_query("SELECT greptime_timestamp FROM metric1", ctx.clone()) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + assert_eq!(batches[0].num_rows(), 2); + let ts_column = batches[0].column(0); + assert_eq!( + ts_column.data_type(), + &DataType::Timestamp(expected_unit, None), + "unexpected time index type" + ); + let stored = |row: usize| match expected_unit { + TimeUnit::Second => ts_column.as_primitive::().value(row), + TimeUnit::Millisecond => ts_column + .as_primitive::() + .value(row), + TimeUnit::Microsecond => ts_column + .as_primitive::() + .value(row), + TimeUnit::Nanosecond => ts_column + .as_primitive::() + .value(row), + }; + assert_eq!(stored(0), expected_first); + assert_eq!(stored(1), expected_second); + + // Negative, non-aligned timestamps must floor towards negative + // infinity on remote read, matching `Timestamp::convert_to` on the + // ingestion path: -1001us/-1001000ns -> -2ms, not -1ms. + if matches!(expected_unit, TimeUnit::Microsecond | TimeUnit::Nanosecond) { + let (ts_datatype, negative_value) = if expected_unit == TimeUnit::Microsecond { + ( + ColumnDataType::TimestampMicrosecond, + ValueData::TimestampMicrosecondValue(-1001), + ) + } else { + ( + ColumnDataType::TimestampNanosecond, + ValueData::TimestampNanosecondValue(-1_001_000), + ) + }; + let negative_request = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "metric1".to_string(), + rows: Some(Rows { + schema: vec![ + api::v1::ColumnSchema { + column_name: "greptime_timestamp".to_string(), + datatype: ts_datatype as i32, + semantic_type: SemanticType::Timestamp as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension: None, + options: None, + }, + api::v1::ColumnSchema { + column_name: "job".to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: SemanticType::Tag as i32, + datatype_extension: None, + options: None, + }, + ], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(negative_value), + }, + Value { + value_data: Some(ValueData::F64Value(3.0)), + }, + Value { + value_data: Some(ValueData::StringValue("spark".to_string())), + }, + ], + }], + }), + }], + }; + instance + .write(negative_request, ctx.clone(), true) + .await + .unwrap(); + + let read_request = ReadRequest { + queries: vec![Query { + start_timestamp_ms: -2, + end_timestamp_ms: 0, + matchers: vec![LabelMatcher { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "metric1".to_string(), + r#type: 0, + }], + ..Default::default() + }], + ..Default::default() + }; + let resp = instance.read(read_request, ctx.clone()).await.unwrap(); + let body = prom_store::snappy_decompress(&resp.body).unwrap(); + let read_response = ReadResponse::decode(&body[..]).unwrap(); + assert_eq!(1, read_response.results.len()); + assert_eq!( + read_response.results[0].timeseries[0].samples, + vec![Sample { + value: 3.0, + timestamp: -2, + }], + "remote read must floor negative non-aligned timestamps" + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_standalone_prom_store_write_existing_logical_table_of_other_physical_unit() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new( + "test_prom_store_write_existing_logical_other_physical_unit", + ) + .build() + .await; + let instance = standalone.fe_instance(); + + // Two physical metric tables with different time index units. + let db = "prometheus_mixed_units"; + let ctx = Arc::new(QueryContext::with(DEFAULT_CATALOG_NAME, db)); + assert!( + SqlQueryHandler::do_query( + instance.as_ref(), + &format!("CREATE DATABASE IF NOT EXISTS {db}"), + ctx.clone(), + ) + .await + .first() + .unwrap() + .is_ok() + ); + let mut output = instance + .do_query( + "CREATE TABLE phy_us (greptime_timestamp TIMESTAMP(6) NOT NULL, \ + greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + // First create the logical table `shared_metric` on the default + // (millisecond) physical table. + let shared_series = TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "shared_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }; + let (row_inserts, _) = to_grpc_row_insert_requests(&WriteRequest { + timeseries: vec![shared_series.clone()], + ..Default::default() + }) + .unwrap(); + instance + .write(row_inserts, ctx.clone(), true) + .await + .unwrap(); + + // Now select the microsecond physical table while writing to BOTH the + // existing millisecond `shared_metric` (bound to the default physical + // table) and a new table: the existing table's request must keep the + // millisecond unit, and only the new table uses microsecond. + let mut hint_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, db); + hint_ctx.set_extension(PHYSICAL_TABLE_PARAM, "phy_us".to_string()); + let hint_ctx = Arc::new(hint_ctx); + let (row_inserts, _) = to_grpc_row_insert_requests(&WriteRequest { + timeseries: vec![ + TimeSeries { + samples: vec![Sample { + value: 1.5, + timestamp: 1500, + }], + ..shared_series + }, + TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "fresh_us_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 2.5, + timestamp: 1500, + }], + ..Default::default() + }, + ], + ..Default::default() + }) + .unwrap(); + instance + .write(row_inserts, hint_ctx.clone(), true) + .await + .unwrap(); + + // The existing table keeps its millisecond unit and precision. + let mut output = instance + .do_query( + "SELECT greptime_timestamp FROM shared_metric ORDER BY greptime_timestamp", + ctx.clone(), + ) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + let shared_ts = batches[0] + .column(0) + .as_primitive::(); + assert_eq!((shared_ts.value(0), shared_ts.value(1)), (1000, 1500)); + + // The new table is created on the selected physical table with the + // microsecond unit. + let mut output = instance + .do_query( + "SELECT greptime_timestamp FROM fresh_us_metric", + ctx.clone(), + ) + .await; + let OutputData::Stream(stream) = output.remove(0).unwrap().data else { + unreachable!() + }; + let batches = common_recordbatch::RecordBatches::try_collect(stream) + .await + .unwrap() + .take(); + let fresh_ts = batches[0] + .column(0) + .as_primitive::(); + assert_eq!(fresh_ts.value(0), 1_500_000); + } + async fn test_prom_store_remote_rw(instance: &Arc, physical_table: Option) { let write_request = WriteRequest { timeseries: prom_store::mock_timeseries(), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index fcafbc0ce52..1b68d3720a7 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -15,18 +15,20 @@ use std::collections::BTreeMap; use std::io::Write; use std::str::FromStr; +use std::sync::Arc; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::metadata::MetricType as RemoteWriteV2MetricType; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Sample as RemoteWriteV2Sample, - TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Request as RemoteWriteV2Request, + Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ Label, LabelMatcher, Query, ReadRequest, ReadResponse, Sample, TimeSeries, WriteRequest, }; +use api::v1::RowInsertRequests; use auth::{UserProviderRef, user_provider_from_option}; use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode}; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -69,7 +71,9 @@ use servers::http::header::constants::{ }; use servers::http::header::{GREPTIME_DB_HEADER_NAME, GREPTIME_TIMEZONE_HEADER_NAME}; use servers::http::otlp::GoogleRpcStatus; -use servers::http::prometheus::{Column, PrometheusJsonResponse, PrometheusResponse}; +use servers::http::prometheus::{ + Column, PromQueryResult, PrometheusJsonResponse, PrometheusResponse, +}; use servers::http::result::error_result::ErrorResponse; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use servers::http::result::influxdb_result_v1::{InfluxdbOutput, InfluxdbV1Response}; @@ -79,11 +83,13 @@ use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store::{self, mock_timeseries_new_label}; use servers::request_memory_limiter::ServerMemoryLimiter; +use session::context::QueryContextRef; use standalone::options::StandaloneOptions; use table::table_name::TableName; use tests_integration::test_util::{ - MockInstanceImpl, StorageType, assert_wal_delta, build_test_prom_server, setup_test_http_app, - setup_test_http_app_with_frontend, setup_test_http_app_with_frontend_and_slow_query_threshold, + MockInstanceImpl, StorageType, TestGuard, assert_wal_delta, build_test_prom_server, + setup_test_http_app, setup_test_http_app_with_frontend, + setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, setup_test_prom_app_with_frontend_batched, }; @@ -157,6 +163,7 @@ macro_rules! http_tests { test_http_analyze_stream_tql, test_http_sql_slow_query, test_prometheus_promql_api, + test_promql_over_non_millisecond_physical_tables, test_prometheus_label_replace_response, test_prom_http_api, test_config_api, @@ -165,6 +172,9 @@ macro_rules! http_tests { test_prometheus_remote_write_v2, test_prometheus_remote_write_v2_native_histogram, test_prometheus_remote_write_batched, + test_prometheus_remote_write_batched_mixed_time_index_units, + test_prometheus_remote_write_batched_interceptor_time_index_units, + test_prometheus_remote_write_v2_batched_interceptor_time_index_units, test_prometheus_remote_special_labels, test_prometheus_remote_schema_labels, test_prometheus_remote_write_with_pipeline, @@ -1713,6 +1723,197 @@ pub async fn test_prom_http_api(store_type: StorageType) { guard.remove_all().await; } +/// PromQL must behave identically regardless of the physical metric table's +/// time index unit: the same samples are remote-written into physical tables +/// pre-created with second/micro/nano time indexes (plus a millisecond +/// baseline), and every PromQL response must match the baseline exactly. +/// Regression test for . +pub async fn test_promql_over_non_millisecond_physical_tables(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_prom_app_with_frontend(store_type, "promql_non_ms_units").await; + let client = TestClient::new(app).await; + + let units = [ + ("ms", "timestamp(3)"), + ("second", "timestamp(0)"), + ("micro", "timestamp(6)"), + ("nano", "timestamp(9)"), + ]; + + for (suffix, ts_type) in units { + let db = format!("promql_units_{suffix}"); + let res = client + .get(&format!("/v1/sql?db=public&sql=create database {db}")) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "create database {db}"); + + let res = client + .get(&format!( + "/v1/sql?db={db}&sql=CREATE TABLE greptime_physical_table \ + (greptime_timestamp {ts_type} NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + )) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::OK, + "create physical table with {ts_type} in {db}" + ); + + // unit_gauge{job="demo"} = 1.0 @ 1000ms, 2.0 @ 2000ms + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "unit_gauge".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![ + Sample { + value: 1.0, + timestamp: 1000, + }, + Sample { + value: 2.0, + timestamp: 2000, + }, + ], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let write_url = format!("/v1/prometheus/write?db={db}"); + let res = client + .post(write_url.as_str()) + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::NO_CONTENT, + "remote write into {db}" + ); + } + + let query_paths = [ + "query?query=unit_gauge&time=2", + "query_range?query=unit_gauge&start=0&end=5&step=1", + "query?query=avg_over_time(unit_gauge[1m])&time=2", + "query?query=rate(unit_gauge[1m])&time=2", + ]; + for path in query_paths { + let mut baseline: Option<(String, PrometheusResponse)> = None; + for (suffix, _) in units { + let db = format!("promql_units_{suffix}"); + let res = client + .get(&format!("/v1/prometheus/api/v1/{path}&db={db}")) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "promql `{path}` on {db}"); + let body = res.json::().await; + assert_eq!(body.status, "success", "promql `{path}` on {db}"); + match &baseline { + None => baseline = Some((suffix.to_string(), body.data)), + Some((baseline_suffix, baseline_data)) => { + assert_eq!( + &body.data, baseline_data, + "promql `{path}` on {db} differs from the {baseline_suffix} baseline" + ); + } + } + } + } + + // Pin the baseline responses so the cross-unit equality above cannot pass + // vacuously on wrong data: the samples are 1.0@1s and 2.0@2s, and + // `rate(unit_gauge[1m])@2s` follows Prometheus's extrapolated rate: + // Δv=1 over a 1s sampled interval extrapolated to 1.5s, divided by the + // 60s window -> 1.5/60 = 0.025 per second. + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=unit_gauge&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "value": [2.0, "2.0"] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query_range?db=promql_units_ms&query=unit_gauge&start=0&end=5&step=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "matrix", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "values": [ + [1.0, "1.0"], [2.0, "2.0"], [3.0, "2.0"], [4.0, "2.0"], [5.0, "2.0"] + ] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=avg_over_time(unit_gauge[1m])&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "value": [2.0, "1.5"] + }] + })) + .unwrap() + ); + + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=rate(unit_gauge[1m])&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + let rate = match body.data { + PrometheusResponse::PromData(data) => match data.result { + PromQueryResult::Vector(v) => v[0].value.as_ref().unwrap().1.clone(), + other => panic!("expected vector, got {other:?}"), + }, + other => panic!("expected prom data, got {other:?}"), + }; + assert_eq!(rate, "0.025"); + + guard.remove_all().await; +} + pub async fn test_metrics_api(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = setup_test_http_app(store_type, "metrics_api").await; @@ -3593,6 +3794,277 @@ async fn write_prometheus_skip_wal_sample( assert_eq!(request.send().await.status(), StatusCode::NO_CONTENT); } +/// Batched remote write against a logical table bound to a non-millisecond +/// physical table: the bulk guard must reject the destination (the bulk +/// encode only produces millisecond batches) and fall back to the ordinary +/// insert path, which converts the requests to the table's unit. +pub async fn test_prometheus_remote_write_batched_mixed_time_index_units(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_prom_app_with_frontend_batched(store_type, "prom_rw_batched_mixed_units").await; + let client = TestClient::new(app).await; + + let res = client + .get("/v1/sql?db=public&sql=CREATE TABLE phy_us (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) ENGINE = metric WITH ('physical_metric_table' = 'true')") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write = |metric: &str, value: f64, timestamp: i64| { + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: metric.to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { value, timestamp }], + ..Default::default() + }], + ..Default::default() + }; + prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap() + }; + + // Create the logical table on the microsecond physical table. The bulk + // guard rejects the microsecond selected physical table, so this write + // takes the ordinary insert path. + let res = client + .post("/v1/prometheus/write?physical_table=phy_us") + .header("Content-Encoding", "snappy") + .body(write("us_metric", 2.5, 1500)) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + // Write the same existing table again while selecting the default + // (millisecond) physical table: the guard must reject the existing + // microsecond destination and fall back to the ordinary insert path — + // without the destination check the bulk path would build millisecond + // arrays against the microsecond schema and fail the write. + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(write("us_metric", 3.5, 2000)) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + // Both samples are stored on the microsecond time index. + validate_data( + "prom_rw_batched_mixed_units", + &client, + "SELECT COUNT(*), MAX(greptime_value) FROM us_metric", + "[[2,3.5]]", + ) + .await; + + guard.remove_all().await; +} + +/// A Prometheus write interceptor that redirects every remote write to the +/// `tenant_redirect` schema, like a per-tenant redirection: the incoming +/// context targets `public` while the destination tables live elsewhere. +struct PromSchemaRedirectInterceptor; + +impl servers::interceptor::PromStoreProtocolInterceptor for PromSchemaRedirectInterceptor { + type Error = servers::error::Error; + + fn pre_write( + &self, + _write_req: &RowInsertRequests, + ctx: QueryContextRef, + ) -> servers::error::Result<()> { + ctx.set_current_schema("tenant_redirect"); + Ok(()) + } +} + +async fn setup_redirecting_batched_prom_app( + store_type: StorageType, + name: &str, +) -> (TestClient, TestGuard) { + let plugins = Plugins::default(); + plugins.insert::>( + Arc::new(PromSchemaRedirectInterceptor), + ); + let standalone = tests_integration::standalone::GreptimeDbStandaloneBuilder::new(name) + .with_default_store_type(store_type) + .with_plugin(plugins) + .build() + .await; + let server = build_test_prom_server(standalone.fe_instance().clone(), true) + .with_greptime_config_options(standalone.opts.datanode_options().to_toml().unwrap()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + (client, standalone.guard) +} + +async fn create_redirect_schema_with_microsecond_table(client: &TestClient, table: &str) { + for sql in [ + "create database if not exists tenant_redirect", + "CREATE TABLE tenant_redirect.phy_us \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('physical_metric_table' = 'true')", + &format!( + "CREATE TABLE tenant_redirect.{table} \ + (ts timestamp(6) time index, val double, host string primary key) \ + engine=metric with ('on_physical_table' = 'phy_us')" + ), + ] { + let res = client + .get(format!("/v1/sql?sql={sql}").as_str()) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK, "setup: {sql}"); + } +} + +/// Regression test: with batching enabled and a pre_write interceptor that +/// redirects the context to a schema holding same-named non-millisecond +/// tables, the bulk eligibility must be evaluated against the redirected +/// destinations (after preflight) and fall back to the ordinary insert path +/// instead of failing in the bulk encode. +pub async fn test_prometheus_remote_write_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_batched_interceptor_units").await; + create_redirect_schema_with_microsecond_table(&client, "intercept_metric").await; + + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "intercept_metric".to_string(), + }, + Label { + name: "job".to_string(), + value: "demo".to_string(), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.intercept_metric", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + +/// The v2 variant of the interceptor regression test: a mixed samples-plus- +/// histograms request whose existing sample destination lives in the +/// redirected schema with a microsecond time index must fall back to the +/// ordinary insert path (both series land), not fail in the bulk encode. +pub async fn test_prometheus_remote_write_v2_batched_interceptor_time_index_units( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (client, mut guard) = + setup_redirecting_batched_prom_app(store_type, "prom_rw_v2_batched_interceptor_units") + .await; + create_redirect_schema_with_microsecond_table(&client, "v2_mixed_sample").await; + + // One sample series targeting the existing microsecond table, plus one + // histogram series for a new (millisecond) table. + let mut symbols = vec![String::new()]; + let mut symbol = |value: &str| { + symbols.push(value.to_string()); + (symbols.len() - 1) as u32 + }; + let name_ref = symbol("__name__"); + let sample_ref = symbol("v2_mixed_sample"); + let job_ref = symbol("job"); + let demo_ref = symbol("demo"); + let name_ref2 = symbol("__name__"); + let histo_ref = symbol("v2_mixed_histo"); + let request = RemoteWriteV2Request { + symbols, + timeseries: vec![ + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref, sample_ref, job_ref, demo_ref], + samples: vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ..Default::default() + }, + RemoteWriteV2TimeSeries { + labels_refs: vec![name_ref2, histo_ref, job_ref, demo_ref], + histograms: vec![Histogram { + count: Some(Count::CountInt(1)), + sum: 1.0, + positive_spans: vec![BucketSpan { + offset: 0, + length: 1, + }], + positive_deltas: vec![1], + timestamp: 1000, + start_timestamp: 1000, + ..Default::default() + }], + ..Default::default() + }, + ], + }; + let compressed = prom_store::snappy_compress(&request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_sample", + "[[1]]", + ) + .await; + validate_data( + "prom_rw_v2_batched_interceptor_units", + &client, + "SELECT COUNT(*) FROM tenant_redirect.v2_mixed_histo", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + /// Covers the batched (pending-rows-batcher) Prometheus remote write path, which /// bypasses `PromStoreProtocolHandler::write`. Verifies the metric table is created /// asynchronously and still carries the Prometheus semantic identity stamped on the From 1d365adb3dce580f1e0185c2afdcb30a230b5744 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 28 Sep 2026 07:54:37 +0000 Subject: [PATCH 053/133] fix(ci): update the shared Actions runner to v2.337.0 (#9376) Signed-off-by: WenyXu --- .github/runner-scale-sets/query-regression/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index 9b6a8d5ebb6..d02425d684e 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -13,7 +13,7 @@ RUN git init /tmp/otelgen \ -o /usr/local/bin/otelgen ./cmd/otelgen FROM ghcr.io/astral-sh/uv@sha256:3d868e555f8f1dbc324afa005066cd11e1053fc4743b9808ca8025283e65efa5 AS uv -FROM ghcr.io/actions/actions-runner@sha256:08c30b0a7105f64bddfc485d2487a22aa03932a791402393352fdf674bda2c29 +FROM ghcr.io/actions/actions-runner:2.337.0@sha256:e5496277be5d09bc968b3d64911b74e219ac4a3f2edce956a3ecf9271bea1ef4 USER root From bc5ad4b416f9b90c27c2bca329e34512af7900c6 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 07:57:04 +0000 Subject: [PATCH 054/133] perf(mito2): batch index page loads and share cached bloom metadata (#9360) * perf(mito2): load missing index pages of all ranges in one read Signed-off-by: Dennis Zhuang * refactor(mito2): drop unused IndexCacheMetrics::merge Signed-off-by: Dennis Zhuang * perf(mito2): share cached bloom filter metadata instead of cloning it Signed-off-by: Dennis Zhuang * fix(mito2): only merge adjacent ranges when batch reading index pages Signed-off-by: Dennis Zhuang * fix(mito2): skip empty ranges when batch loading index pages Signed-off-by: Dennis Zhuang * perf(mito2): fetch index ranges concurrently Signed-off-by: Dennis Zhuang * test(index): return Arc from RecordingReader Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/index/src/bloom_filter/applier.rs | 5 +- src/index/src/bloom_filter/reader.rs | 7 +- src/mito2/src/cache/index.rs | 90 ++++++++++++++++--- .../src/cache/index/bloom_filter_index.rs | 35 +++----- src/mito2/src/cache/index/inverted_index.rs | 90 +++++++++++++++---- src/mito2/src/engine/puffin_index.rs | 7 +- src/mito2/src/sst/index/store.rs | 12 ++- 7 files changed, 186 insertions(+), 60 deletions(-) diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index ebd44cf84e6..fe1aafec2d6 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -14,6 +14,7 @@ use std::collections::BTreeSet; use std::ops::Range; +use std::sync::Arc; use fastbloom::BloomFilter; use greptime_proto::v1::index::BloomFilterMeta; @@ -38,7 +39,7 @@ pub struct InListPredicate { pub struct BloomFilterApplier { reader: Box, - meta: BloomFilterMeta, + meta: Arc, } impl BloomFilterApplier { @@ -401,7 +402,7 @@ mod tests { async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result { + ) -> Result> { self.inner.metadata(metrics).await } } diff --git a/src/index/src/bloom_filter/reader.rs b/src/index/src/bloom_filter/reader.rs index a9e08694e74..04eae84f005 100644 --- a/src/index/src/bloom_filter/reader.rs +++ b/src/index/src/bloom_filter/reader.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::ops::{Range, Rem}; +use std::sync::Arc; use std::time::{Duration, Instant}; use async_trait::async_trait; @@ -164,7 +165,7 @@ pub trait BloomFilterReader: Sync { async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result; + ) -> Result>; /// Reads a bloom filter with the given location. async fn bloom_filter( @@ -265,13 +266,13 @@ impl BloomFilterReader for BloomFilterReaderImpl { async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result { + ) -> Result> { let metadata = self.reader.metadata().await.context(IoSnafu)?; let file_size = metadata.content_length; let mut meta_reader = BloomFilterMetaReader::new(&self.reader, file_size, Some(DEFAULT_PREFETCH_SIZE)); - meta_reader.metadata(metrics).await + meta_reader.metadata(metrics).await.map(Arc::new) } } diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index 0cab159a7db..26f9a3d5f96 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -16,6 +16,7 @@ pub mod bloom_filter_index; pub mod inverted_index; pub mod result_cache; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::hash::Hash; use std::ops::Range; @@ -44,16 +45,6 @@ pub struct IndexCacheMetrics { pub page_bytes: u64, } -impl IndexCacheMetrics { - /// Merges another set of metrics into this one. - pub fn merge(&mut self, other: &Self) { - self.cache_hit += other.cache_hit; - self.cache_miss += other.cache_miss; - self.num_pages += other.num_pages; - self.page_bytes += other.page_bytes; - } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct PageKey { page_id: u64, @@ -243,6 +234,85 @@ where )) } + /// Like [`IndexCache::get_or_load`] for several ranges, but loads the missing pages of + /// all ranges with a single `load` call so a cold read costs one round trip instead of + /// one per range. + async fn get_or_load_vec( + &self, + key: K, + file_size: u64, + ranges: &[Range], + load: F, + ) -> Result<(Vec, IndexCacheMetrics), E> + where + F: FnOnce(Vec>) -> Fut, + Fut: Future, E>>, + E: std::error::Error, + { + let mut metrics = IndexCacheMetrics::default(); + let mut pages: HashMap = HashMap::new(); + let mut missing: Vec = Vec::new(); + let mut seen = HashSet::new(); + for range in ranges.iter().filter(|r| r.end > r.start) { + let size = (range.end - range.start) as u32; + for page_key in PageKey::generate_page_keys(range.start, size, self.page_size) { + if !seen.insert(page_key.page_id) { + continue; + } + metrics.num_pages += 1; + match self.get_page(key, page_key) { + Some(page) => { + CACHE_HIT.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); + metrics.cache_hit += 1; + metrics.page_bytes += page.len() as u64; + pages.insert(page_key.page_id, page); + } + None => { + CACHE_MISS.with_label_values(&[INDEX_CONTENT_TYPE]).inc(); + metrics.cache_miss += 1; + missing.push(page_key.page_id); + } + } + } + } + + if !missing.is_empty() { + missing.sort_unstable(); + let load_ranges = missing + .iter() + .map(|page_id| { + let start = page_id * self.page_size; + start..(start + self.page_size).min(file_size) + }) + .collect(); + let loaded = load(load_ranges).await?; + debug_assert_eq!(loaded.len(), missing.len()); + for (page_id, page) in missing.into_iter().zip(loaded) { + metrics.page_bytes += page.len() as u64; + self.put_page(key, PageKey { page_id }, page.clone()); + pages.insert(page_id, page); + } + } + + let result = ranges + .iter() + .map(|range| { + let size = (range.end - range.start) as u32; + if size == 0 { + return Bytes::new(); + } + let buffer = Buffer::from_iter( + PageKey::generate_page_keys(range.start, size, self.page_size) + .map(|p| pages[&p.page_id].clone()), + ); + buffer + .slice(PageKey::calculate_range(range.start, size, self.page_size)) + .to_bytes() + }) + .collect(); + Ok((result, metrics)) + } + fn get_page(&self, key: K, page_key: PageKey) -> Option { self.index.get(&(key, page_key)) } diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 707afb18a65..57f0103c1bb 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -154,25 +154,16 @@ impl BloomFilterReader for CachedBloomFilterIndexBl ) -> Result> { let start = metrics.as_ref().map(|_| Instant::now()); - let mut pages = Vec::with_capacity(ranges.len()); - let mut total_cache_metrics = crate::cache::index::IndexCacheMetrics::default(); - for range in ranges { - let inner = &self.inner; - let (page, cache_metrics) = self - .cache - .get_or_load( - (self.file_id, self.index_version, self.column_id, self.tag), - self.blob_size, - range.start, - (range.end - range.start) as u32, - move |ranges| async move { inner.read_vec(&ranges, None).await }, - ) - .await?; - - total_cache_metrics.merge(&cache_metrics); - pages.push(Bytes::from(page)); - } - + let inner = &self.inner; + let (pages, total_cache_metrics) = self + .cache + .get_or_load_vec( + (self.file_id, self.index_version, self.column_id, self.tag), + self.blob_size, + ranges, + move |ranges| async move { inner.read_vec(&ranges, None).await }, + ) + .await?; if let Some(m) = metrics { m.total_ranges += total_cache_metrics.num_pages; m.total_bytes += total_cache_metrics.page_bytes; @@ -190,7 +181,7 @@ impl BloomFilterReader for CachedBloomFilterIndexBl async fn metadata( &self, metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result { + ) -> Result> { if let Some(cached) = self.cache .get_metadata((self.file_id, self.index_version, self.column_id, self.tag)) @@ -199,12 +190,12 @@ impl BloomFilterReader for CachedBloomFilterIndexBl if let Some(m) = metrics { m.cache_hit += 1; } - Ok((*cached).clone()) + Ok(cached) } else { let meta = self.inner.metadata(metrics).await?; self.cache.put_metadata( (self.file_id, self.index_version, self.column_id, self.tag), - Arc::new(meta.clone()), + meta.clone(), ); CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc(); Ok(meta) diff --git a/src/mito2/src/cache/index/inverted_index.rs b/src/mito2/src/cache/index/inverted_index.rs index 4c1b07126c8..675ed736110 100644 --- a/src/mito2/src/cache/index/inverted_index.rs +++ b/src/mito2/src/cache/index/inverted_index.rs @@ -130,25 +130,16 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead ) -> Result> { let start = metrics.as_ref().map(|_| Instant::now()); - let mut pages = Vec::with_capacity(ranges.len()); - let mut total_cache_metrics = crate::cache::index::IndexCacheMetrics::default(); - for range in ranges { - let inner = &self.inner; - let (page, cache_metrics) = self - .cache - .get_or_load( - (self.file_id, self.index_version), - self.blob_size, - range.start, - (range.end - range.start) as u32, - move |ranges| async move { inner.read_vec(&ranges, None).await }, - ) - .await?; - - total_cache_metrics.merge(&cache_metrics); - pages.push(Bytes::from(page)); - } - + let inner = &self.inner; + let (pages, total_cache_metrics) = self + .cache + .get_or_load_vec( + (self.file_id, self.index_version), + self.blob_size, + ranges, + move |ranges| async move { inner.read_vec(&ranges, None).await }, + ) + .await?; if let Some(m) = metrics { m.total_bytes += total_cache_metrics.page_bytes; m.total_ranges += total_cache_metrics.num_pages; @@ -472,4 +463,65 @@ mod test { assert_eq!(read, expected); } } + + #[tokio::test] + async fn test_get_or_load_vec_loads_missing_pages_once() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let mut rng = rand::rng(); + let mut data = vec![0u8; 64 * 1024]; + rng.fill_bytes(&mut data); + let file_size = data.len() as u64; + let cache = InvertedIndexCache::new(1024 * 1024, 1024 * 1024, 1000); + let key = (FileId::random(), 0); + + for _ in 0..FUZZ_REPEAT_TIMES { + cache.invalidate_file(key.0); + let ranges = (0..rng.random_range(1..20)) + .map(|_| { + let start = rng.random_range(0..file_size); + start..rng.random_range(start..=file_size.min(start + 5000)) + }) + .collect::>(); + let expected = ranges + .iter() + .map(|r| bytes::Bytes::copy_from_slice(&data[r.start as usize..r.end as usize])) + .collect::>(); + + let loads = AtomicUsize::new(0); + let load = |pages: Vec>| { + loads.fetch_add(1, Ordering::Relaxed); + // Pages are requested once each, in file order. + assert!(pages.windows(2).all(|w| w[0].end <= w[1].start)); + let data = &data; + async move { + Ok::<_, std::io::Error>( + pages + .iter() + .map(|r| { + bytes::Bytes::copy_from_slice( + &data[r.start as usize..r.end as usize], + ) + }) + .collect(), + ) + } + }; + let (cold, cold_metrics) = cache + .get_or_load_vec(key, file_size, &ranges, load) + .await + .unwrap(); + assert_eq!(cold, expected); + let cold_loads = usize::from(cold_metrics.cache_miss > 0); + assert_eq!(loads.load(Ordering::Relaxed), cold_loads); + + let (warm, metrics) = cache + .get_or_load_vec(key, file_size, &ranges, load) + .await + .unwrap(); + assert_eq!(warm, expected); + assert_eq!(metrics.cache_miss, 0); + assert_eq!(loads.load(Ordering::Relaxed), cold_loads); + } + } } diff --git a/src/mito2/src/engine/puffin_index.rs b/src/mito2/src/engine/puffin_index.rs index 5624405dd7e..278d9017e7f 100644 --- a/src/mito2/src/engine/puffin_index.rs +++ b/src/mito2/src/engine/puffin_index.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::convert::TryFrom; +use std::sync::Arc; use common_base::range_read::RangeReader; use common_telemetry::warn; @@ -111,7 +112,7 @@ pub(crate) async fn collect_index_entries_from_puffin( ) .await; - let bloom_value = bloom_meta.as_ref().map(bloom_meta_value); + let bloom_value = bloom_meta.as_deref().map(bloom_meta_value); let (target_type, target_json) = decode_target_info(target_key); let meta_json = build_meta_json(bloom_value, None, None); let entry = build_index_entry( @@ -137,7 +138,7 @@ pub(crate) async fn collect_index_entries_from_puffin( ) .await; - let bloom_value = bloom_meta.as_ref().map(bloom_meta_value); + let bloom_value = bloom_meta.as_deref().map(bloom_meta_value); let fulltext_value = Some(fulltext_meta_value(blob)); let (target_type, target_json) = decode_target_info(target_key); let meta_json = build_meta_json(bloom_value, fulltext_value, None); @@ -294,7 +295,7 @@ async fn try_read_bloom_meta( cache: Option<&BloomFilterIndexCacheRef>, tag: Tag, context: &IndexEntryContext<'_>, -) -> Option { +) -> Option> { let column_id = decode_column_id(target_key); // Failures are logged but do not abort the overall metadata collection. diff --git a/src/mito2/src/sst/index/store.rs b/src/mito2/src/sst/index/store.rs index e737252652d..0a99744986a 100644 --- a/src/mito2/src/sst/index/store.rs +++ b/src/mito2/src/sst/index/store.rs @@ -262,6 +262,9 @@ impl AsyncWrite for InstrumentedAsyncWrite<'_, W> } } +/// Ranges fetched in parallel by one `read_vec`, as for parquet reads. +const FETCH_CONCURRENCY: usize = 8; + /// Implements `RangeReader` for `ObjectStore` and record metrics. pub(crate) struct InstrumentedRangeReader<'a> { store: ObjectStore, @@ -309,9 +312,16 @@ impl RangeReader for InstrumentedRangeReader<'_> { } async fn read_vec(&self, ranges: &[Range]) -> io::Result> { + // OpenDAL merges ranges less than 1 MiB apart by default without bounding the merged + // length, so sparse page reads could fetch most of the file. Only merge adjacent + // ranges, and fetch the rest in parallel like parquet reads do. + // TODO: decide whether to merge across small gaps based on object store latency, + // bytes read and query latency. let bufs = self .store - .reader(&self.path) + .reader_with(&self.path) + .gap(0) + .concurrent(FETCH_CONCURRENCY) .await? .fetch(ranges.to_owned()) .await?; From ffdd6d09a67f7c71bc109229c41f91889c1df2a0 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 08:22:42 +0000 Subject: [PATCH 055/133] perf(index): cut allocations when building bloom and inverted indexes (#9359) * perf(index): build bloom filters from element hashes without per-token allocation Signed-off-by: Dennis Zhuang * perf(index): speed up inverted index building with hashed buffers and sync pushes Signed-off-by: Dennis Zhuang * chore(index): use BuildHasher::hash_one for element hashes Signed-off-by: Dennis Zhuang * chore(index): require callers to act on spill requests Signed-off-by: Dennis Zhuang * fix(index): insert bloom hashes one by one to keep segment set capacity bounded Signed-off-by: Dennis Zhuang * test(index): add index build and bloom search benchmarks Signed-off-by: Dennis Zhuang * test(index): keep applier setup out of the bloom search benchmark Signed-off-by: Dennis Zhuang * fix(index): skip empty spills and keep the old inverted sort memory estimate Signed-off-by: Dennis Zhuang * fix(index): stream fulltext token hashes and test spill dispatch Signed-off-by: Dennis Zhuang * docs(index): note non-ASCII tokens still allocate in analyze_text_hashes Signed-off-by: Dennis Zhuang * docs(index): narrow the allocation note to case-insensitive non-ASCII tokens Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.lock | 1 + src/index/Cargo.toml | 5 + src/index/benches/index_build_bench.rs | 424 ++++++++++++++++++ src/index/src/bloom_filter.rs | 68 +++ src/index/src/bloom_filter/applier.rs | 17 +- src/index/src/bloom_filter/creator.rs | 164 ++++--- .../bloom_filter/creator/finalize_segment.rs | 21 +- src/index/src/bloom_filter/reader.rs | 7 +- .../src/fulltext_index/create/bloom_filter.rs | 5 +- src/index/src/fulltext_index/tokenizer.rs | 47 ++ src/index/src/inverted_index/create.rs | 34 +- src/index/src/inverted_index/create/sort.rs | 16 +- .../create/sort/external_sort.rs | 131 ++++-- .../src/inverted_index/create/sort_create.rs | 115 +++-- .../src/sst/index/inverted_index/creator.rs | 52 ++- 15 files changed, 879 insertions(+), 228 deletions(-) create mode 100644 src/index/benches/index_build_bench.rs diff --git a/Cargo.lock b/Cargo.lock index 597ad270305..bb7f69d1551 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6813,6 +6813,7 @@ dependencies = [ name = "index" version = "1.3.0-alpha.1" dependencies = [ + "ahash 0.8.12", "async-trait", "asynchronous-codec", "bytemuck", diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index bb1225c12c8..078f9e42c56 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true workspace = true [dependencies] +ahash.workspace = true async-trait.workspace = true asynchronous-codec = "0.7.0" bytemuck.workspace = true @@ -57,3 +58,7 @@ harness = false [[bench]] name = "bytes_to_u64_vec" harness = false + +[[bench]] +name = "index_build_bench" +harness = false diff --git a/src/index/benches/index_build_bench.rs b/src/index/benches/index_build_bench.rs new file mode 100644 index 00000000000..4c6c3fac020 --- /dev/null +++ b/src/index/benches/index_build_bench.rs @@ -0,0 +1,424 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Index build and bloom search benchmarks over synthetic observability data. + +use std::collections::{BTreeSet, HashMap}; +use std::hint::black_box; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; + +use async_trait::async_trait; +use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main}; +use futures::{AsyncRead, AsyncReadExt}; +use index::bitmap::BitmapType; +use index::bloom_filter::applier::{BloomFilterApplier, InListPredicate}; +use index::bloom_filter::creator::BloomFilterCreator; +use index::bloom_filter::reader::BloomFilterReaderImpl; +use index::external_provider::{ExternalTempFileProvider, Reader, Writer}; +use index::fulltext_index::Config; +use index::fulltext_index::create::{BloomFilterFulltextIndexCreator, FulltextIndexCreator}; +use index::fulltext_index::tokenizer::{Analyzer, EnglishTokenizer}; +use index::inverted_index::create::InvertedIndexCreator; +use index::inverted_index::create::sort::external_sort::ExternalSorter; +use index::inverted_index::create::sort_create::SortIndexCreator; +use index::inverted_index::format::writer::InvertedIndexBlobWriter; +use puffin::puffin_manager::{PuffinWriter, PutOptions}; +use rand::seq::IndexedRandom; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +const ROWS: usize = 100_000; +const SEGMENT_ROWS: usize = 10240; +const ROW_GROUP_ROWS: usize = 102400; + +/// Drains blobs so the bloom fulltext creator can finish without a puffin file. +struct DrainPuffinWriter; + +#[async_trait] +impl PuffinWriter for DrainPuffinWriter { + async fn put_blob( + &mut self, + _key: &str, + raw_data: R, + _options: PutOptions, + _properties: HashMap, + ) -> puffin::error::Result + where + R: AsyncRead + Send, + { + let mut buf = Vec::new(); + Box::pin(raw_data).read_to_end(&mut buf).await.unwrap(); + Ok(buf.len() as u64) + } + + async fn put_dir( + &mut self, + _key: &str, + _dir: PathBuf, + _options: PutOptions, + _properties: HashMap, + ) -> puffin::error::Result { + unreachable!("bloom fulltext index only writes blobs") + } + + fn set_footer_lz4_compressed(&mut self, _lz4_compressed: bool) {} + + async fn finish(self) -> puffin::error::Result { + Ok(0) + } +} + +/// The benchmarks set no memory limit, so nothing is spilled. +struct NoSpill; + +#[async_trait] +impl ExternalTempFileProvider for NoSpill { + async fn create(&self, _: &str, _: &str) -> Result { + unreachable!("no memory limit is set") + } + + async fn read_all(&self, _: &str) -> Result, index::error::Error> { + Ok(vec![]) + } +} + +fn uuid(rng: &mut ChaCha8Rng) -> String { + let h = format!("{:032x}", rng.random::()); + format!( + "{}-{}-{}-{}-{}", + &h[0..8], + &h[8..12], + &h[12..16], + &h[16..20], + &h[20..32] + ) +} + +/// Access, auth, timeout and error log lines with request ids, IPs and numbers: most +/// tokens of a segment are distinct. +fn service_logs(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(1); + let paths = ["users", "orders", "items", "carts", "payments", "sessions"]; + (0..rows) + .map(|_| { + let ip = format!( + "10.{}.{}.{}", + rng.random_range(0..16), + rng.random_range(0..256), + rng.random_range(1..255) + ); + match rng.random_range(0..10) { + 0..=5 => format!( + "INFO GET /api/v1/{}/{}/detail?page={} 200 {}ms request_id={} client={}", + paths.choose(&mut rng).unwrap(), + rng.random_range(0..200000), + rng.random_range(0..50), + rng.random_range(1..2000), + uuid(&mut rng), + ip + ), + 6..=7 => format!( + "WARN connection to {}:{} timed out after {}ms retry={}", + ip, + rng.random_range(1024..65535), + rng.random_range(100..30000), + rng.random_range(0..5) + ), + _ => format!( + "ERROR failed to process message {} from queue orders-{}: deadline exceeded", + uuid(&mut rng), + rng.random_range(0..64) + ), + } + }) + .collect() +} + +/// Java stack traces: long lines where a few frame names repeat many times. +fn stack_trace_logs(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(2); + let frames = [ + "at org.apache.kafka.clients.consumer.KafkaConsumer.poll(KafkaConsumer.java:1250)", + "at com.example.orders.OrderService.process(OrderService.java:88)", + "at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)", + "at java.base/java.lang.Thread.run(Thread.java:833)", + ]; + (0..rows) + .map(|_| { + let mut line = format!( + "ERROR request {} failed: java.lang.IllegalStateException: retry exhausted", + uuid(&mut rng) + ); + for _ in 0..40 { + line.push(' '); + line.push_str(frames.choose(&mut rng).unwrap()); + } + line + }) + .collect() +} + +/// Trace ids with temporal locality: each trace spreads over about a hundred rows. +fn trace_ids(rows: usize) -> Vec { + let mut rng = ChaCha8Rng::seed_from_u64(3); + let traces = (0..rows / 10) + .map(|_| format!("{:032x}", rng.random::())) + .collect::>(); + (0..rows) + .map(|i| { + let j = (i / 10) as i64 + rng.random_range(-50i64..=50); + traces[j.clamp(0, traces.len() as i64 - 1) as usize].clone() + }) + .collect() +} + +/// Prometheus-like series sorted by labels; each series has `samples` rows. +fn series(num_series: usize) -> Vec<[String; 6]> { + let mut rng = ChaCha8Rng::seed_from_u64(4); + let mut series = (0..num_series) + .map(|s| { + let pod = s / 2; + let ns = pod % 400; + [ + format!("region-{}", ns % 4), + format!("az-{}", ns % 12), + format!("cluster-{}", ns % 40), + format!("namespace-{ns}"), + format!("pod-{:05x}-{pod}", rng.random_range(0..0xfffff)), + format!("container-{}", s % 2), + ] + }) + .collect::>(); + series.sort(); + series +} + +fn bloom_creator() -> BloomFilterCreator { + BloomFilterCreator::new( + SEGMENT_ROWS, + 0.01, + Arc::new(NoSpill), + Arc::new(AtomicUsize::new(0)), + None, + ) +} + +fn bench_bloom_fulltext_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_fulltext_build"); + group.sample_size(10); + + for (name, lines) in [ + ("service_logs", service_logs(ROWS)), + ("stack_traces", stack_trace_logs(ROWS / 10)), + ] { + group.throughput(Throughput::Elements(lines.len() as u64)); + group.bench_function(name, |b| { + b.iter_batched( + || { + BloomFilterFulltextIndexCreator::new( + Config::default(), + SEGMENT_ROWS, + 0.01, + Arc::new(NoSpill), + Arc::new(AtomicUsize::new(0)), + None, + ) + }, + |mut creator| { + let lines = &lines; + rt.block_on(async move { + for line in lines { + creator.push_text(line).await.unwrap(); + } + let written = creator + .finish(&mut DrainPuffinWriter, "blob", PutOptions::default()) + .await + .unwrap(); + black_box(written) + }) + }, + BatchSize::LargeInput, + ) + }); + } + group.finish(); +} + +fn bench_bloom_skipping_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_skipping_build"); + group.sample_size(10); + + let trace_ids = trace_ids(ROWS * 10); + group.throughput(Throughput::Elements(trace_ids.len() as u64)); + group.bench_function("trace_id", |b| { + b.iter_batched( + bloom_creator, + |mut creator| { + let trace_ids = &trace_ids; + rt.block_on(async move { + for id in trace_ids { + creator + .push_n_row_elem(1, Some(id.as_bytes())) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + + // A sorted tag pushed as one run per series, as the flat SST path does. + let series = series(ROWS); + let samples = 10; + group.throughput(Throughput::Elements((series.len() * samples) as u64)); + group.bench_function("sorted_tag_runs", |b| { + b.iter_batched( + bloom_creator, + |mut creator| { + let series = &series; + rt.block_on(async move { + for s in series { + creator + .push_n_row_elem(samples, Some(s[4].as_bytes())) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + group.finish(); +} + +fn bench_inverted_build(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("inverted_build"); + group.sample_size(10); + + let series = series(ROWS); + let samples = 10; + let names = ["region", "az", "cluster", "namespace", "pod", "container"]; + group.throughput(Throughput::Elements((series.len() * samples) as u64)); + group.bench_function("six_tags", |b| { + b.iter_batched( + || { + let factory = ExternalSorter::factory( + Arc::new(NoSpill), + None, + Arc::new(AtomicUsize::new(0)), + None, + ); + SortIndexCreator::new(factory, NonZeroUsize::new(1024).unwrap()) + }, + |mut creator| { + let series = &series; + rt.block_on(async move { + for s in series { + for (name, value) in names.iter().zip(s) { + let spill = + creator.push_with_name_n(name, Some(value.as_bytes()), samples); + assert!(!spill, "no memory limit is set"); + } + } + let mut out = Vec::new(); + let mut writer = InvertedIndexBlobWriter::new(&mut out); + creator + .finish(&mut writer, BitmapType::Roaring) + .await + .unwrap(); + black_box(out) + }) + }, + BatchSize::LargeInput, + ) + }); + group.finish(); +} + +#[allow(clippy::single_range_in_vec_init)] +fn bench_bloom_search(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bloom_search"); + + let lines = service_logs(ROWS * 10); + let analyzer = Analyzer::new(Box::new(EnglishTokenizer), false); + let blob = rt.block_on(async { + let mut creator = bloom_creator(); + for line in &lines { + creator + .push_row_elems(analyzer.analyze_text(line).unwrap()) + .await + .unwrap(); + } + let mut out = Vec::new(); + creator.finish(&mut out).await.unwrap(); + out + }); + + let num_rows = lines.len(); + // Each row group tests every segment's filter against the three probes. + let predicates = ["timed", "deadline", "nonexistent"] + .iter() + .map(|term| InListPredicate { + list: BTreeSet::from([term.as_bytes().to_vec()]), + }) + .collect::>(); + group.bench_function("three_terms_per_row_group", |b| { + b.iter_batched( + || { + let reader = BloomFilterReaderImpl::new(blob.clone()); + rt.block_on(BloomFilterApplier::new(Box::new(reader))) + .unwrap() + }, + |mut applier| { + rt.block_on(async { + let mut matched = 0; + for start in (0..num_rows).step_by(ROW_GROUP_ROWS) { + let end = (start + ROW_GROUP_ROWS).min(num_rows); + matched += applier + .search(&predicates, &[start..end], None) + .await + .unwrap() + .len(); + } + black_box(matched) + }) + }, + BatchSize::SmallInput, + ) + }); + group.finish(); +} + +criterion_group!( + benches, + bench_bloom_fulltext_build, + bench_bloom_skipping_build, + bench_inverted_build, + bench_bloom_search +); +criterion_main!(benches); diff --git a/src/index/src/bloom_filter.rs b/src/index/src/bloom_filter.rs index eb818a0f5a0..156be09e69e 100644 --- a/src/index/src/bloom_filter.rs +++ b/src/index/src/bloom_filter.rs @@ -17,5 +17,73 @@ pub mod creator; pub mod error; pub mod reader; +use std::hash::{BuildHasher, BuildHasherDefault, Hasher}; +use std::sync::LazyLock; + /// The seed used for the Bloom filter. pub const SEED: u128 = 42; + +static ELEMENT_HASHER: LazyLock = + LazyLock::new(|| fastbloom::DefaultHasher::seeded(&SEED.to_be_bytes())); + +/// Returns the hash fastbloom derives for `elem` with the persisted [`SEED`]. +/// +/// A filter built from these hashes with [`PrehashedBuildHasher`] has exactly the same +/// bits as one built by inserting the elements themselves, so files stay readable by +/// both paths. +pub fn element_hash(elem: &[u8]) -> u64 { + ELEMENT_HASHER.hash_one(elem) +} + +/// Hasher that returns an already computed [`element_hash`] unchanged. +#[derive(Default)] +pub struct PrehashedHasher(u64); + +impl Hasher for PrehashedHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, _bytes: &[u8]) { + unreachable!("PrehashedHasher only accepts u64 hashes") + } + + fn write_u64(&mut self, hash: u64) { + self.0 = hash; + } +} + +pub type PrehashedBuildHasher = BuildHasherDefault; + +/// A persisted bloom filter probed by [`element_hash`] values. +pub type PrehashedBloomFilter = fastbloom::BloomFilter<512, PrehashedBuildHasher>; + +#[cfg(test)] +mod tests { + use fastbloom::BloomFilter; + + use super::*; + + #[test] + fn test_prehashed_filter_matches_seeded_filter() { + // Persisted filters are read back with `.seed(&SEED)`, so building them from + // precomputed hashes must set exactly the same bits. + for count in [0usize, 1, 7, 100, 5000] { + let elems = (0..count) + .map(|i| format!("elem-{i}").into_bytes()) + .chain([Vec::new()]) + .collect::>(); + let mut seeded = BloomFilter::with_false_pos(0.01) + .seed(&SEED) + .expected_items(elems.len()); + let mut prehashed = BloomFilter::with_false_pos(0.01) + .hasher(PrehashedBuildHasher::default()) + .expected_items(elems.len()); + for elem in &elems { + seeded.insert(elem); + prehashed.insert(&element_hash(elem)); + } + assert_eq!(seeded.as_slice(), prehashed.as_slice()); + } + } +} diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index fe1aafec2d6..96a1c6ff17f 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -16,13 +16,13 @@ use std::collections::BTreeSet; use std::ops::Range; use std::sync::Arc; -use fastbloom::BloomFilter; use greptime_proto::v1::index::BloomFilterMeta; use itertools::Itertools; use crate::Bytes; use crate::bloom_filter::error::Result; use crate::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; +use crate::bloom_filter::{PrehashedBloomFilter, element_hash}; /// Filter bytes one batch of [`BloomFilterApplier::search_groups`] reads. A single row /// group larger than this is still read as one batch, so this is not a memory limit; a @@ -190,7 +190,7 @@ impl BloomFilterApplier { &mut self, segments: &[usize], metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result<(Vec<(u64, usize)>, Vec)> { + ) -> Result<(Vec<(u64, usize)>, Vec)> { let segment_locations = segments .iter() .map(|&seg| (self.meta.segment_loc_indices[seg], seg)) @@ -215,11 +215,15 @@ impl BloomFilterApplier { fn find_matching_rows( &self, segment_locations: Vec<(u64, usize)>, - bloom_filters: Vec, + bloom_filters: Vec, predicates: &[InListPredicate], ) -> Vec> { let rows_per_segment = self.meta.rows_per_segment as usize; let mut matching_row_ranges = Vec::with_capacity(bloom_filters.len()); + let predicate_hashes = predicates + .iter() + .map(|p| p.list.iter().map(|v| element_hash(v)).collect::>()) + .collect::>(); // Group segments by their location index (since they have the same bloom filter) and check if they match all predicates for ((_loc_index, group), bloom_filter) in segment_locations @@ -229,12 +233,9 @@ impl BloomFilterApplier { .zip(bloom_filters.iter()) { // Check if this bloom filter matches each predicate (AND semantics) - let matches_all_predicates = predicates.iter().all(|predicate| { + let matches_all_predicates = predicate_hashes.iter().all(|hashes| { // For each predicate, at least one probe must match (OR semantics) - predicate - .list - .iter() - .any(|probe| bloom_filter.contains(probe)) + hashes.iter().any(|hash| bloom_filter.contains(hash)) }); if !matches_all_predicates { diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index 07f4f977caf..83c0ae0c697 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/src/index/src/bloom_filter/creator.rs @@ -26,8 +26,8 @@ use prost::Message; use snafu::ResultExt; use crate::Bytes; -use crate::bloom_filter::SEED; use crate::bloom_filter::error::{IoSnafu, Result}; +use crate::bloom_filter::{PrehashedBuildHasher, element_hash}; use crate::external_provider::ExternalTempFileProvider; /// `BloomFilterCreator` is responsible for creating and managing bloom filters @@ -52,8 +52,11 @@ pub struct BloomFilterCreator { /// Row count that added to the bloom filter so far. accumulated_row_count: usize, - /// A set of distinct elements in the current segment. - cur_seg_distinct_elems: HashSet, + /// Distinct element hashes (see [`element_hash`]) in the current segment. + /// + /// Elements with equal hashes set the same bits, so deduplicating by hash loses + /// nothing and avoids copying the values. + cur_seg_distinct_elems: HashSet, /// The memory usage of the current segment's distinct elements. cur_seg_distinct_elems_mem_usage: usize, @@ -106,98 +109,36 @@ impl BloomFilterCreator { /// reaches `rows_per_segment`, it finalizes the current segment. pub async fn push_n_row_elems( &mut self, - mut nrows: usize, + nrows: usize, elems: impl IntoIterator, ) -> Result<()> { - if nrows == 0 { - return Ok(()); - } - if nrows == 1 { - return self.push_row_elems(elems).await; - } - - let elems = elems.into_iter().collect::>(); - while nrows > 0 { - let rows_to_seg_end = - self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); - let rows_to_push = nrows.min(rows_to_seg_end); - nrows -= rows_to_push; - - self.accumulated_row_count += rows_to_push; - - let mut mem_diff = 0; - for elem in &elems { - let len = elem.len(); - let is_new = self.cur_seg_distinct_elems.insert(elem.clone()); - if is_new { - mem_diff += len; - } - } - self.cur_seg_distinct_elems_mem_usage += mem_diff; - self.global_memory_usage - .fetch_add(mem_diff, Ordering::Relaxed); - - if self - .accumulated_row_count - .is_multiple_of(self.rows_per_segment) - { - self.finalize_segment().await?; - self.finalized_row_count = self.accumulated_row_count; - } - } - - Ok(()) + let hashes = elems + .into_iter() + .map(|e| element_hash(&e)) + .collect::>(); + self.push_n_row_hashes(nrows, &hashes).await } - /// Adds `nrows` copies of a single borrowed value (or null), copying it only when - /// it is new to a segment. Row counts advance for nulls as well. - pub async fn push_n_row_elem(&mut self, mut nrows: usize, elem: Option<&[u8]>) -> Result<()> { - while nrows > 0 { - let rows_to_seg_end = - self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); - let rows_to_push = nrows.min(rows_to_seg_end); - nrows -= rows_to_push; - self.accumulated_row_count += rows_to_push; - - if let Some(elem) = elem { - let old_len = self.cur_seg_distinct_elems.len(); - // Only allocate when the value is absent from the current segment. - if !self.cur_seg_distinct_elems.contains(elem) { - self.cur_seg_distinct_elems.insert(elem.to_vec()); - } - if self.cur_seg_distinct_elems.len() != old_len { - self.cur_seg_distinct_elems_mem_usage += elem.len(); - self.global_memory_usage - .fetch_add(elem.len(), Ordering::Relaxed); - } - } - if self - .accumulated_row_count - .is_multiple_of(self.rows_per_segment) - { - self.finalize_segment().await?; - self.finalized_row_count = self.accumulated_row_count; - } + /// Adds `nrows` copies of a single borrowed value (or null). Row counts advance for + /// nulls as well. + pub async fn push_n_row_elem(&mut self, nrows: usize, elem: Option<&[u8]>) -> Result<()> { + match elem { + Some(elem) => self.push_n_row_hashes(nrows, &[element_hash(elem)]).await, + None => self.push_n_row_hashes(nrows, &[]).await, } - Ok(()) } /// Adds a row of elements to the bloom filter. If the number of accumulated rows /// reaches `rows_per_segment`, it finalizes the current segment. pub async fn push_row_elems(&mut self, elems: impl IntoIterator) -> Result<()> { - self.accumulated_row_count += 1; + self.push_row_hashes(elems.into_iter().map(|e| element_hash(&e))) + .await + } - let mut mem_diff = 0; - for elem in elems.into_iter() { - let len = elem.len(); - let is_new = self.cur_seg_distinct_elems.insert(elem); - if is_new { - mem_diff += len; - } - } - self.cur_seg_distinct_elems_mem_usage += mem_diff; - self.global_memory_usage - .fetch_add(mem_diff, Ordering::Relaxed); + /// Adds a row of element hashes computed by [`element_hash`]. + pub async fn push_row_hashes(&mut self, hashes: impl IntoIterator) -> Result<()> { + self.accumulated_row_count += 1; + self.insert_hashes(hashes); if self .accumulated_row_count @@ -210,6 +151,43 @@ impl BloomFilterCreator { Ok(()) } + /// Adds `nrows` rows that all contain the element hashes in `hashes`. + pub async fn push_n_row_hashes(&mut self, mut nrows: usize, hashes: &[u64]) -> Result<()> { + while nrows > 0 { + let rows_to_seg_end = + self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); + let rows_to_push = nrows.min(rows_to_seg_end); + nrows -= rows_to_push; + self.accumulated_row_count += rows_to_push; + self.insert_hashes(hashes.iter().copied()); + + if self + .accumulated_row_count + .is_multiple_of(self.rows_per_segment) + { + self.finalize_segment().await?; + self.finalized_row_count = self.accumulated_row_count; + } + } + + Ok(()) + } + + fn insert_hashes(&mut self, hashes: impl IntoIterator) { + let old_len = self.cur_seg_distinct_elems.len(); + // Not `extend`: it reserves for the iterator's length, which counts duplicate + // tokens, and the capacity survives `drain` at segment boundaries. + for hash in hashes { + self.cur_seg_distinct_elems.insert(hash); + } + let mem_diff = (self.cur_seg_distinct_elems.len() - old_len) * size_of::(); + if mem_diff > 0 { + self.cur_seg_distinct_elems_mem_usage += mem_diff; + self.global_memory_usage + .fetch_add(mem_diff, Ordering::Relaxed); + } + } + /// Finalizes any remaining segments and writes the bloom filters and metadata to the provided writer. pub async fn finish(&mut self, mut writer: impl AsyncWrite + Unpin) -> Result<()> { if self.accumulated_row_count > self.finalized_row_count { @@ -286,6 +264,7 @@ mod tests { use futures::io::Cursor; use super::*; + use crate::bloom_filter::SEED; use crate::external_provider::MockExternalTempFileProvider; /// Converts a slice of bytes to a vector of `u64`. @@ -296,6 +275,23 @@ mod tests { .collect() } + #[tokio::test] + async fn test_duplicate_hashes_do_not_grow_segment_set() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + creator + .push_row_hashes(std::iter::repeat_n(7, 1_000_000)) + .await + .unwrap(); + assert_eq!(creator.cur_seg_distinct_elems.len(), 1); + assert!(creator.cur_seg_distinct_elems.capacity() < 16); + } + #[tokio::test] async fn test_bloom_filter_creator() { let mut writer = Cursor::new(Vec::new()); diff --git a/src/index/src/bloom_filter/creator/finalize_segment.rs b/src/index/src/bloom_filter/creator/finalize_segment.rs index dd7d03864e4..4cbce66edce 100644 --- a/src/index/src/bloom_filter/creator/finalize_segment.rs +++ b/src/index/src/bloom_filter/creator/finalize_segment.rs @@ -22,8 +22,7 @@ use futures::stream::StreamExt; use futures::{AsyncWriteExt, Stream, stream}; use snafu::ResultExt; -use crate::Bytes; -use crate::bloom_filter::creator::SEED; +use crate::bloom_filter::PrehashedBuildHasher; use crate::bloom_filter::creator::intermediate_codec::IntermediateBloomFilterCodecV1; use crate::bloom_filter::error::{IntermediateSnafu, IoSnafu, Result}; use crate::external_provider::ExternalTempFileProvider; @@ -98,14 +97,14 @@ impl FinalizedBloomFilterStorage { /// If the memory usage exceeds the threshold, flushes the in-memory Bloom filters to disk. pub async fn add( &mut self, - elems: impl IntoIterator, + elem_hashes: impl IntoIterator, element_count: usize, ) -> Result<()> { let mut bf = BloomFilter::with_false_pos(self.false_positive_rate) - .seed(&SEED) + .hasher(PrehashedBuildHasher::default()) .expected_items(element_count); - for elem in elems.into_iter() { - bf.insert(&elem); + for hash in elem_hashes.into_iter() { + bf.insert(&hash); } let fbf = FinalizedBloomFilterSegment::from(bf, element_count); @@ -231,7 +230,7 @@ pub struct FinalizedBloomFilterSegment { } impl FinalizedBloomFilterSegment { - fn from(bf: BloomFilter, elem_count: usize) -> Self { + fn from(bf: BloomFilter<512, S>, elem_count: usize) -> Self { let bf_slice = bf.as_slice(); let mut bloom_filter_bytes = Vec::with_capacity(std::mem::size_of_val(bf_slice)); for &x in bf_slice { @@ -256,6 +255,7 @@ mod tests { use super::*; use crate::bloom_filter::creator::tests::u64_vec_from_bytes; + use crate::bloom_filter::{SEED, element_hash}; use crate::external_provider::MockExternalTempFileProvider; #[tokio::test] @@ -300,11 +300,12 @@ mod tests { let dup_batch = 200; for i in 0..(batch - dup_batch) { - let elems = (elem_count * i..elem_count * (i + 1)).map(|x| x.to_string().into_bytes()); + let elems = (elem_count * i..elem_count * (i + 1)) + .map(|x| element_hash(x.to_string().as_bytes())); storage.add(elems, elem_count).await.unwrap(); } for _ in 0..dup_batch { - storage.add(Some(vec![]), 1).await.unwrap(); + storage.add(Some(element_hash(&[])), 1).await.unwrap(); } // Flush happens. @@ -356,7 +357,7 @@ mod tests { let batch = 1000; for _ in 0..batch { - storage.add(Some(vec![]), 1).await.unwrap(); + storage.add(Some(element_hash(&[])), 1).await.unwrap(); } // Drain the storage. diff --git a/src/index/src/bloom_filter/reader.rs b/src/index/src/bloom_filter/reader.rs index 04eae84f005..15cf957b7b6 100644 --- a/src/index/src/bloom_filter/reader.rs +++ b/src/index/src/bloom_filter/reader.rs @@ -25,10 +25,10 @@ use greptime_proto::v1::index::{BloomFilterLoc, BloomFilterMeta}; use prost::Message; use snafu::{ResultExt, ensure}; -use crate::bloom_filter::SEED; use crate::bloom_filter::error::{ DecodeProtoSnafu, FileSizeTooSmallSnafu, IoSnafu, Result, UnexpectedMetaSizeSnafu, }; +use crate::bloom_filter::{PrehashedBloomFilter, PrehashedBuildHasher, SEED}; /// Minimum size of the bloom filter, which is the size of the length of the bloom filter. const BLOOM_META_LEN_SIZE: u64 = 4; @@ -181,11 +181,12 @@ pub trait BloomFilterReader: Sync { Ok(bm) } + /// Reads multiple bloom filters; probe them with [`crate::bloom_filter::element_hash`]. async fn bloom_filter_vec( &self, locs: &[BloomFilterLoc], metrics: Option<&mut BloomFilterReadMetrics>, - ) -> Result> { + ) -> Result> { let ranges = locs .iter() .map(|l| l.offset..l.offset + l.size) @@ -196,7 +197,7 @@ pub trait BloomFilterReader: Sync { for (bs, loc) in bss.into_iter().zip(locs.iter()) { let vec = bytes_to_u64_vec(&bs); let bm = BloomFilter::from_vec(vec) - .seed(&SEED) + .hasher(PrehashedBuildHasher::default()) .expected_items(loc.element_count as _); result.push(bm); } diff --git a/src/index/src/fulltext_index/create/bloom_filter.rs b/src/index/src/fulltext_index/create/bloom_filter.rs index 28648f9b540..a2db31418eb 100644 --- a/src/index/src/fulltext_index/create/bloom_filter.rs +++ b/src/index/src/fulltext_index/create/bloom_filter.rs @@ -74,11 +74,12 @@ impl BloomFilterFulltextIndexCreator { #[async_trait] impl FulltextIndexCreator for BloomFilterFulltextIndexCreator { async fn push_text(&mut self, text: &str) -> Result<()> { - let tokens = self.analyzer.analyze_text(text)?; + let mut token_buf = Vec::new(); + let hashes = self.analyzer.analyze_text_hashes(text, &mut token_buf); self.inner .as_mut() .context(AbortedSnafu)? - .push_row_elems(tokens) + .push_row_hashes(hashes) .await .map_err(BoxedError::new) .context(ExternalSnafu)?; diff --git a/src/index/src/fulltext_index/tokenizer.rs b/src/index/src/fulltext_index/tokenizer.rs index 3afc826e6f1..e6752a40420 100644 --- a/src/index/src/fulltext_index/tokenizer.rs +++ b/src/index/src/fulltext_index/tokenizer.rs @@ -13,6 +13,7 @@ // limitations under the License. use crate::Bytes; +use crate::bloom_filter::element_hash; use crate::fulltext_index::error::Result; lazy_static::lazy_static! { @@ -140,6 +141,30 @@ impl Analyzer { } } + /// Returns the bloom filter hash of each token in the given text. + /// + /// Equivalent to hashing every token returned by [`Analyzer::analyze_text`] with + /// [`element_hash`]. Only case-insensitive non-ASCII tokens allocate, in `to_lowercase`; + /// case-insensitive ASCII tokens are lowercased in `buf`. + pub fn analyze_text_hashes<'a>( + &self, + text: &'a str, + buf: &'a mut Vec, + ) -> impl Iterator + use<'a> { + let case_sensitive = self.case_sensitive; + self.tokenizer.tokenize(text).into_iter().map(move |token| { + if case_sensitive { + element_hash(token.as_bytes()) + } else if token.is_ascii() { + buf.clear(); + buf.extend(token.bytes().map(|b| b.to_ascii_lowercase())); + element_hash(buf) + } else { + element_hash(token.to_lowercase().as_bytes()) + } + }) + } + /// Analyzes the given text into a list of tokens. pub fn analyze_text(&self, text: &str) -> Result> { let res = self @@ -162,6 +187,28 @@ impl Analyzer { mod tests { use super::*; + #[test] + fn test_analyze_text_hashes_matches_analyze_text() { + let text = "Hello, WORLD ship_Ship 清æ´è¡¨é¢ ÄÖÜ straße Ä°stanbul x"; + for (tokenizer, case_sensitive) in [ + (Box::new(EnglishTokenizer) as Box, false), + (Box::new(EnglishTokenizer), true), + (Box::new(ChineseTokenizer), false), + ] { + let analyzer = Analyzer::new(tokenizer, case_sensitive); + let expected = analyzer + .analyze_text(text) + .unwrap() + .iter() + .map(|t| element_hash(t)) + .collect::>(); + let hashes = analyzer + .analyze_text_hashes(text, &mut Vec::new()) + .collect::>(); + assert_eq!(expected, hashes); + } + } + #[test] fn test_english_tokenizer() { let tokenizer = EnglishTokenizer; diff --git a/src/index/src/inverted_index/create.rs b/src/index/src/inverted_index/create.rs index dade9705ac8..74da3612232 100644 --- a/src/index/src/inverted_index/create.rs +++ b/src/index/src/inverted_index/create.rs @@ -27,30 +27,24 @@ use crate::inverted_index::format::writer::InvertedIndexWriter; pub trait InvertedIndexCreator: Send { /// Adds a value to the named index. A `None` value represents an absence of data (null) /// - /// - `index_name`: Identifier for the index being built - /// - `value`: The data to be indexed, or `None` for a null entry - /// /// It should be equivalent to calling `push_with_name_n` with `n = 1` - async fn push_with_name( - &mut self, - index_name: &str, - value: Option>, - ) -> Result<()> { - self.push_with_name_n(index_name, value, 1).await + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_with_name(&mut self, index_name: &str, value: Option>) -> bool { + self.push_with_name_n(index_name, value, 1) } - /// Adds `n` identical values to the named index. `None` values represent absence of data (null) + /// Buffers `n` identical values for the named index. `None` values represent absence of + /// data (null). /// - /// - `index_name`: Identifier for the index being built - /// - `value`: The data to be indexed, or `None` for a null entry - /// - /// It should be equivalent to calling `push_with_name` `n` times - async fn push_with_name_n( - &mut self, - index_name: &str, - value: Option>, - n: usize, - ) -> Result<()>; + /// Returns true when buffered data exceeds the memory limit; the caller must then call + /// [`InvertedIndexCreator::spill`] before pushing more. Pushing is synchronous so the + /// per-row path does not allocate a future. + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_with_name_n(&mut self, index_name: &str, value: Option>, n: usize) + -> bool; + + /// Moves the buffers that asked for it to external storage. + async fn spill(&mut self) -> Result<()>; /// Finalizes the index creation process, ensuring all data is properly indexed and stored /// in the provided writer diff --git a/src/index/src/inverted_index/create/sort.rs b/src/index/src/inverted_index/create/sort.rs index bdc1ec21f62..fbf81adbaa4 100644 --- a/src/index/src/inverted_index/create/sort.rs +++ b/src/index/src/inverted_index/create/sort.rs @@ -42,15 +42,15 @@ pub struct SortOutput { /// Handles data sorting, supporting incremental input and retrieval of sorted output #[async_trait] pub trait Sorter: Send { - /// Inputs a non-null or null value into the sorter. - /// Should be equivalent to calling `push_n` with n = 1 - async fn push(&mut self, value: Option>) -> Result<()> { - self.push_n(value, 1).await - } + /// Buffers `n` identical non-null or null values in memory. + /// + /// Returns true when the buffer should be spilled with [`Sorter::spill`] before more + /// values are pushed. Kept synchronous so the per-row path does not allocate a future. + #[must_use = "a true result requires calling `spill` before pushing more"] + fn push_n(&mut self, value: Option>, n: usize) -> bool; - /// Pushing n identical non-null or null values into the sorter. - /// Should be equivalent to calling `push` n times - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()>; + /// Moves the in-memory buffer to external storage. + async fn spill(&mut self) -> Result<()>; /// Completes the sorting process and returns the sorted data async fn output(&mut self) -> Result; diff --git a/src/index/src/inverted_index/create/sort/external_sort.rs b/src/index/src/inverted_index/create/sort/external_sort.rs index 12c49ca89de..cef07df9f8f 100644 --- a/src/index/src/inverted_index/create/sort/external_sort.rs +++ b/src/index/src/inverted_index/create/sort/external_sort.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{BTreeMap, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::mem; use std::num::NonZeroUsize; use std::ops::RangeInclusive; @@ -22,6 +22,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use async_trait::async_trait; use common_telemetry::{debug, error}; use futures::stream; +use roaring::RoaringBitmap; use snafu::ResultExt; use crate::bitmap::Bitmap; @@ -35,6 +36,28 @@ use crate::inverted_index::create::sort_create::SorterFactory; use crate::inverted_index::error::{IntermediateSnafu, Result}; use crate::{Bytes, BytesRef}; +/// Segments of one value. Rows arrive in order, so segment ids only grow. +struct Posting { + segments: RoaringBitmap, + last_segment: u32, +} + +/// Estimated cost of one segment id in a roaring array container. +const SEGMENT_SIZE: usize = size_of::(); + +impl Posting { + /// Adds `start..=end` and returns the estimated memory growth. + fn push(&mut self, start: u32, end: u32) -> usize { + if end <= self.last_segment { + return 0; + } + let start = start.max(self.last_segment + 1); + self.segments.insert_range(start..=end); + self.last_segment = end; + (end - start + 1) as usize * SEGMENT_SIZE + } +} + /// `ExternalSorter` manages the sorting of data using both in-memory structures and external files. /// It dumps data to external files when the in-memory buffer crosses a certain memory threshold. pub struct ExternalSorter { @@ -48,7 +71,7 @@ pub struct ExternalSorter { segment_null_bitmap: Bitmap, /// In-memory buffer to hold values and their corresponding bitmaps until memory threshold is exceeded - values_buffer: BTreeMap, + values_buffer: HashMap, /// Count of all rows ingested so far total_row_count: usize, @@ -78,11 +101,11 @@ pub struct ExternalSorter { #[async_trait] impl Sorter for ExternalSorter { - /// Pushes n identical values into the sorter, adding them to the in-memory buffer and dumping - /// the buffer to an external file if necessary - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()> { + /// Pushes n identical values into the in-memory buffer; returns whether it should be + /// spilled. + fn push_n(&mut self, value: Option>, n: usize) -> bool { if n == 0 { - return Ok(()); + return false; } let segment_index_range = self.segment_index_range(n); @@ -90,13 +113,17 @@ impl Sorter for ExternalSorter { if let Some(value) = value { let memory_diff = self.push_not_null(value, segment_index_range); - self.may_dump_buffer(memory_diff).await + self.account_memory(memory_diff) } else { self.segment_null_bitmap.insert_range(segment_index_range); - Ok(()) + false } } + async fn spill(&mut self) -> Result<()> { + self.dump_buffer().await + } + /// Finalizes the sorting operation, merging data from both in-memory buffer and external files /// into a sorted stream async fn output(&mut self) -> Result { @@ -110,9 +137,7 @@ impl Sorter for ExternalSorter { let mut tree_nodes: VecDeque = VecDeque::with_capacity(readers.len() + 1); tree_nodes.push_back(Box::new(stream::iter( - mem::take(&mut self.values_buffer) - .into_iter() - .map(|(value, (bitmap, _))| Ok((value, bitmap))), + Self::sorted(mem::take(&mut self.values_buffer)).map(Ok), ))); for (_, reader) in readers { tree_nodes.push_back(IntermediateReader::new(reader).into_stream().await?); @@ -149,7 +174,7 @@ impl ExternalSorter { temp_file_provider, segment_null_bitmap: Bitmap::new_bitvec(), // bitvec is more efficient for many null values - values_buffer: BTreeMap::new(), + values_buffer: HashMap::default(), total_row_count: 0, segment_row_count, @@ -188,51 +213,65 @@ impl ExternalSorter { value: BytesRef<'_>, segment_index_range: RangeInclusive, ) -> usize { + let (start, end) = ( + *segment_index_range.start() as u32, + *segment_index_range.end() as u32, + ); match self.values_buffer.get_mut(value) { - Some((bitmap, mem_usage)) => { - bitmap.insert_range(segment_index_range); - let new_usage = bitmap.memory_usage() + value.len(); - let diff = new_usage - *mem_usage; - *mem_usage = new_usage; - - diff - } + Some(posting) => posting.push(start, end), None => { - let mut bitmap = Bitmap::new_roaring(); - bitmap.insert_range(segment_index_range); - - let mem_usage = bitmap.memory_usage() + value.len(); - self.values_buffer - .insert(value.to_vec(), (bitmap, mem_usage)); - - mem_usage + let mut segments = RoaringBitmap::new(); + segments.insert_range(start..=end); + let posting = Posting { + segments, + last_segment: end, + }; + self.values_buffer.insert(value.to_vec(), posting); + value.len() + (end - start + 1) as usize * SEGMENT_SIZE } } } - /// Checks if the in-memory buffer exceeds the threshold and offloads it to external storage if necessary - async fn may_dump_buffer(&mut self, memory_diff: usize) -> Result<()> { + /// Drains `values` sorted by value. + fn sorted( + values: HashMap, + ) -> impl Iterator { + let mut values = values.into_iter().collect::>(); + values.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + values + .into_iter() + .map(|(value, posting)| (value, Bitmap::Roaring(posting.segments))) + } + + /// Records `memory_diff` and returns whether the buffer exceeds the thresholds and + /// should be offloaded to external storage. + fn account_memory(&mut self, memory_diff: usize) -> bool { self.current_memory_usage += memory_diff; let memory_usage = self.current_memory_usage; self.global_memory_usage .fetch_add(memory_diff, Ordering::Relaxed); - if self.global_memory_usage_sort_limit.is_none() { - return Ok(()); + let Some(limit) = self.global_memory_usage_sort_limit else { + return false; + }; + if self.global_memory_usage.load(Ordering::Relaxed) < limit { + return false; } - - if self.global_memory_usage.load(Ordering::Relaxed) - < self.global_memory_usage_sort_limit.unwrap() - { - return Ok(()); - } - if let Some(current_threshold) = self.current_memory_usage_threshold && memory_usage < current_threshold { + return false; + } + true + } + + async fn dump_buffer(&mut self) -> Result<()> { + // A second spill request before any new value would write an empty file with the + // same id and replace the first one. + if self.values_buffer.is_empty() { return Ok(()); } - + let memory_usage = self.current_memory_usage; let file_id = &format!("{:012}", self.total_row_count); let index_name = &self.index_name; let writer = self @@ -247,7 +286,7 @@ impl ExternalSorter { self.current_memory_usage = 0; let entries = values.len(); - IntermediateWriter::new(writer).write_all(values.into_iter().map(|(k, (b, _))| (k, b))).await.inspect(|_| + IntermediateWriter::new(writer).write_all(Self::sorted(values)).await.inspect(|_| debug!("Dumped {entries} entries ({memory_usage} bytes) to intermediate file {file_id} for index {index_name}") ).inspect_err(|e| error!(e; "Failed to dump {entries} entries to intermediate file {file_id} for index {index_name}") @@ -271,7 +310,7 @@ impl ExternalSorter { #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{BTreeMap, HashMap}; use std::iter; use std::sync::Mutex; @@ -329,7 +368,9 @@ mod tests { dictionary_values_and_sorted_result(row_count, segment_row_count); for (value, n) in dic_values { - sorter.push_n(value.as_deref(), n).await.unwrap(); + if sorter.push_n(value.as_deref(), n) { + sorter.spill().await.unwrap(); + } } sorted_result @@ -338,7 +379,9 @@ mod tests { shuffle_values_and_sorted_result(row_count, segment_row_count); for value in mock_values { - sorter.push(value.as_deref()).await.unwrap(); + if sorter.push_n(value.as_deref(), 1) { + sorter.spill().await.unwrap(); + } } sorted_result diff --git a/src/index/src/inverted_index/create/sort_create.rs b/src/index/src/inverted_index/create/sort_create.rs index cb22179a698..9218dd76ab6 100644 --- a/src/index/src/inverted_index/create/sort_create.rs +++ b/src/index/src/inverted_index/create/sort_create.rs @@ -42,6 +42,9 @@ pub struct SortIndexCreator { /// Number of rows in each segment, used to produce sorters segment_row_count: NonZeroUsize, + + /// Indexes whose sorter asked to spill its buffer. + pending_spills: Vec, } #[async_trait] @@ -50,22 +53,33 @@ impl InvertedIndexCreator for SortIndexCreator { /// /// If the index does not exist, a new index is created even if `n` is 0. /// Caller may leverage this behavior to create indexes with no data. - async fn push_with_name_n( + fn push_with_name_n( &mut self, index_name: &str, value: Option>, n: usize, - ) -> Result<()> { - match self.sorters.get_mut(index_name) { - Some(sorter) => sorter.push_n(value, n).await, + ) -> bool { + let sorter = match self.sorters.get_mut(index_name) { + Some(sorter) => sorter, None => { - let index_name = index_name.to_string(); - let mut sorter = (self.sorter_factory)(index_name.clone(), self.segment_row_count); - sorter.push_n(value, n).await?; - self.sorters.insert(index_name, sorter); - Ok(()) + let sorter = (self.sorter_factory)(index_name.to_string(), self.segment_row_count); + self.sorters.entry(index_name.to_string()).or_insert(sorter) + } + }; + if sorter.push_n(value, n) { + self.pending_spills.push(index_name.to_string()); + return true; + } + false + } + + async fn spill(&mut self) -> Result<()> { + for index_name in std::mem::take(&mut self.pending_spills) { + if let Some(sorter) = self.sorters.get_mut(&index_name) { + sorter.spill().await?; } } + Ok(()) } /// Finalizes the sorting for all indexes and writes them using the inverted index writer @@ -110,6 +124,7 @@ impl SortIndexCreator { sorter_factory, sorters: HashMap::new(), segment_row_count, + pending_spills: Vec::new(), } } } @@ -117,6 +132,7 @@ impl SortIndexCreator { #[cfg(test)] mod tests { use std::collections::BTreeMap; + use std::sync::{Arc, Mutex}; use common_base::BitVec; use futures::{StreamExt, stream}; @@ -140,10 +156,7 @@ mod tests { for (index_name, values) in index_values { for value in values { - creator - .push_with_name(index_name, Some(value)) - .await - .unwrap(); + assert!(!creator.push_with_name(index_name, Some(value))); } } @@ -189,10 +202,7 @@ mod tests { for (index_name, values) in index_values { for value in values { - creator - .push_with_name(index_name, Some(value)) - .await - .unwrap(); + assert!(!creator.push_with_name(index_name, Some(value))); } } @@ -221,9 +231,9 @@ mod tests { let mut creator = SortIndexCreator::new(NaiveSorter::factory(), NonZeroUsize::new(1).unwrap()); - creator.push_with_name_n("a", None, 0).await.unwrap(); - creator.push_with_name_n("b", None, 0).await.unwrap(); - creator.push_with_name_n("c", None, 0).await.unwrap(); + assert!(!creator.push_with_name_n("a", None, 0)); + assert!(!creator.push_with_name_n("b", None, 0)); + assert!(!creator.push_with_name_n("c", None, 0)); let mut mock_writer = MockInvertedIndexWriter::new(); mock_writer @@ -250,6 +260,51 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_sort_index_creator_spills_requested_sorters() { + let spilled = Arc::new(Mutex::new(Vec::new())); + let factory: SorterFactory = { + let spilled = spilled.clone(); + Box::new(move |index_name, _| { + Box::new(SpillRecordingSorter { + index_name, + spilled: spilled.clone(), + }) + }) + }; + let mut creator = SortIndexCreator::new(factory, NonZeroUsize::new(1).unwrap()); + + assert!(creator.push_with_name("a", Some(b"1"))); + assert!(!creator.push_with_name("b", Some(b"1"))); + creator.spill().await.unwrap(); + assert_eq!(*spilled.lock().unwrap(), vec!["a"]); + + creator.spill().await.unwrap(); + assert_eq!(*spilled.lock().unwrap(), vec!["a"]); + } + + /// Requests a spill on every push to index `a` and records the spills it receives. + struct SpillRecordingSorter { + index_name: String, + spilled: Arc>>, + } + + #[async_trait] + impl Sorter for SpillRecordingSorter { + fn push_n(&mut self, _value: Option>, _n: usize) -> bool { + self.index_name == "a" + } + + async fn spill(&mut self) -> Result<()> { + self.spilled.lock().unwrap().push(self.index_name.clone()); + Ok(()) + } + + async fn output(&mut self) -> Result { + unreachable!() + } + } + fn set_bit(bit_vec: &mut BitVec, index: usize) { if index >= bit_vec.len() { bit_vec.resize(index + 1, false); @@ -277,20 +332,18 @@ mod tests { #[async_trait] impl Sorter for NaiveSorter { - async fn push(&mut self, value: Option>) -> Result<()> { - let segment_index = self.total_row_count / self.segment_row_count; - self.total_row_count += 1; + fn push_n(&mut self, value: Option>, n: usize) -> bool { + for _ in 0..n { + let segment_index = self.total_row_count / self.segment_row_count; + self.total_row_count += 1; - let bitmap = self.values.entry(value.map(Into::into)).or_default(); - set_bit(bitmap, segment_index); - - Ok(()) + let bitmap = self.values.entry(value.map(Into::into)).or_default(); + set_bit(bitmap, segment_index); + } + false } - async fn push_n(&mut self, value: Option>, n: usize) -> Result<()> { - for _ in 0..n { - self.push(value).await?; - } + async fn spill(&mut self) -> Result<()> { Ok(()) } diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 9ea5258ee3b..d1dee226515 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -204,10 +204,12 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(EncodeSnafu)?; - self.index_creator - .push_with_name_n(target_key, elem, count) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name_n(target_key, elem, count) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { if self.codec.pk_col_info(*col_id).is_some() { @@ -233,10 +235,15 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(DecodeSnafu)?; - self.index_creator + if self + .index_creator .push_with_name_n(target_key, value, count) - .await - .context(PushIndexValueSnafu)?; + { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } } @@ -307,10 +314,12 @@ impl InvertedIndexer { }) .transpose()?; - self.index_creator - .push_with_name_n(target_key, value, n) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name_n(target_key, value, n) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } // fields None => { @@ -326,10 +335,12 @@ impl InvertedIndexer { self.value_buf.clear(); let value = values.data.get_ref(i); if value.is_null() { - self.index_creator - .push_with_name(target_key, None) - .await - .context(PushIndexValueSnafu)?; + if self.index_creator.push_with_name(target_key, None) { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } else { IndexValueCodec::encode_nonnull_value( value, @@ -337,10 +348,15 @@ impl InvertedIndexer { &mut self.value_buf, ) .context(EncodeSnafu)?; - self.index_creator + if self + .index_creator .push_with_name(target_key, Some(&self.value_buf)) - .await - .context(PushIndexValueSnafu)?; + { + self.index_creator + .spill() + .await + .context(PushIndexValueSnafu)?; + } } } } From 75bd8e9ce6fa554c2da3fa7d72dc63c8b50e99ce Mon Sep 17 00:00:00 2001 From: Minghan Jiang Date: Mon, 28 Sep 2026 08:39:48 +0000 Subject: [PATCH 056/133] feat: add HDFS object storage backend (#8701) * feat: add HDFS object storage backend Signed-off-by: Minghan2005 * fix: make HDFS storage operations durable Gate the native HDFS backend behind an explicit feature. Publish writes through same-directory temporary files and atomic HDFS Rename2 replacement, and provide streaming copy fallback for COPY_REGION. Add regression coverage for interrupted writes and the region-copy path. Signed-off-by: Minghan2005 * ci: run HDFS object store tests Signed-off-by: jeremyhi * docs: note HDFS temporary file cleanup follow-up Signed-off-by: jeremyhi * feat: enable HDFS object storage by default Signed-off-by: jeremyhi * docs: remove redundant HDFS build feature notes Signed-off-by: jeremyhi --------- Signed-off-by: Minghan2005 Signed-off-by: jeremyhi Co-authored-by: Minghan2005 Co-authored-by: jeremyhi --- .github/workflows/nightly-ci.yml | 2 +- .github/workflows/rust.yml | 4 +- Cargo.lock | 241 ++++++++++-- Cargo.toml | 1 + config/config.md | 12 +- config/datanode.example.toml | 22 +- config/standalone.example.toml | 22 +- src/cmd/Cargo.toml | 2 + src/cmd/tests/load_config_test.rs | 38 ++ src/mito2/Cargo.toml | 1 + src/mito2/src/region/utils.rs | 63 ++++ src/object-store/Cargo.toml | 2 + src/object-store/src/config.rs | 109 ++++++ src/object-store/src/factory.rs | 61 +++ src/object-store/src/layers.rs | 4 + src/object-store/src/layers/hdfs.rs | 551 ++++++++++++++++++++++++++++ src/object-store/src/lib.rs | 2 + 17 files changed, 1090 insertions(+), 47 deletions(-) create mode 100644 src/object-store/src/layers/hdfs.rs diff --git a/.github/workflows/nightly-ci.yml b/.github/workflows/nightly-ci.yml index a948dbcc20f..f7f955233a1 100644 --- a/.github/workflows/nightly-ci.yml +++ b/.github/workflows/nightly-ci.yml @@ -102,7 +102,7 @@ jobs: working-directory: tests-integration/fixtures run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index d4ab7e1d7b3..e29b69678c1 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -233,7 +233,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 @@ -290,7 +290,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions + run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F ai_functions -F hdfs-object-store env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/Cargo.lock b/Cargo.lock index bb7f69d1551..4ced42fc860 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,10 +40,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] +[[package]] +name = "aes" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "aes-siv" version = "0.7.0" @@ -51,10 +62,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e08d0cdb774acd1e4dac11478b1a0c0d203134b2aab0ba25eb430de9b18f8b9" dependencies = [ "aead", - "aes", - "cipher", + "aes 0.8.4", + "cipher 0.4.4", "cmac", - "ctr", + "ctr 0.9.2", "dbl", "digest 0.10.7", "zeroize", @@ -1400,6 +1411,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-padding" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "710f1dd022ef4e93f8a438b4ba958de7f64308434fa6a87104481645cc30068b" +dependencies = [ + "hybrid-array", +] + [[package]] name = "blocking" version = "1.6.1" @@ -1534,9 +1554,9 @@ checksum = "40e38929add23cdf8a366df9b0e088953150724bcbe5fc330b0d8eb3b328eec8" [[package]] name = "bumpalo" -version = "3.19.0" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytecheck" @@ -1730,7 +1750,16 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "cbc" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2dc9ee5f88d11e0beb842c88b33c8a5cf0d1329c4b19494af42b07dbfe8896" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -1775,7 +1804,7 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "738b8d467867f80a71351933f70461f5b56f24d5c93e0cf216e59229c968d330" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -1827,7 +1856,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] @@ -1850,7 +1879,7 @@ checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" dependencies = [ "aead", "chacha20 0.9.1", - "cipher", + "cipher 0.4.4", "poly1305", "zeroize", ] @@ -1943,10 +1972,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" dependencies = [ "crypto-common 0.1.6", - "inout", + "inout 0.1.4", "zeroize", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.0", + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1955,7 +1995,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" dependencies = [ "glob", "libc", - "libloading", + "libloading 0.8.8", ] [[package]] @@ -2118,7 +2158,7 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8543454e3c3f5126effff9cd44d562af4e31fb8ce1cc0d3dcd8f084515dbc1aa" dependencies = [ - "cipher", + "cipher 0.4.4", "dbl", "digest 0.10.7", ] @@ -2228,7 +2268,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -3218,6 +3258,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -3238,9 +3284,9 @@ dependencies = [ [[package]] name = "crc" -version = "3.3.0" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9710d3b3739c2e349eb44fe848ad0b7c8cb1e42bd87ee49371df2f7acaf3e675" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" dependencies = [ "crc-catalog", ] @@ -3438,7 +3484,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" dependencies = [ "aead", - "cipher", + "cipher 0.4.4", "generic-array", "poly1305", "salsa20", @@ -3483,7 +3529,16 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "ctr" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaca1c4b237092596f64d571e9db6ce4109c4ef9742e27590f1709594461f21" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -4742,6 +4797,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "des" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "916a94e407b54f9034d71dd748234cd1e516ced6284009906ae246f177eafe5a" +dependencies = [ + "cipher 0.5.2", +] + [[package]] name = "diff" version = "0.1.13" @@ -5819,6 +5883,34 @@ dependencies = [ "byteorder", ] +[[package]] +name = "g2gen" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5a7e0eb46f83a20260b850117d204366674e85d3a908d90865c78df9a6b1dfc" +dependencies = [ + "g2poly", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "g2p" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "539e2644c030d3bf4cd208cb842d2ce2f80e82e6e8472390bcef83ceba0d80ad" +dependencies = [ + "g2gen", + "g2poly", +] + +[[package]] +name = "g2poly" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "312d2295c7302019c395cfb90dacd00a82a2eabd700429bba9c7a3f38dbbe11b" + [[package]] name = "generator" version = "0.8.5" @@ -6167,6 +6259,47 @@ dependencies = [ "hashbrown 0.15.4", ] +[[package]] +name = "hdfs-native" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd181084003308224efddf737417832186839ce2882d2468ddad0114cfb3551" +dependencies = [ + "aes 0.9.3", + "base64 0.22.1", + "bitflags 2.12.1", + "bumpalo", + "bytes", + "cbc 0.2.1", + "chrono", + "cipher 0.5.2", + "crc", + "ctr 0.10.1", + "des", + "dns-lookup", + "futures", + "g2p", + "hex", + "hmac 0.13.0", + "libc", + "libloading 0.9.0", + "log", + "md-5 0.11.0", + "num-traits", + "once_cell", + "prost 0.14.1", + "prost-types 0.14.1", + "rand 0.10.1", + "regex", + "roxmltree", + "socket2 0.6.4", + "thiserror 2.0.17", + "tokio", + "url", + "uuid", + "whoami 2.1.3", +] + [[package]] name = "hdrhistogram" version = "7.5.4" @@ -6549,7 +6682,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -6620,7 +6753,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -6996,10 +7129,20 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" dependencies = [ - "block-padding", + "block-padding 0.3.3", "generic-array", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "block-padding 0.4.2", + "hybrid-array", +] + [[package]] name = "instant" version = "0.1.13" @@ -7061,7 +7204,7 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96e4f67dbfc0f75d7b65953ecf0be3fd84ee0cb1ae72a00a4aa9a2f5518a2c80" dependencies = [ - "aes", + "aes 0.8.4", ] [[package]] @@ -7810,6 +7953,16 @@ dependencies = [ "windows-targets 0.53.5", ] +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link 0.2.1", +] + [[package]] name = "liblzma" version = "0.4.6" @@ -9300,7 +9453,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 1.3.1", + "proc-macro-crate 3.3.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9407,6 +9560,7 @@ dependencies = [ "common-test-util", "derive_builder 0.20.2", "futures", + "hdfs-native", "humantime-serde", "object-store", "object_store", @@ -9482,7 +9636,7 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cc40678e045ff4eb1666ea6c0f994b133c31f673c09aed292261b6d5b6963a0" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -9555,6 +9709,7 @@ dependencies = [ "opendal-service-azblob", "opendal-service-fs", "opendal-service-gcs", + "opendal-service-hdfs-native", "opendal-service-http", "opendal-service-mysql", "opendal-service-oss", @@ -9745,6 +9900,20 @@ dependencies = [ "uuid", ] +[[package]] +name = "opendal-service-hdfs-native" +version = "0.59.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b84340e11a3bcd7f1f0b55df032d040b2907483c546c3021485aba749c3072c" +dependencies = [ + "bytes", + "futures", + "hdfs-native", + "log", + "opendal-core", + "serde", +] + [[package]] name = "opendal-service-http" version = "0.59.2" @@ -10792,8 +10961,8 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" dependencies = [ - "aes", - "cbc", + "aes 0.8.4", + "cbc 0.1.2", "der", "pbkdf2", "scrypt", @@ -11313,7 +11482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11333,8 +11502,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.10.5", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11382,7 +11551,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11395,7 +11564,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -12988,7 +13157,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -13778,7 +13947,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -16119,13 +16288,13 @@ version = "0.33.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "772206835b5bca8719825d2ab9e6fb160cbfce8ac1dc0e5dffe22cbe77254468" dependencies = [ - "aes", + "aes 0.8.4", "aes-siv", "base16", "base62", "base64-simd", "bytes", - "cbc", + "cbc 0.1.2", "cfb-mode", "cfg-if", "chacha20poly1305", @@ -16141,7 +16310,7 @@ dependencies = [ "crc", "crypto_secretbox", "csv", - "ctr", + "ctr 0.9.2", "digest 0.10.7", "dns-lookup", "domain", diff --git a/Cargo.toml b/Cargo.toml index e057db65826..5f497a3faff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -161,6 +161,7 @@ fst = "0.4.7" futures = "0.3" futures-util = "0.3" greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "9533078d6fd4bdaddbca0750cb9bf7ad64b2e568" } +hdfs-native = "0.14.6" hex = "0.4" hostname = "0.4.0" http = "1" diff --git a/config/config.md b/config/config.md index e380f075548..cebc5091b4c 100644 --- a/config/config.md +++ b/config/config.md @@ -148,9 +148,11 @@ | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | | `storage.copy_root` | String | `./greptimedb_data/copy` | Root directory for standalone SQL access to local files.
Relative SQL paths are resolved below this directory. Absolute paths are accepted only when
they are inside this directory. Defaults to `/copy`.
Distributed deployments always reject SQL access to local files.
Upgrade note: COPY commands and existing external tables that reference paths outside this
directory will fail. Move those files below the copy root, set this option to a dedicated
directory containing them, or migrate the files to object storage before upgrading. | -| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS. | +| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | -| `storage.root` | String | Unset | The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`.
**It's only used when the storage type is `S3`, `Oss` and `Azblob`**. | +| `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | +| `storage.name_node` | String | Unset | The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`.
**It's only used when the storage type is `Hdfs`**. | +| `storage.options` | InlineTable | Unset | Additional options passed to the native HDFS client.
**It's only used when the storage type is `Hdfs`**. | | `storage.access_key_id` | String | Unset | The access key id of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3` and `Oss`**. | | `storage.secret_access_key` | String | Unset | The secret access key of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3`**. | | `storage.access_key_secret` | String | Unset | The secret access key of the aliyun account.
**It's only used when the storage type is `Oss`**. | @@ -609,9 +611,11 @@ | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | -| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS. | +| `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS.
- `Hdfs`: the data is stored in the Hadoop Distributed File System. | | `storage.bucket` | String | Unset | The S3 bucket name.
**It's only used when the storage type is `S3`, `Oss` and `Gcs`**. | -| `storage.root` | String | Unset | The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`.
**It's only used when the storage type is `S3`, `Oss` and `Azblob`**. | +| `storage.root` | String | Unset | The directory or object prefix under which data is stored.
**It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. | +| `storage.name_node` | String | Unset | The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`.
**It's only used when the storage type is `Hdfs`**. | +| `storage.options` | InlineTable | Unset | Additional options passed to the native HDFS client.
**It's only used when the storage type is `Hdfs`**. | | `storage.access_key_id` | String | Unset | The access key id of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3` and `Oss`**. | | `storage.secret_access_key` | String | Unset | The secret access key of the aws account.
It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key.
**It's only used when the storage type is `S3`**. | | `storage.access_key_secret` | String | Unset | The secret access key of the aliyun account.
**It's only used when the storage type is `Oss`**. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index da2d353fef8..253cc3f2a57 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -300,6 +300,13 @@ overwrite_entry_start_id = false # credential = "base64-credential" # endpoint = "https://storage.googleapis.com" +# Example of using HDFS as the storage with the native Rust client. +# [storage] +# type = "Hdfs" +# root = "/greptimedb" +# name_node = "hdfs://127.0.0.1:9000" +# options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + ## The query engine options. [query] ## Parallelism of the query engine. @@ -348,6 +355,7 @@ data_home = "./greptimedb_data" ## - `Gcs`: the data is stored in the Google Cloud Storage. ## - `Azblob`: the data is stored in the Azure Blob Storage. ## - `Oss`: the data is stored in the Aliyun OSS. +## - `Hdfs`: the data is stored in the Hadoop Distributed File System. type = "File" ## The S3 bucket name. @@ -355,11 +363,21 @@ type = "File" ## @toml2docs:none-default bucket = "greptimedb" -## The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`. -## **It's only used when the storage type is `S3`, `Oss` and `Azblob`**. +## The directory or object prefix under which data is stored. +## **It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. ## @toml2docs:none-default root = "greptimedb" +## The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +name_node = "hdfs://127.0.0.1:9000" + +## Additional options passed to the native HDFS client. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +options = {} + ## The access key id of the aws account. ## It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key. ## **It's only used when the storage type is `S3` and `Oss`**. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e0a8deb9c22..fad4f64f08a 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -509,6 +509,13 @@ max_running_procedures = 128 # credential = "base64-credential" # endpoint = "https://storage.googleapis.com" +# Example of using HDFS as the storage with the native Rust client. +# [storage] +# type = "Hdfs" +# root = "/greptimedb" +# name_node = "hdfs://127.0.0.1:9000" +# options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + ## The query engine options. [query] ## Parallelism of the query engine. @@ -566,6 +573,7 @@ data_home = "./greptimedb_data" ## - `Gcs`: the data is stored in the Google Cloud Storage. ## - `Azblob`: the data is stored in the Azure Blob Storage. ## - `Oss`: the data is stored in the Aliyun OSS. +## - `Hdfs`: the data is stored in the Hadoop Distributed File System. type = "File" ## The S3 bucket name. @@ -573,11 +581,21 @@ type = "File" ## @toml2docs:none-default bucket = "greptimedb" -## The S3 data will be stored in the specified prefix, for example, `s3://${bucket}/${root}`. -## **It's only used when the storage type is `S3`, `Oss` and `Azblob`**. +## The directory or object prefix under which data is stored. +## **It's only used when the storage type is `S3`, `Oss`, `Gcs`, `Azblob` and `Hdfs`**. ## @toml2docs:none-default root = "greptimedb" +## The HDFS NameNode URI, for example, `hdfs://127.0.0.1:9000`. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +name_node = "hdfs://127.0.0.1:9000" + +## Additional options passed to the native HDFS client. +## **It's only used when the storage type is `Hdfs`**. +## @toml2docs:none-default +options = {} + ## The access key id of the aws account. ## It's **highly recommended** to use AWS IAM roles instead of hardcoding the access key id and secret key. ## **It's only used when the storage type is `S3` and `Oss`**. diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index a67dd895999..2c925de57c8 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -22,6 +22,7 @@ required-features = ["dev-tools"] [features] default = [ "ai_functions", + "hdfs-object-store", "servers/pprof", "servers/mem-prof", "meta-srv/pg_kvbackend", @@ -32,6 +33,7 @@ enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterpr # Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] +hdfs-object-store = ["mito2/hdfs-object-store", "object-store/hdfs-object-store"] mysql-object-store = ["object-store/mysql-object-store"] tokio-console = ["common-telemetry/tokio-console"] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index d4b19812df2..cc095c5c92e 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -125,6 +125,44 @@ fn test_load_runtime_options_without_max_blocking_threads() { ); } +#[test] +#[cfg(feature = "hdfs-object-store")] +fn test_load_datanode_hdfs_config() { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + r#" + [storage] + type = "Hdfs" + name = "local-hdfs" + root = "/greptimedb" + name_node = "hdfs://127.0.0.1:9000" + enable_read_cache = false + options = { "dfs.client.block.write.replace-datanode-on-failure.enable" = "true" } + "#, + ) + .unwrap(); + + let options = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let object_store::config::ObjectStoreConfig::Hdfs(hdfs) = options.component.storage.store + else { + unreachable!() + }; + + assert_eq!("local-hdfs", hdfs.name); + assert_eq!("/greptimedb", hdfs.connection.root); + assert_eq!("hdfs://127.0.0.1:9000", hdfs.connection.name_node); + assert_eq!( + Some(&"true".to_string()), + hdfs.connection + .options + .get("dfs.client.block.write.replace-datanode-on-failure.enable") + ); + assert!(!hdfs.cache.enable_read_cache); +} + #[allow(deprecated)] #[test] fn test_load_datanode_example_config() { diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 9ab2dee10df..46032deef38 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -10,6 +10,7 @@ test = ["common-test-util", "rstest", "rstest_reuse", "rskafka"] testing = ["test"] test-shared-fs-region-migration = [] enterprise = [] +hdfs-object-store = ["object-store/hdfs-object-store"] [lints] workspace = true diff --git a/src/mito2/src/region/utils.rs b/src/mito2/src/region/utils.rs index 6917315757a..7dfaad72aa7 100644 --- a/src/mito2/src/region/utils.rs +++ b/src/mito2/src/region/utils.rs @@ -265,7 +265,20 @@ impl RegionFileCopier { #[cfg(test)] mod tests { + #[cfg(feature = "hdfs-object-store")] + use object_store::ObjectStore; + #[cfg(feature = "hdfs-object-store")] + use object_store::layers::HdfsCompatibilityLayer; + #[cfg(feature = "hdfs-object-store")] + use object_store::services::Fs; + use super::*; + #[cfg(feature = "hdfs-object-store")] + use crate::access_layer::AccessLayer; + #[cfg(feature = "hdfs-object-store")] + use crate::sst::index::intermediate::IntermediateManager; + #[cfg(feature = "hdfs-object-store")] + use crate::sst::index::puffin_manager::PuffinManagerFactory; #[test] fn test_build_copy_file_paths() { @@ -344,4 +357,54 @@ mod tests { format!("/table_dir/1_0000000002/index/{}.1.puffin", file_id) ); } + + #[cfg(feature = "hdfs-object-store")] + #[tokio::test] + async fn test_copy_region_files_with_hdfs_fallback() { + let (temp_dir, puffin_manager) = + PuffinManagerFactory::new_for_test_async("hdfs-copy-region").await; + let intermediate_manager = IntermediateManager::init_fs(temp_dir.path().to_string_lossy()) + .await + .unwrap(); + let storage_dir = temp_dir.path().join("storage"); + std::fs::create_dir(&storage_dir).unwrap(); + let object_store = ObjectStore::new(Fs::default().root(storage_dir.to_str().unwrap())) + .unwrap() + .layer(HdfsCompatibilityLayer::new_for_test()); + let access_layer = Arc::new(AccessLayer::new( + "table_dir", + PathType::Bare, + object_store.clone(), + puffin_manager, + intermediate_manager, + )); + let copier = RegionFileCopier::new(access_layer); + let source_region_id = RegionId::new(1, 1); + let target_region_id = RegionId::new(1, 2); + let file_id = FileId::random(); + let descriptor = FileDescriptor::Data { file_id, size: 8 }; + let (source_path, target_path) = build_copy_file_paths( + source_region_id, + target_region_id, + descriptor, + "table_dir", + PathType::Bare, + ); + object_store.write(&source_path, "contents").await.unwrap(); + + copier + .copy_files(source_region_id, target_region_id, vec![descriptor], 1) + .await + .unwrap(); + + assert_eq!( + b"contents", + object_store + .read(&target_path) + .await + .unwrap() + .to_bytes() + .as_ref() + ); + } } diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index 4083271a967..15e2069e3bd 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true workspace = true [features] +hdfs-object-store = ["dep:hdfs-native", "opendal/services-hdfs-native", "uuid"] mysql-object-store = ["opendal/services-mysql"] services-memory = ["opendal/services-memory"] testing = ["derive_builder", "uuid"] @@ -21,6 +22,7 @@ common-macro.workspace = true common-runtime.workspace = true common-telemetry.workspace = true derive_builder = { workspace = true, optional = true } +hdfs-native = { workspace = true, optional = true } humantime-serde.workspace = true opendal = { version = "0.59.2", features = [ "layers-tracing", diff --git a/src/object-store/src/config.rs b/src/object-store/src/config.rs index 1308db92fab..2917bdbe817 100644 --- a/src/object-store/src/config.rs +++ b/src/object-store/src/config.rs @@ -12,10 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "hdfs-object-store")] +use std::collections::HashMap; use std::time::Duration; use common_base::readable_size::ReadableSize; use common_base::secrets::{ExposeSecret, SecretString}; +#[cfg(feature = "hdfs-object-store")] +use opendal::services::HdfsNative; #[cfg(feature = "mysql-object-store")] use opendal::services::Mysql; use opendal::services::{Azblob, Gcs, Oss, S3}; @@ -34,6 +38,8 @@ pub enum ObjectStoreConfig { Oss(OssConfig), Azblob(AzblobConfig), Gcs(GcsConfig), + #[cfg(feature = "hdfs-object-store")] + Hdfs(HdfsConfig), #[cfg(feature = "mysql-object-store")] Mysql(MysqlConfig), } @@ -53,6 +59,8 @@ impl ObjectStoreConfig { Self::Oss(_) => "Oss", Self::Azblob(_) => "Azblob", Self::Gcs(_) => "Gcs", + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(_) => "Hdfs", #[cfg(feature = "mysql-object-store")] Self::Mysql(_) => "Mysql", } @@ -72,6 +80,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => &oss.name, Self::Azblob(az) => &az.name, Self::Gcs(gcs) => &gcs.name, + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => &hdfs.name, #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => &mysql.name, }; @@ -91,6 +101,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => Some(&oss.cache), Self::Azblob(az) => Some(&az.cache), Self::Gcs(gcs) => Some(&gcs.cache), + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => Some(&hdfs.cache), #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => Some(&mysql.cache), } @@ -104,6 +116,8 @@ impl ObjectStoreConfig { Self::Oss(oss) => Some(&mut oss.cache), Self::Azblob(az) => Some(&mut az.cache), Self::Gcs(gcs) => Some(&mut gcs.cache), + #[cfg(feature = "hdfs-object-store")] + Self::Hdfs(hdfs) => Some(&mut hdfs.cache), #[cfg(feature = "mysql-object-store")] Self::Mysql(mysql) => Some(&mut mysql.cache), } @@ -289,6 +303,42 @@ impl From<&GcsConnection> for Gcs { } } +/// Connection options for a Hadoop Distributed File System backend. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +#[serde(default)] +#[cfg(feature = "hdfs-object-store")] +pub struct HdfsConnection { + /// Working directory for all object-store operations. + pub root: String, + /// HDFS NameNode URI, for example `hdfs://127.0.0.1:9000`. + pub name_node: String, + /// Additional options passed to the native HDFS client. + pub options: HashMap, +} + +#[cfg(feature = "hdfs-object-store")] +impl From<&HdfsConnection> for HdfsNative { + fn from(connection: &HdfsConnection) -> Self { + let root = util::normalize_dir(&connection.root); + HdfsNative::default() + .root(&root) + .name_node(&connection.name_node) + .options(connection.options.clone()) + } +} + +/// Hadoop Distributed File System object storage configuration. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +#[serde(default)] +#[cfg(feature = "hdfs-object-store")] +pub struct HdfsConfig { + pub name: String, + #[serde(flatten)] + pub connection: HdfsConnection, + #[serde(flatten)] + pub cache: ObjectStorageCacheConfig, +} + #[cfg(feature = "mysql-object-store")] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] #[serde(default)] @@ -411,6 +461,20 @@ mod tests { assert_eq!("test", s3_config.config_name()); assert_eq!("S3", s3_config.provider_name()); + #[cfg(feature = "hdfs-object-store")] + { + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig::default()); + assert_eq!("Hdfs", hdfs_config.config_name()); + assert_eq!("Hdfs", hdfs_config.provider_name()); + + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig { + name: "test".to_string(), + ..Default::default() + }); + assert_eq!("test", hdfs_config.config_name()); + assert_eq!("Hdfs", hdfs_config.provider_name()); + } + #[cfg(feature = "mysql-object-store")] { let mysql_config = ObjectStoreConfig::Mysql(MysqlConfig::default()); @@ -438,6 +502,11 @@ mod tests { assert!(gcs_config.is_object_storage()); let azblob_config = ObjectStoreConfig::Azblob(AzblobConfig::default()); assert!(azblob_config.is_object_storage()); + #[cfg(feature = "hdfs-object-store")] + { + let hdfs_config = ObjectStoreConfig::Hdfs(HdfsConfig::default()); + assert!(hdfs_config.is_object_storage()); + } #[cfg(feature = "mysql-object-store")] { let mysql_config = ObjectStoreConfig::Mysql(MysqlConfig::default()); @@ -445,6 +514,46 @@ mod tests { } } + #[cfg(feature = "hdfs-object-store")] + #[test] + fn test_hdfs_config_serde() { + let config: ObjectStoreConfig = toml::from_str( + r#" +type = "Hdfs" +name = "hdfs-store" +root = "/greptimedb" +name_node = "hdfs://127.0.0.1:9000" + +[options] +"dfs.client.block.write.replace-datanode-on-failure.enable" = "true" +"#, + ) + .unwrap(); + + let ObjectStoreConfig::Hdfs(hdfs_config) = config else { + unreachable!() + }; + + assert_eq!("hdfs-store", hdfs_config.name); + assert_eq!("/greptimedb", hdfs_config.connection.root); + assert_eq!("hdfs://127.0.0.1:9000", hdfs_config.connection.name_node); + assert_eq!( + Some(&"true".to_string()), + hdfs_config + .connection + .options + .get("dfs.client.block.write.replace-datanode-on-failure.enable") + ); + + let serialized = toml::to_string(&hdfs_config).unwrap(); + assert!(serialized.contains("name_node = \"hdfs://127.0.0.1:9000\"")); + assert!( + serialized.contains( + "\"dfs.client.block.write.replace-datanode-on-failure.enable\" = \"true\"" + ) + ); + } + #[cfg(feature = "mysql-object-store")] #[test] fn test_mysql_config_connection_string_serde() { diff --git a/src/object-store/src/factory.rs b/src/object-store/src/factory.rs index 9bdd656a07c..356d3aecb21 100644 --- a/src/object-store/src/factory.rs +++ b/src/object-store/src/factory.rs @@ -15,15 +15,21 @@ use std::{fs, path}; use common_telemetry::info; +#[cfg(feature = "hdfs-object-store")] +use opendal::services::HdfsNative; #[cfg(feature = "mysql-object-store")] use opendal::services::Mysql; use opendal::services::{Fs, Gcs, Oss, S3}; use snafu::prelude::*; +#[cfg(feature = "hdfs-object-store")] +use crate::config::HdfsConfig; #[cfg(feature = "mysql-object-store")] use crate::config::MysqlConfig; use crate::config::{AzblobConfig, FileConfig, GcsConfig, ObjectStoreConfig, OssConfig, S3Config}; use crate::error::{self, Result}; +#[cfg(feature = "hdfs-object-store")] +use crate::layers::HdfsCompatibilityLayer; use crate::services::Azblob; use crate::util::{build_http_context, clean_temp_dir, join_dir, normalize_dir}; use crate::{ATOMIC_WRITE_DIR, OLD_ATOMIC_WRITE_DIR, ObjectStore, util}; @@ -39,11 +45,36 @@ pub async fn new_raw_object_store( ObjectStoreConfig::Oss(oss_config) => new_oss_object_store(oss_config).await, ObjectStoreConfig::Azblob(azblob_config) => new_azblob_object_store(azblob_config).await, ObjectStoreConfig::Gcs(gcs_config) => new_gcs_object_store(gcs_config).await, + #[cfg(feature = "hdfs-object-store")] + ObjectStoreConfig::Hdfs(hdfs_config) => new_hdfs_object_store(hdfs_config).await, #[cfg(feature = "mysql-object-store")] ObjectStoreConfig::Mysql(mysql_config) => new_mysql_object_store(mysql_config).await, } } +/// Creates an object store backed by a native HDFS client. +#[cfg(feature = "hdfs-object-store")] +pub async fn new_hdfs_object_store(hdfs_config: &HdfsConfig) -> Result { + let root = util::normalize_dir(&hdfs_config.connection.root); + info!( + "The HDFS NameNode is: {}, root is: {}", + hdfs_config.connection.name_node, root + ); + + let builder = HdfsNative::from(&hdfs_config.connection); + let compatibility_layer = HdfsCompatibilityLayer::new( + &hdfs_config.connection.name_node, + &hdfs_config.connection.root, + &hdfs_config.connection.options, + ) + .context(error::InitBackendSnafu)?; + let operator = ObjectStore::new(builder) + .context(error::InitBackendSnafu)? + .layer(compatibility_layer); + + Ok(operator) +} + #[cfg(feature = "mysql-object-store")] pub async fn new_mysql_object_store(mysql_config: &MysqlConfig) -> Result { let root = util::normalize_dir(&mysql_config.root); @@ -144,3 +175,33 @@ pub async fn new_s3_object_store(s3_config: &S3Config) -> Result { Ok(operator) } + +#[cfg(all(test, feature = "hdfs-object-store"))] +mod tests { + use opendal::services::HDFS_NATIVE_SCHEME; + + use super::*; + use crate::config::HdfsConnection; + + #[tokio::test] + async fn test_new_hdfs_object_store() { + let config = HdfsConfig { + connection: HdfsConnection { + root: "/greptimedb".to_string(), + name_node: "hdfs://127.0.0.1:9000".to_string(), + ..Default::default() + }, + ..Default::default() + }; + + let store = new_hdfs_object_store(&config).await.unwrap(); + assert_eq!(HDFS_NATIVE_SCHEME, store.info().scheme()); + assert_eq!("/greptimedb/", store.info().root()); + } + + #[tokio::test] + async fn test_new_hdfs_object_store_requires_name_node() { + let result = new_hdfs_object_store(&HdfsConfig::default()).await; + assert!(result.is_err()); + } +} diff --git a/src/object-store/src/layers.rs b/src/object-store/src/layers.rs index cc9fa9f4df5..77c894dbfcd 100644 --- a/src/object-store/src/layers.rs +++ b/src/object-store/src/layers.rs @@ -12,9 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "hdfs-object-store")] +mod hdfs; #[cfg(feature = "testing")] pub mod mock; +#[cfg(feature = "hdfs-object-store")] +pub use hdfs::HdfsCompatibilityLayer; pub use opendal::layers::*; pub use prometheus::build_prometheus_metrics_layer; diff --git a/src/object-store/src/layers/hdfs.rs b/src/object-store/src/layers/hdfs.rs new file mode 100644 index 00000000000..4e00098fbf7 --- /dev/null +++ b/src/object-store/src/layers/hdfs.rs @@ -0,0 +1,551 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt::{self, Debug}; +use std::sync::Arc; + +use hdfs_native::{Client, ClientBuilder}; +use opendal::raw::oio::{Delete as _, Read as _, ReadStream as _, Write as _}; +use opendal::raw::{ + Layer, OpCompose, OpCopy, OpCreateDir, OpDelete, OpList, OpPresign, OpRead, OpRename, OpStat, + OpWrite, RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, Servicer, oio, +}; +use opendal::{Buffer, Capability, ErrorKind, Metadata, OperationContext, Result}; +use uuid::Uuid; + +/// Adds atomic writes and streaming copies to the native HDFS backend. +#[derive(Debug, Clone)] +pub struct HdfsCompatibilityLayer { + renamer: AtomicRenamer, +} + +impl HdfsCompatibilityLayer { + /// Creates a compatibility layer for the HDFS connection. + pub fn new( + name_node: &str, + root: &str, + options: &std::collections::HashMap, + ) -> Result { + let mut config = std::collections::HashMap::new(); + let namenodes = name_node + .split(',') + .filter_map(|value| { + let value = value + .trim() + .trim_start_matches("hdfs://") + .trim_end_matches('/'); + (!value.is_empty()).then_some(value) + }) + .collect::>(); + for (index, namenode) in namenodes.iter().enumerate() { + config.insert( + format!("dfs.namenode.rpc-address.nameservice.nn{index}"), + (*namenode).to_string(), + ); + } + config.insert( + "dfs.ha.namenodes.nameservice".to_string(), + (0..namenodes.len()) + .map(|index| format!("nn{index}")) + .collect::>() + .join(","), + ); + config.extend(options.clone()); + let client = ClientBuilder::new() + .with_url("hdfs://nameservice") + .with_config(config) + .build() + .map_err(hdfs_error)?; + Ok(Self { + renamer: AtomicRenamer::Native { + client, + root: opendal::raw::normalize_root(root), + }, + }) + } + + /// Creates a compatibility layer backed by the inner service's rename. + #[cfg(any(test, feature = "testing"))] + pub fn new_for_test() -> Self { + Self { + renamer: AtomicRenamer::Raw, + } + } +} + +#[derive(Debug, Clone)] +enum AtomicRenamer { + Native { + client: Client, + root: String, + }, + #[cfg(any(test, feature = "testing"))] + Raw, +} + +impl AtomicRenamer { + async fn rename( + &self, + _inner: &Servicer, + _ctx: &OperationContext, + from: &str, + to: &str, + ) -> Result<()> { + match self { + Self::Native { client, root } => { + // OpenDAL's HDFS rename removes an existing destination before + // renaming. Use HDFS Rename2 with overwrite to keep replacement atomic. + client + .rename( + &opendal::raw::build_rooted_abs_path(root, from), + &opendal::raw::build_rooted_abs_path(root, to), + true, + ) + .await + .map_err(hdfs_error) + } + #[cfg(any(test, feature = "testing"))] + Self::Raw => _inner + .rename(_ctx, from, to, OpRename::new()) + .await + .map(|_| ()), + } + } +} + +fn hdfs_error(error: hdfs_native::HdfsError) -> opendal::Error { + opendal::Error::new(ErrorKind::Unexpected, "native HDFS operation failed").set_source(error) +} + +impl Layer for HdfsCompatibilityLayer { + fn apply_service(&self, inner: Servicer) -> Servicer { + Arc::new(HdfsCompatibilityService { + inner, + renamer: self.renamer.clone(), + }) + } +} + +struct HdfsCompatibilityService { + inner: Servicer, + renamer: AtomicRenamer, +} + +impl Debug for HdfsCompatibilityService { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HdfsCompatibilityService") + .field("inner", &self.inner) + .finish() + } +} + +/// A writer that publishes non-append writes with an atomic rename. +struct HdfsWriter(HdfsWriterInner); + +enum HdfsWriterInner { + Direct(oio::Writer), + Atomic { + inner: Servicer, + context: OperationContext, + renamer: AtomicRenamer, + writer: Option, + temporary_path: String, + target_path: String, + }, +} + +impl oio::Write for HdfsWriter { + async fn write(&mut self, buffer: Buffer) -> Result<()> { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.write(buffer).await, + HdfsWriterInner::Atomic { writer, .. } => { + writer + .as_mut() + .ok_or_else(writer_unavailable)? + .write(buffer) + .await + } + } + } + + async fn close(&mut self) -> Result { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.close().await, + HdfsWriterInner::Atomic { + inner, + context, + renamer, + writer, + temporary_path, + target_path, + } => { + let mut writer = writer.take().ok_or_else(writer_unavailable)?; + let metadata = match writer.close().await { + Ok(metadata) => metadata, + Err(error) => { + drop(writer); + let _ = delete_path(inner, context, temporary_path).await; + return Err(error); + } + }; + drop(writer); + + if let Err(error) = renamer + .rename(inner, context, temporary_path, target_path) + .await + { + let _ = delete_path(inner, context, temporary_path).await; + return Err(error); + } + + Ok(metadata) + } + } + } + + async fn abort(&mut self) -> Result<()> { + match &mut self.0 { + HdfsWriterInner::Direct(writer) => writer.abort().await, + HdfsWriterInner::Atomic { + inner, + context, + writer, + temporary_path, + .. + } => { + let abort_result = if let Some(mut writer) = writer.take() { + let result = writer.abort().await; + drop(writer); + result + } else { + Ok(()) + }; + let cleanup_result = delete_path(inner, context, temporary_path).await; + + match abort_result { + Err(error) if error.kind() != ErrorKind::Unsupported => Err(error), + _ => cleanup_result, + } + } + } + } +} + +fn writer_unavailable() -> opendal::Error { + opendal::Error::new( + ErrorKind::Unexpected, + "HDFS writer is unavailable after close or abort", + ) +} + +impl Service for HdfsCompatibilityService { + type Reader = oio::Reader; + type Writer = HdfsWriter; + type Lister = oio::Lister; + type Deleter = oio::Deleter; + type Copier = oio::OneShotCopier; + type Composer = oio::Composer; + + fn info(&self) -> ServiceInfo { + self.inner.info() + } + + fn capability(&self) -> Capability { + let mut capability = self.inner.capability(); + capability.copy = true; + capability + } + + async fn create_dir( + &self, + ctx: &OperationContext, + path: &str, + args: OpCreateDir, + ) -> Result { + self.inner.create_dir(ctx, path, args).await + } + + async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result { + self.inner.stat(ctx, path, args).await + } + + fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result { + self.inner.read(ctx, path, args) + } + + fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result { + if args.append() { + return self + .inner + .write(ctx, path, args) + .map(|writer| HdfsWriter(HdfsWriterInner::Direct(writer))); + } + + let temporary_path = temporary_path(path); + let writer = self.inner.write(ctx, &temporary_path, args)?; + Ok(HdfsWriter(HdfsWriterInner::Atomic { + inner: Arc::clone(&self.inner), + context: ctx.clone(), + renamer: self.renamer.clone(), + writer: Some(writer), + temporary_path, + target_path: path.to_string(), + })) + } + + fn copy( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpCopy, + ) -> Result { + if args.if_not_exists() || args.if_match().is_some() { + return Err(opendal::Error::new( + ErrorKind::Unsupported, + "conditional copy is not supported by the HDFS fallback", + )); + } + + let inner = Arc::clone(&self.inner); + let context = ctx.clone(); + let renamer = self.renamer.clone(); + let from = from.to_string(); + let to = to.to_string(); + Ok(oio::OneShotCopier::new(async move { + copy_via_read_write(inner, &context, renamer, &from, &to).await + })) + } + + fn delete(&self, ctx: &OperationContext) -> Result { + self.inner.delete(ctx) + } + + fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result { + self.inner.list(ctx, path, args) + } + + fn compose(&self, ctx: &OperationContext, to: &str, args: OpCompose) -> Result { + self.inner.compose(ctx, to, args) + } + + async fn rename( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpRename, + ) -> Result { + self.inner.rename(ctx, from, to, args).await + } + + async fn presign( + &self, + ctx: &OperationContext, + path: &str, + args: OpPresign, + ) -> Result { + self.inner.presign(ctx, path, args).await + } +} + +async fn copy_via_read_write( + inner: Servicer, + context: &OperationContext, + renamer: AtomicRenamer, + source_path: &str, + target_path: &str, +) -> Result { + let reader = inner.read(context, source_path, OpRead::new())?; + let (_, mut reader) = reader.open(opendal::BytesRange::from(..)).await?; + let temporary_path = temporary_path(target_path); + let mut writer = inner.write(context, &temporary_path, OpWrite::new())?; + + loop { + let buffer = match reader.read().await { + Ok(buffer) => buffer, + Err(error) => { + abort_and_delete(&inner, context, writer, &temporary_path).await; + return Err(error); + } + }; + if buffer.is_empty() { + break; + } + if let Err(error) = writer.write(buffer).await { + abort_and_delete(&inner, context, writer, &temporary_path).await; + return Err(error); + } + } + + let metadata = match writer.close().await { + Ok(metadata) => metadata, + Err(error) => { + drop(writer); + let _ = delete_path(&inner, context, &temporary_path).await; + return Err(error); + } + }; + drop(writer); + + if let Err(error) = renamer + .rename(&inner, context, &temporary_path, target_path) + .await + { + let _ = delete_path(&inner, context, &temporary_path).await; + return Err(error); + } + + Ok(metadata) +} + +async fn abort_and_delete( + inner: &Servicer, + context: &OperationContext, + mut writer: oio::Writer, + path: &str, +) { + let _ = writer.abort().await; + drop(writer); + let _ = delete_path(inner, context, path).await; +} + +async fn delete_path(inner: &Servicer, context: &OperationContext, path: &str) -> Result<()> { + let mut deleter = inner.delete(context)?; + deleter.delete(path, OpDelete::new()).await?; + deleter.close().await +} + +// TODO(fengjiachun): Clean up temporary files left behind after a process crash. +fn temporary_path(path: &str) -> String { + let suffix = format!(".greptime-{}.tmp", Uuid::new_v4()); + match path.rsplit_once('/') { + Some((parent, name)) => format!("{parent}/.{name}{suffix}"), + None => format!(".{path}{suffix}"), + } +} + +#[cfg(test)] +mod tests { + use opendal::services::Fs; + use opendal::{Operator, Writer}; + use tempfile::TempDir; + + use super::*; + + fn test_store() -> (TempDir, Operator) { + let directory = tempfile::tempdir().unwrap(); + let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap())) + .unwrap() + .layer(HdfsCompatibilityLayer::new_for_test()); + (directory, store) + } + + #[tokio::test] + async fn test_service_operations() { + let (_directory, store) = test_store(); + store.create_dir("data/").await.unwrap(); + store.write("data/source", "contents").await.unwrap(); + assert_eq!(8, store.stat("data/source").await.unwrap().content_length()); + assert_eq!( + b"onte", + store + .read_with("data/source") + .range(1..5) + .await + .unwrap() + .to_bytes() + .as_ref() + ); + store.rename("data/source", "data/target").await.unwrap(); + let entries = store.list("data/").await.unwrap(); + assert!(entries.iter().any(|entry| entry.path() == "data/target")); + assert!(!store.exists("data/source").await.unwrap()); + store.delete("data/target").await.unwrap(); + assert!(!store.exists("data/target").await.unwrap()); + } + + #[tokio::test] + async fn test_atomic_write_keeps_old_data_after_abort() { + let (_directory, store) = test_store(); + store.write("manifest.json", "old").await.unwrap(); + + let mut writer: Writer = store.writer("manifest.json").await.unwrap(); + writer.write("new").await.unwrap(); + writer.abort().await.unwrap(); + + assert_eq!( + b"old", + store + .read("manifest.json") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + assert!( + store + .list("") + .await + .unwrap() + .iter() + .all(|entry| !entry.path().contains(".greptime-")) + ); + } + + #[tokio::test] + async fn test_atomic_write_replaces_on_close() { + let (_directory, store) = test_store(); + store.write("manifest.json", "old").await.unwrap(); + store.write("manifest.json", "new").await.unwrap(); + + assert_eq!( + b"new", + store + .read("manifest.json") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + assert!( + store + .list("") + .await + .unwrap() + .iter() + .all(|entry| !entry.path().contains(".greptime-")) + ); + } + + #[tokio::test] + async fn test_copy_fallback_streams_to_target() { + let (_directory, store) = test_store(); + store.write("source.parquet", "contents").await.unwrap(); + store + .copy("source.parquet", "target.parquet") + .await + .unwrap(); + + assert_eq!( + b"contents", + store + .read("target.parquet") + .await + .unwrap() + .to_bytes() + .as_ref() + ); + } +} diff --git a/src/object-store/src/lib.rs b/src/object-store/src/lib.rs index 2de2b570082..14a2c3d1c68 100644 --- a/src/object-store/src/lib.rs +++ b/src/object-store/src/lib.rs @@ -30,6 +30,8 @@ pub mod secure_fs; pub mod test_util; pub mod util; +#[cfg(feature = "hdfs-object-store")] +pub use config::HdfsConnection; pub use config::{AzblobConnection, GcsConnection, OssConnection, S3Connection}; /// The default object cache directory name. From a310ca2bcf0b70bdc2a8184b5251778446e84f3f Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Mon, 28 Sep 2026 09:16:23 +0000 Subject: [PATCH 057/133] ci(query-regression): bump RUNNER_IMAGE_EPOCH to 6 (#9381) ci(query-regression): bump RUNNER_IMAGE_EPOCH to 6 for image m-0xidkbbavm3suxqd0y1m Co-authored-by: greptimedb-ci --- .github/workflows/query-regression.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 8ce77827de7..cc8afaf8091 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -553,7 +553,7 @@ jobs: readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" readonly RUNNER_IMAGE_DIGEST="sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11" - readonly RUNNER_IMAGE_EPOCH="5" + readonly RUNNER_IMAGE_EPOCH="6" require_expected_root() { local name="$1" From 03823a9a01bfafb3c2eaa0b5e6db3a795e37a14e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 28 Sep 2026 09:17:45 +0000 Subject: [PATCH 058/133] feat(log-store): add the enqueued acknowledgement mode to the object store WAL (#9358) * feat(log-store): add the enqueued acknowledgement mode to the object store WAL Add `ack_mode` (`durable` by default, or `enqueued`) and the backlog thresholds `max_unpersisted_bytes` and `max_unpersisted_age` to the object store WAL config, validated by the datanode and the store. In the `enqueued` mode `append_batch` returns on admission with the entry ids assigned and the object is created in the background. At a backlog threshold the next append is held back until an upload completes. A transient create failure is repeated under the same sequence with the same bytes; any conflicting object poisons the store. `stop` uploads the backlog, or returns the error that dropped it once stop began. `obsolete` clamps the watermark to the durable entry id, and an id the store handed out needs no sequence floor. Add `LogStore::wait_durable` with a default that returns at once. The object store WAL answers it once the region is durable and indexed through the entry id, and fails it after a backlog was lost. Signed-off-by: jeremyhi * fix(log-store): poison an enqueued store on permanent create failures Repeat a failed create in the enqueued mode only when the storage error is retryable; any other storage error poisons the store. A conflicting object poisons an enqueued store without reading its epoch, so a failed header read cannot turn the conflict into a retry. A durability wait for an id above the highest id the store handed out now waits for the handed-out ids of the region below it instead of returning at once. Signed-off-by: jeremyhi * fix(log-store): poison on permanent create failures after stop begins A create that fails with a storage error that is not retryable poisons an enqueued store even after stop began; only a transient failure drops the backlog without poisoning. Durability waiters whose callers stopped waiting are pruned before a new waiter is queued. The backlog age test no longer depends on a follow-up append finishing within the threshold. Signed-off-by: jeremyhi * fix(log-store): answer a durability wait once no earlier entry is pending A durability wait now returns once the region holds no entry at or below the target that is handed out but not durable, instead of waiting for the largest id handed out to the region. A later object of the region that is still being created no longer holds back a wait whose target it does not cover. Signed-off-by: jeremyhi * test(log-store): order the pending durability wait check after the actor Signed-off-by: jeremyhi * docs(store-api): state that the default wait_durable keeps each store's guarantee The default `LogStore::wait_durable` returns at once, which keeps each log store's own acknowledgement guarantee; Raft Engine with `sync_write = false` acknowledges before its periodic sync, so the documentation no longer claims that every entry id a caller holds is durable. The object store WAL configuration test now also serializes the new options and reads them back. Signed-off-by: jeremyhi * docs(log-store): limit the acknowledgement guarantees to the durable mode Signed-off-by: jeremyhi * fix(log-store): repeat enqueued creates that a retry layer marks persistent An object store wrapped in the OpenDAL retry layer reports a temporary error that outlasted its retries as persistent rather than temporary. The enqueued mode now repeats a create after any storage error that is not permanent, so a transient outage behind the retry layer no longer poisons the store and drops the acknowledged backlog; after stop began such a failure still drops the backlog without poisoning. Signed-off-by: jeremyhi * test(log-store): cover a persistent create failure after stop begins Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- config/config.md | 10 +- config/datanode.example.toml | 17 +- config/standalone.example.toml | 17 +- src/common/wal/src/config.rs | 20 +- src/common/wal/src/config/object_store.rs | 24 + src/datanode/src/datanode.rs | 32 + src/log-store/src/object_store_wal/batch.rs | 4 + src/log-store/src/object_store_wal/store.rs | 1069 +++++++++++++++++-- src/store-api/src/logstore.rs | 18 + 9 files changed, 1140 insertions(+), 71 deletions(-) diff --git a/config/config.md b/config/config.md index cebc5091b4c..89b3e13dcdd 100644 --- a/config/config.md +++ b/config/config.md @@ -125,7 +125,10 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | +| `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | | `metadata_store` | -- | -- | Metadata storage options. | | `metadata_store.file_size` | String | `64MB` | The size of the metadata store log file. | @@ -599,7 +602,10 @@ | `wal.overwrite_entry_start_id` | Bool | `false` | Ignore missing entries during read WAL.
**It's only used when the provider is `kafka`**.

This option ensures that when Kafka messages are deleted, the system
can still successfully replay memtable data without throwing an
out-of-range error.
However, enabling this option might lead to unexpected data loss,
as the system will skip over missing entries instead of treating
them as critical errors. | | `wal.storage_provider` | String | `""` | The name of the storage provider that holds the WAL objects, an empty name selects the default object store.
**It's only used when the provider is `experimental_object_store`**. | | `wal.prefix` | String | `wal` | The path prefix of the WAL objects inside the storage provider.
The objects are written under `/datanodes//epochs/`, which is derived from this prefix.
**It's only used when the provider is `experimental_object_store`**. | -| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.flush_interval` | String | `100ms` | The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`.
Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters.
When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.ack_mode` | String | `durable` | When an append to the object store WAL returns.
- `durable`: an append returns after the object holding its entries is durable (the default).
- `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog.
**It's only used when the provider is `experimental_object_store`**. | +| `wal.max_unpersisted_bytes` | String | `64MB` | The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode.
Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | +| `wal.max_unpersisted_age` | String | `8s` | The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode.
**It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. | | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 253cc3f2a57..82b93fd6ccc 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -231,10 +231,25 @@ overwrite_entry_start_id = false ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. ## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. -## When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node. +## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" +## When an append to the object store WAL returns. +## - `durable`: an append returns after the object holding its entries is durable (the default). +## - `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog. +## **It's only used when the provider is `experimental_object_store`**. +#+ ack_mode = "durable" + +## The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode. +## Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_bytes = "64MB" + +## The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_age = "8s" + ## What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry. ## - `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default). ## - `fail`: the read fails, so the region does not open. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index fad4f64f08a..e5eaa86c4c5 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -414,10 +414,25 @@ overwrite_entry_start_id = false ## The interval of flushing buffered entries to the object store, at least `10ms`, defaults to `100ms`. ## Each non-empty timer-triggered flush creates one object, and a batch that reaches `max_batch_bytes` is flushed immediately, so under sustained load the batch seals on size and the interval no longer matters. -## When writes are sparse, a shorter interval lowers the acknowledgement latency of appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node. +## When writes are sparse, a shorter interval lowers the acknowledgement latency of `durable` appends and raises the number of object requests: timer-triggered sealing creates at most one object per interval per node, and in `enqueued` mode the backlog thresholds can seal earlier. An `enqueued` append does not wait for its batch to seal once it is admitted, but the backlog thresholds can delay admission. ## **It's only used when the provider is `experimental_object_store`**. #+ flush_interval = "100ms" +## When an append to the object store WAL returns. +## - `durable`: an append returns after the object holding its entries is durable (the default). +## - `enqueued`: an append returns once its entries are admitted and their ids are assigned, and the object is created in the background. A crash loses the unpersisted backlog. +## **It's only used when the provider is `experimental_object_store`**. +#+ ack_mode = "durable" + +## The size of the unpersisted backlog at which new appends stall until an upload completes, in `enqueued` mode. +## Nothing is dropped and nothing is rejected; the threshold does not bound what a crash during an outage can lose. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_bytes = "64MB" + +## The age of the oldest unpersisted entry at which new appends stall until an upload completes, in `enqueued` mode. +## **It's only used when the provider is `experimental_object_store` and `ack_mode` is `enqueued`**. +#+ max_unpersisted_age = "8s" + ## What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry. ## - `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default). ## - `fail`: the read fails, so the region does not open. diff --git a/src/common/wal/src/config.rs b/src/common/wal/src/config.rs index f79a28607cc..0cbebfb20ab 100644 --- a/src/common/wal/src/config.rs +++ b/src/common/wal/src/config.rs @@ -138,7 +138,7 @@ mod tests { use super::*; use crate::TopicSelectorType; - use crate::config::object_store::CorruptedSegmentAction; + use crate::config::object_store::{AckMode, CorruptedSegmentAction}; use crate::config::{DatanodeKafkaConfig, MetasrvKafkaConfig}; #[test] @@ -294,6 +294,9 @@ mod tests { assert_eq!(config.prefix, "wal"); assert_eq!(config.flush_interval, Duration::from_millis(100)); assert_eq!(config.max_batch_bytes, ReadableSize::mb(8)); + assert_eq!(config.ack_mode, AckMode::Durable); + assert_eq!(config.max_unpersisted_bytes, ReadableSize::mb(64)); + assert_eq!(config.max_unpersisted_age, Duration::from_secs(8)); assert_eq!(config.on_corrupted_segment, CorruptedSegmentAction::Skip); assert!(matches!( MetasrvWalConfig::try_from(datanode_wal_config).unwrap_err(), @@ -306,6 +309,9 @@ mod tests { prefix = "cluster-a/wal" flush_interval = "500ms" max_batch_bytes = "4MB" + ack_mode = "enqueued" + max_unpersisted_bytes = "32MB" + max_unpersisted_age = "4s" on_corrupted_segment = "fail" "#; let datanode_wal_config: DatanodeWalConfig = toml::from_str(toml_str).unwrap(); @@ -314,12 +320,24 @@ mod tests { prefix: "cluster-a/wal".to_string(), flush_interval: Duration::from_millis(500), max_batch_bytes: ReadableSize::mb(4), + ack_mode: AckMode::Enqueued, + max_unpersisted_bytes: ReadableSize::mb(32), + max_unpersisted_age: Duration::from_secs(4), on_corrupted_segment: CorruptedSegmentAction::Fail, }; assert_eq!( datanode_wal_config, DatanodeWalConfig::ObjectStore(expected) ); + let serialized = toml::to_string(&datanode_wal_config).unwrap(); + let table: toml::Table = toml::from_str(&serialized).unwrap(); + assert_eq!(table["ack_mode"].as_str(), Some("enqueued")); + assert_eq!(table["max_unpersisted_bytes"].as_str(), Some("32MiB")); + assert_eq!(table["max_unpersisted_age"].as_str(), Some("4s")); + assert_eq!( + datanode_wal_config, + toml::from_str::(&serialized).unwrap() + ); // The persisted tag is not accepted as a config provider. let toml_str = r#" diff --git a/src/common/wal/src/config/object_store.rs b/src/common/wal/src/config/object_store.rs index 04a212fb53f..dd2c4544129 100644 --- a/src/common/wal/src/config/object_store.rs +++ b/src/common/wal/src/config/object_store.rs @@ -24,6 +24,18 @@ use serde::{Deserialize, Serialize}; /// node id and the generation from the metasrv. pub const STANDALONE_GENERATION: u64 = 0; +/// When an append to the object store WAL returns. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum AckMode { + /// An append returns after the object holding its entries is durable. + #[default] + Durable, + /// An append returns once its entries are admitted and their ids are + /// assigned; the object is created in the background. + Enqueued, +} + /// What a read does with a segment that still does not decode after a /// second fetch, because its checksum does not match or its content /// disagrees with its footer entry. @@ -56,6 +68,15 @@ pub struct ObjectStoreWalConfig { pub flush_interval: Duration, /// The max size of a single batch object, defaults to 8MiB. pub max_batch_bytes: ReadableSize, + /// When an append returns, defaults to `durable`. + pub ack_mode: AckMode, + /// Size of the unpersisted backlog at which appends stall until an + /// upload completes in `enqueued` mode, defaults to 64MiB. + pub max_unpersisted_bytes: ReadableSize, + /// Age of the oldest unpersisted entry at which appends stall until an + /// upload completes in `enqueued` mode, defaults to 8s. + #[serde(with = "humantime_serde")] + pub max_unpersisted_age: Duration, /// What a read does with a segment that still does not decode after a /// second fetch, because its checksum does not match or its content /// disagrees with its footer entry, defaults to `skip`. @@ -69,6 +90,9 @@ impl Default for ObjectStoreWalConfig { prefix: "wal".to_string(), flush_interval: Duration::from_millis(100), max_batch_bytes: ReadableSize::mb(8), + ack_mode: AckMode::Durable, + max_unpersisted_bytes: ReadableSize::mb(64), + max_unpersisted_age: Duration::from_secs(8), on_corrupted_segment: CorruptedSegmentAction::Skip, } } diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index ae4b0723e1a..b3d2c981f1e 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -780,6 +780,22 @@ fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> reason: "must be greater than 0", } ); + ensure!( + config.max_unpersisted_bytes.as_bytes() > 0, + InvalidObjectStoreWalConfigSnafu { + field: "max_unpersisted_bytes", + value: config.max_unpersisted_bytes.to_string(), + reason: "must be greater than 0", + } + ); + ensure!( + config.max_unpersisted_age > Duration::ZERO, + InvalidObjectStoreWalConfigSnafu { + field: "max_unpersisted_age", + value: format!("{:?}", config.max_unpersisted_age), + reason: "must be greater than 0", + } + ); Ok(()) } @@ -1163,12 +1179,28 @@ mod tests { }, "max_batch_bytes", ), + ( + ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(0), + ..Default::default() + }, + "max_unpersisted_bytes", + ), + ( + ObjectStoreWalConfig { + max_unpersisted_age: Duration::ZERO, + ..Default::default() + }, + "max_unpersisted_age", + ), ]; for (config, expected_field) in cases { let expected_value = match expected_field { "prefix" => config.prefix.clone(), "flush_interval" => format!("{:?}", config.flush_interval), + "max_unpersisted_bytes" => config.max_unpersisted_bytes.to_string(), + "max_unpersisted_age" => format!("{:?}", config.max_unpersisted_age), _ => config.max_batch_bytes.to_string(), }; let data_home = create_temp_dir("object-store-wal-invalid-config"); diff --git a/src/log-store/src/object_store_wal/batch.rs b/src/log-store/src/object_store_wal/batch.rs index b07ac8e8cb0..078af921def 100644 --- a/src/log-store/src/object_store_wal/batch.rs +++ b/src/log-store/src/object_store_wal/batch.rs @@ -145,6 +145,10 @@ impl OpenBatch { self.entries.is_empty() } + pub(crate) fn holds_region(&self, region_id: RegionId) -> bool { + self.positions.contains_key(®ion_id) + } + /// Returns the estimated size of the admitted entries. pub(crate) fn estimated_bytes(&self) -> usize { self.estimated_bytes diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index ee925ae0a38..13d4220c6ab 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -19,12 +19,12 @@ use std::fmt; use std::ops::Range; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError, RwLock}; -use std::time::Duration; +use std::time::{Duration, Instant}; use async_stream::try_stream; use bytes::Bytes; use common_telemetry::info; -use common_wal::config::object_store::ObjectStoreWalConfig; +use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig}; use futures::future::BoxFuture; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt}; @@ -46,7 +46,7 @@ use crate::error::{ StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, WalObjectSequenceUnsettledSnafu, }; -use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, sequence_floor}; +use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ ChainLink, EncodedObject, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, @@ -60,6 +60,9 @@ const COMMAND_BUFFER: usize = 1024; const APPEND_BUFFER: usize = 16; const MIN_FLUSH_INTERVAL: Duration = Duration::from_millis(10); const MAX_IN_FLIGHT_CREATES: usize = 4; +/// Delay before a create that failed transiently is attempted again in the +/// `enqueued` acknowledgement mode, where no caller is left to retry it. +const CREATE_RETRY_DELAY: Duration = Duration::from_millis(100); /// Number of sealed batches that may wait to be created or indexed. One /// admission can seal two batches, the open batch before an append that would /// exhaust its positions and the append's own batch, so the actor takes an @@ -79,13 +82,15 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// when it reaches the size limit or the flush interval elapses and creates /// the object under the next sequence while it keeps admitting entries into /// the next batch; up to [`MAX_IN_FLIGHT_CREATES`] creates run at a time. -/// Objects are indexed in the catalog and their batches acknowledged in -/// sequence order, so an acknowledged entry never has a missing predecessor. -/// An append returns once its object is durable and indexed. While -/// [`MAX_SEALED_BATCHES`] batches wait to become durable no append is -/// admitted. +/// Objects are indexed in the catalog in sequence order. In the `durable` +/// acknowledgement mode an append returns once its object is durable and +/// indexed, so an acknowledged entry never has a missing predecessor; in the +/// `enqueued` mode it returns on admission and the object is created in the +/// background. While [`MAX_SEALED_BATCHES`] batches +/// wait to become durable no append is admitted. pub(crate) struct ObjectStoreLogStore { prefix: String, + ack_mode: AckMode, io: Arc, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, @@ -109,6 +114,7 @@ impl fmt::Debug for ObjectStoreLogStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ObjectStoreLogStore") .field("prefix", &self.prefix) + .field("ack_mode", &self.ack_mode) .finish_non_exhaustive() } } @@ -145,6 +151,16 @@ impl ObjectStoreLogStore { ); let max_batch_bytes = positive_bytes(config.max_batch_bytes.as_bytes(), "max batch bytes")?; + let max_unpersisted_bytes = positive_bytes( + config.max_unpersisted_bytes.as_bytes(), + "max unpersisted bytes", + )?; + ensure!( + config.max_unpersisted_age > Duration::ZERO, + InvalidWalObjectStoreSnafu { + reason: "max unpersisted age is zero", + } + ); let Recovered { mut catalog, next_object_seq, @@ -196,12 +212,18 @@ impl ObjectStoreLogStore { stopped: stopped.clone(), command_rx, append_rx, + ack_mode: config.ack_mode, + max_unpersisted_bytes, + max_unpersisted_age: config.max_unpersisted_age, open_batch: OpenBatch::new(max_batch_bytes), issued_entry_ids: durable_entry_ids, pending: Vec::new(), sealed: VecDeque::new(), creates: FuturesUnordered::new(), + stalled: None, + durable_waiters: Vec::new(), stop: Vec::new(), + stop_error: None, next_object_seq: start .object_seq .checked_add(1) @@ -219,6 +241,7 @@ impl ObjectStoreLogStore { common_runtime::spawn_global(actor.run()); Ok(Arc::new(Self { prefix, + ack_mode: config.ack_mode, io, catalog, obsolete_entry_ids, @@ -235,6 +258,16 @@ impl ObjectStoreLogStore { })) } + /// Returns the largest entry id of the provider's region whose object is + /// durable and indexed, or zero for a region without such entries. In the + /// `enqueued` acknowledgement mode an entry id that an append returned + /// stays above this value until the object holding it is created. + pub(crate) fn durable_entry_id(&self, provider: &Provider) -> Result { + self.check_terminal()?; + let region_id = self.region_of(provider)?; + Ok(self.durable_entry_id_of(region_id)) + } + fn durable_entry_id_of(&self, region_id: RegionId) -> EntryId { let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); catalog.region_max_entry_id(region_id).unwrap_or(0) @@ -362,7 +395,8 @@ impl LogStore for ObjectStoreLogStore { type Error = Error; /// Stops the store. Creates in flight run to completion and acknowledge - /// their entries if they succeed. + /// their entries if they succeed; in the `enqueued` mode the remaining + /// backlog is uploaded first, and a failure of that upload is returned. async fn stop(&self) -> Result<()> { self.stopped.store(true, Ordering::Release); let (response_tx, response_rx) = oneshot::channel(); @@ -496,12 +530,17 @@ impl LogStore for ObjectStoreLogStore { .collect()) } - /// Moves the obsolete watermark of the region up to `entry_id` and makes - /// every id the region is assigned from now on greater than `entry_id`: - /// the sequence of the next object is raised above the object `entry_id` - /// names unless it is there already, so a watermark assigned under - /// another prefix, or one whose object this prefix no longer holds, is - /// never passed by a new id. Zero names no object and needs no floor. + /// Moves the obsolete watermark of the region up to `entry_id`. In the + /// `enqueued` acknowledgement mode the watermark never passes the durable + /// entry id: an entry that is not durable yet is replayed after a crash, + /// and hiding it would skip it. + /// + /// Every id the region is assigned from now on is greater than + /// `entry_id`, whatever watermark is recorded: the sequence of the next + /// object is raised above the object `entry_id` names unless it is there + /// already, so a watermark assigned under another prefix, or one whose + /// object this prefix no longer holds, is never passed by a new id. Zero + /// names no object and needs no floor. /// /// The watermark and the floor are applied together by the actor: a call /// cancelled before its command is queued publishes neither, while a @@ -509,7 +548,8 @@ impl LogStore for ObjectStoreLogStore { /// ids under the next sequence the floor cannot move, and the call fails /// with [`Error::WalObjectSequenceUnsettled`] and records neither. The /// watermark is kept in memory; the caller re-establishes it after a - /// restart. + /// restart. In the `enqueued` mode an `entry_id` this store handed out + /// needs no floor: it is never handed out again while the store runs. async fn obsolete( &self, provider: &Provider, @@ -517,10 +557,15 @@ impl LogStore for ObjectStoreLogStore { entry_id: EntryId, ) -> Result<()> { self.check_region(provider, region_id)?; + let watermark = match self.ack_mode { + AckMode::Durable => entry_id, + AckMode::Enqueued => entry_id.min(self.durable_entry_id_of(region_id)), + }; let (response_tx, response_rx) = oneshot::channel(); let command = Command::Obsolete { region_id, entry_id, + watermark, response: response_tx, }; let answered = match self.command_tx.send(command).await { @@ -533,7 +578,7 @@ impl LogStore for ObjectStoreLogStore { // still queued behind the stop: nothing is assigned an id any // more, so the watermark alone is consistent. None => { - record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + record_obsolete(&self.obsolete_entry_ids, region_id, watermark); Ok(()) } } @@ -564,9 +609,29 @@ impl LogStore for ObjectStoreLogStore { } fn latest_entry_id(&self, provider: &Provider) -> Result { + self.durable_entry_id(provider) + } + + /// Waits until every entry of the provider's region with an id at or + /// below `entry_id` is durable and indexed. Returns at once in the + /// `durable` acknowledgement mode, where a caller only holds durable ids. + async fn wait_durable(&self, provider: &Provider, entry_id: EntryId) -> Result<()> { self.check_terminal()?; let region_id = self.region_of(provider)?; - Ok(self.durable_entry_id_of(region_id)) + if entry_id <= self.durable_entry_id_of(region_id) { + return Ok(()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(Command::WaitDurable { + region_id, + entry_id, + response: response_tx, + }) + .await + .ok() + .context(ObjectStoreWalStoppedSnafu)?; + response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? } } @@ -575,12 +640,19 @@ type AppendResponse = oneshot::Sender>; type QueuedAppend = (Vec, AppendResponse); enum Command { + /// Answered once the region is durable through `entry_id`. + WaitDurable { + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + }, /// Answered once the watermark is recorded and no id of the region at /// or below `entry_id` can be assigned, or with the reason neither was /// done. Obsolete { region_id: RegionId, entry_id: EntryId, + watermark: EntryId, response: oneshot::Sender>, }, Stop { @@ -598,9 +670,17 @@ struct PendingAppend { response: AppendResponse, } +/// A caller waiting until a region is durable through an entry id. +struct DurableWaiter { + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, +} + /// Where the conditional create of a sealed batch stands. enum CreateState { - /// The create has not started because no slot was free. + /// The create has not started: no slot was free, or in the `enqueued` + /// mode an earlier attempt failed transiently and is repeated. Pending, InFlight, Created, @@ -612,6 +692,9 @@ struct SealedBatch { object_seq: u64, bytes: Bytes, footer: Vec, + first_admitted_at: Instant, + /// Number of creates that were attempted for the batch. + attempts: u32, waiters: Vec, #[cfg(any(test, feature = "testing"))] seal_waiters: Vec>>, @@ -619,6 +702,10 @@ struct SealedBatch { } impl SealedBatch { + fn is_in_flight(&self) -> bool { + matches!(self.state, CreateState::InFlight) + } + fn fail(self, error: impl Fn() -> Error) { for waiter in self.waiters { let _ = waiter.response.send(Err(error())); @@ -646,18 +733,22 @@ type CreateOutcome = (u64, Result); /// | Situation | Sequence | Waiters | Store | /// | --- | --- | --- | --- | /// | created, or identical retry | advances | acknowledged in order | healthy | -/// | transient error | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | -/// | an object of an earlier epoch holds the sequence | never reused | as for a transient error | healthy: that object can never be on the chain | -/// | an object of the same or a later epoch holds the sequence, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | +/// | transient error, `durable` mode | never reused: a create that reported an error may have written its object | every batch that is not indexed fails, created ones included; retries are assigned new ids | healthy: the next batch links to the last indexed object, so no chain reaches a failed batch | +/// | transient error, `enqueued` mode | unchanged | already acknowledged | healthy: the create is repeated with the same bytes after [`CREATE_RETRY_DELAY`], which an identical retry accepts | +/// | transient error, `enqueued` mode after `stop` began | never reused | already acknowledged | the backlog is dropped and `stop` reports the error | +/// | an object of an earlier epoch holds the sequence, `durable` mode | never reused | as for a transient error | healthy: that object can never be on the chain | +/// | an object of the same or a later epoch holds the sequence, any conflicting object in the `enqueued` mode, encoding or catalog error | unchanged | every batch that is not indexed fails | poisoned | /// | created at the last representable sequence | cannot advance | acknowledged | poisoned: no later batch can be allocated a sequence | /// /// A create still in flight when its batch failed runs to completion and its /// outcome is ignored: whatever it stored is off the chain. After `stop` began -/// nothing is admitted and no create starts; creates in flight run to -/// completion and acknowledge if they succeed. +/// nothing is admitted and no create starts, except that the `enqueued` mode +/// uploads its backlog; creates in flight run to completion and acknowledge +/// if they succeed. /// /// Appends arrive on their own channel, which the actor does not read while -/// [`MAX_SEALED_BATCHES`] batches wait, so commands such as `stop` are handled +/// [`MAX_SEALED_BATCHES`] batches wait or an append is held back at a backlog +/// threshold of the `enqueued` mode, so commands such as `stop` are handled /// while admission is held back. struct Actor { io: Arc, @@ -667,16 +758,29 @@ struct Actor { stopped: Arc, command_rx: mpsc::Receiver, append_rx: mpsc::Receiver, + ack_mode: AckMode, + max_unpersisted_bytes: usize, + max_unpersisted_age: Duration, open_batch: OpenBatch, + /// Largest entry id ever handed out per region, whether it became + /// durable or failed. Unlike the accepted ids of the open batch it never + /// moves down. issued_entry_ids: HashMap, - /// Waiters of the open batch. + /// Waiters of the open batch in the `durable` mode. pending: Vec, /// Batches that are not durable yet, in sequence order. sealed: VecDeque, /// Creates in flight, including those of batches that already failed. creates: FuturesUnordered>, + /// The append held back in the `enqueued` mode while the unpersisted + /// backlog is at a threshold. Later appends wait in their channel. + stalled: Option, + durable_waiters: Vec, /// Callers of `stop`, answered once nothing is in flight. stop: Vec>>, + /// The failure `stop` reports in the `enqueued` mode once an + /// acknowledged backlog was dropped, recorded when it happens. + stop_error: Option>, /// Sequence of the next sealed batch, `None` once the sequence is /// exhausted. The open batch assigns its entry ids from it. It is above /// every sealed batch and never moves back while the store runs. @@ -705,7 +809,8 @@ impl Actor { loop { tokio::select! { _ = interval.tick() => { - // Nothing starts after stop began. + // Nothing starts after stop began; the `enqueued` mode + // sealed its backlog when stop was requested. if !self.is_stopped() { self.flush_open_batch(); } @@ -713,12 +818,15 @@ impl Actor { Some((object_seq, result)) = self.creates.next(), if !self.creates.is_empty() => { self.on_create_completed(object_seq, result); } - Some((entries, response)) = self.append_rx.recv(), if self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { + Some((entries, response)) = self.append_rx.recv(), if self.stalled.is_none() && self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { self.handle_append(entries, response); } command = self.command_rx.recv() => match command { - Some(Command::Obsolete { region_id, entry_id, response }) => { - self.handle_obsolete(region_id, entry_id, response); + Some(Command::WaitDurable { region_id, entry_id, response }) => { + self.handle_wait_durable(region_id, entry_id, response); + } + Some(Command::Obsolete { region_id, entry_id, watermark, response }) => { + self.handle_obsolete(region_id, entry_id, watermark, response); } Some(Command::Stop { response }) => { self.handle_stop(response); @@ -749,11 +857,17 @@ impl Actor { let _ = response.send(Err(shared(&error))); return; } + if self.ack_mode == AckMode::Enqueued && self.backlog_at_threshold() { + self.stalled = Some((entries, response)); + self.ensure_create_in_flight(); + return; + } self.admit(entries, response); } /// Admits `entries` into the open batch, which assigns their ids under - /// the next object sequence. The caller waits for the object. + /// the next object sequence. In the `durable` mode the caller waits for + /// the object, in the `enqueued` mode it is answered now. fn admit(&mut self, entries: Vec, response: AppendResponse) { // A region would run past the position range of the open batch: the // batch is sealed and the entries open the next one. The size limit @@ -785,10 +899,21 @@ impl Actor { return; } }; - self.pending.push(PendingAppend { - last_entry_ids, - response, - }); + for (region_id, entry_id) in &last_entry_ids { + self.issued_entry_ids + .entry(*region_id) + .and_modify(|issued| *issued = (*issued).max(*entry_id)) + .or_insert(*entry_id); + } + match self.ack_mode { + AckMode::Durable => self.pending.push(PendingAppend { + last_entry_ids, + response, + }), + AckMode::Enqueued => { + let _ = response.send(Ok(AppendBatchResponse { last_entry_ids })); + } + } #[cfg(any(test, feature = "testing"))] self.admitted_appends.send_modify(|count| *count += 1); if self.open_batch.should_seal() { @@ -796,6 +921,59 @@ impl Actor { } } + /// Returns true once the unpersisted backlog, the open batch and every + /// sealed batch that is not durable, reaches the size or the age threshold. + fn backlog_at_threshold(&self) -> bool { + let bytes = self.open_batch.estimated_bytes() + + self + .sealed + .iter() + .map(|batch| batch.bytes.len()) + .sum::(); + if bytes >= self.max_unpersisted_bytes { + return true; + } + let oldest = self + .sealed + .front() + .map(|batch| batch.first_admitted_at) + .or_else(|| self.open_batch.first_admitted_at()); + oldest.is_some_and(|admitted_at| admitted_at.elapsed() >= self.max_unpersisted_age) + } + + /// Seals the open batch when no create is in flight, so that a stalled + /// append has an upload to wait for. + fn ensure_create_in_flight(&mut self) { + if !self.sealed.iter().any(SealedBatch::is_in_flight) { + self.flush_open_batch(); + } + } + + /// Admits the stalled append once the backlog is below the thresholds and + /// the sealed batches leave room for the two an admission can seal. + fn release_stalled(&mut self) { + if self.stalled.is_none() { + return; + } + if self.is_stopped() { + if let Some((_, response)) = self.stalled.take() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + } + return; + } + if self.backlog_at_threshold() { + // The append needs an upload to complete. + self.ensure_create_in_flight(); + return; + } + if self.sealed.len() + 2 > MAX_SEALED_BATCHES { + return; + } + if let Some((entries, response)) = self.stalled.take() { + self.admit(entries, response); + } + } + /// Seals the open batch as the object `next_object_seq` and starts its /// create when a slot is free. Returns whether a batch was sealed. fn flush_open_batch(&mut self) -> bool { @@ -816,7 +994,7 @@ impl Actor { return false; }; - let (entries, _) = self.open_batch.seal(); + let (entries, first_admitted_at) = self.open_batch.seal(); let header = Header { object_seq, epoch: self.epoch, @@ -854,6 +1032,8 @@ impl Actor { object_seq, bytes: encoded.bytes, footer: encoded.footer, + first_admitted_at, + attempts: 0, waiters: std::mem::take(&mut self.pending), #[cfg(any(test, feature = "testing"))] seal_waiters: Vec::new(), @@ -865,9 +1045,10 @@ impl Actor { /// Starts the creates of pending batches in sequence order while fewer /// than [`MAX_IN_FLIGHT_CREATES`] are in flight, counting the creates of - /// batches that already failed. Nothing starts once stop began. + /// batches that already failed. Nothing starts once stop began, except + /// the backlog of the `enqueued` mode. fn start_creates(&mut self) { - if self.is_stopped() { + if self.is_stopped() && self.ack_mode == AckMode::Durable { return; } let mut in_flight = self.creates.len(); @@ -880,15 +1061,25 @@ impl Actor { } batch.state = CreateState::InFlight; in_flight += 1; + let delay = if batch.attempts == 0 { + Duration::ZERO + } else { + CREATE_RETRY_DELAY + }; + batch.attempts += 1; let io = self.io.clone(); let object_seq = batch.object_seq; let bytes = batch.bytes.clone(); let epoch = self.epoch; + let ack_mode = self.ack_mode; #[cfg(any(test, feature = "testing"))] let mut creates_held = self.creates_held.clone(); #[cfg(any(test, feature = "testing"))] let creates_fail = self.creates_fail.clone(); self.creates.push(Box::pin(async move { + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } // The store was dropped while the create was parked: it never // runs. #[cfg(any(test, feature = "testing"))] @@ -908,8 +1099,12 @@ impl Actor { }); return (object_seq, result); } + // Any conflict poisons an `enqueued` store, so only the + // `durable` mode reads the epoch of the existing object. let result = match io.put_if_absent(object_seq, bytes).await { - Err(error @ Error::WalObjectConflict { .. }) => { + Err(error @ Error::WalObjectConflict { .. }) + if ack_mode == AckMode::Durable => + { stale_conflict(io.as_ref(), object_seq, epoch, error).await } result => result, @@ -932,10 +1127,28 @@ impl Actor { }; match result { Ok(_) => self.sealed[index].state = CreateState::Created, + // Nobody is left to retry in the `enqueued` mode, so the store + // repeats a create that failed transiently under the same sequence + // with the same bytes, which an identical retry accepts. + Err(ref error) + if self.ack_mode == AckMode::Enqueued + && !self.is_stopped() + && is_transient(error) => + { + self.sealed[index].state = CreateState::Pending; + } // The object store did not confirm the object, which may still // exist or land later, or an earlier epoch holds the sequence and - // can never be on the chain. The caller retries the append itself. - Err(error @ (Error::WalObjectStore { .. } | Error::StaleWalObject { .. })) => { + // can never be on the chain. A caller of the `durable` mode + // retries the append itself; after stop began a transient failure + // drops the `enqueued` backlog. + Err(error @ Error::WalObjectStore { .. }) + if self.ack_mode == AckMode::Durable + || (self.is_stopped() && is_transient(&error)) => + { + self.roll_back(Arc::new(error)) + } + Err(error @ Error::StaleWalObject { .. }) if self.ack_mode == AckMode::Durable => { self.roll_back(Arc::new(error)) } Err(error) => { @@ -955,6 +1168,7 @@ impl Actor { } } self.start_creates(); + self.release_stalled(); } /// Indexes the created object at the front and acknowledges its waiters. @@ -988,9 +1202,45 @@ impl Actor { for waiter in batch.seal_waiters { let _ = waiter.send(Ok(())); } + self.resolve_durable_waiters(); true } + fn resolve_durable_waiters(&mut self) { + let waiters = std::mem::take(&mut self.durable_waiters); + for waiter in waiters { + if self.is_durable_through(waiter.region_id, waiter.entry_id) { + let _ = waiter.response.send(Ok(())); + } else { + self.durable_waiters.push(waiter); + } + } + } + + /// Returns true when no id of the region at or below `entry_id` waits to + /// become durable. An id of a batch that failed will never be durable and + /// no longer waits. + fn is_durable_through(&self, region_id: RegionId, entry_id: EntryId) -> bool { + self.lowest_pending_entry_id(region_id) + .is_none_or(|pending| pending > entry_id) + } + + /// Returns the lowest id of the region that was handed out and is not + /// durable yet. The sealed batches are in sequence order and the open + /// batch is above all of them. + fn lowest_pending_entry_id(&self, region_id: RegionId) -> Option { + self.sealed + .iter() + .flat_map(|batch| batch.footer.iter()) + .find(|entry| entry.region_id == region_id) + .map(|entry| entry.min_entry_id) + .or_else(|| { + self.next_object_seq + .filter(|_| self.open_batch.holds_region(region_id)) + .map(|object_seq| entry_id(object_seq, 1)) + }) + } + /// Fails every batch that is not indexed after a create failed /// transiently, including batches already created and batches whose /// create is in flight. Their sequences are not reused: a create that @@ -1013,14 +1263,20 @@ impl Actor { } self.reset_open_batch(); self.fail_unacknowledged(failure); + // An acknowledged backlog was dropped: `stop` reports it, whether + // its caller has arrived yet or not. + if self.ack_mode == AckMode::Enqueued { + self.stop_error.get_or_insert(error); + } } /// Records `error` as terminal and fails every waiter that is not /// acknowledged with it, or with the stopped error if the store was /// stopped meanwhile. Creates in flight run to completion, but their - /// outcome is ignored: the entries of an object they create were never - /// acknowledged, like those of a crash between creation and - /// acknowledgement. Returns the recorded error. + /// outcome is ignored: in the `durable` mode the entries of an object they + /// create were never acknowledged, like those of a crash between creation + /// and acknowledgement; in the `enqueued` mode the acknowledged backlog is + /// discarded and `stop` reports it. Returns the recorded error. fn poison(&mut self, error: Error) -> Arc { let error = set_terminal(&self.terminal_error, error); let stopped = self.is_stopped(); @@ -1036,6 +1292,9 @@ impl Actor { } self.reset_open_batch(); self.fail_unacknowledged(failure); + if self.ack_mode == AckMode::Enqueued { + self.stop_error.get_or_insert(error.clone()); + } error } @@ -1045,11 +1304,56 @@ impl Actor { self.open_batch.reset(); } - /// Fails the waiters of the open batch. + /// Fails the waiters of the open batch, the stalled append and the + /// durability waiters. fn fail_unacknowledged(&mut self, error: impl Fn() -> Error) { for pending in self.pending.drain(..) { let _ = pending.response.send(Err(error())); } + if let Some((_, response)) = self.stalled.take() { + let _ = response.send(Err(error())); + } + for waiter in self.durable_waiters.drain(..) { + let _ = waiter.response.send(Err(error())); + } + } + + fn handle_wait_durable( + &mut self, + region_id: RegionId, + entry_id: EntryId, + response: oneshot::Sender>, + ) { + if let Some(error) = terminal(&self.terminal_error) { + let _ = response.send(Err(shared(&error))); + return; + } + let durable = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + catalog.region_max_entry_id(region_id).unwrap_or(0) + }; + if entry_id <= durable { + let _ = response.send(Ok(())); + return; + } + // An acknowledged backlog was dropped: no entry that is not durable + // can be certified any more, whether it was in that backlog or not. + if let Some(error) = &self.stop_error { + let _ = response.send(Err(shared(error))); + return; + } + if self.is_durable_through(region_id, entry_id) { + let _ = response.send(Ok(())); + return; + } + // A caller that stopped waiting leaves a closed response behind. + self.durable_waiters + .retain(|waiter| !waiter.response.is_closed()); + self.durable_waiters.push(DurableWaiter { + region_id, + entry_id, + response, + }); } /// Makes sure no id of the region at or below `entry_id` is assigned @@ -1059,10 +1363,11 @@ impl Actor { &mut self, region_id: RegionId, entry_id: EntryId, + watermark: EntryId, response: oneshot::Sender>, ) { - let result = self.raise_sequence_floor(entry_id).map(|_| { - record_obsolete(&self.obsolete_entry_ids, region_id, entry_id); + let result = self.raise_sequence_floor(region_id, entry_id).map(|_| { + record_obsolete(&self.obsolete_entry_ids, region_id, watermark); }); let _ = response.send(result); } @@ -1071,7 +1376,7 @@ impl Actor { /// `entry_id` unless it is there already. The move fails while the open /// batch has handed out ids under the next sequence. A floor that does /// not fit an entry id poisons the store like an exhausted sequence. - fn raise_sequence_floor(&mut self, entry_id: EntryId) -> Result<()> { + fn raise_sequence_floor(&mut self, region_id: RegionId, entry_id: EntryId) -> Result<()> { // Nothing is assigned an id after stop began. if self.is_stopped() { return Ok(()); @@ -1079,6 +1384,16 @@ impl Actor { if let Some(error) = terminal(&self.terminal_error) { return Err(shared(&error)); } + // In the `enqueued` mode an id the store handed out is never handed + // out again while it runs: a create that fails transiently is + // repeated under its sequence and a permanent failure poisons the + // store. Every later id of the region is greater, so the id needs no + // floor even before it is durable. + if self.ack_mode == AckMode::Enqueued + && entry_id <= self.issued_entry_ids.get(®ion_id).copied().unwrap_or(0) + { + return Ok(()); + } let sequence_floor = sequence_floor(entry_id); let Some(next_object_seq) = self.next_object_seq else { return Ok(()); @@ -1107,24 +1422,35 @@ impl Actor { } } - /// Begins stopping. Nothing is admitted from now on; the open batch and - /// the batches whose create has not started are dropped. Stop is answered by [`finish_stop`](Self::finish_stop) - /// once nothing is in flight. + /// Begins stopping. Nothing is admitted from now on; the `durable` mode + /// drops the open batch and the batches whose create has not started, + /// the `enqueued` mode seals its backlog so it is uploaded. Stop is + /// answered by [`finish_stop`](Self::finish_stop) once nothing is in flight. fn handle_stop(&mut self, response: oneshot::Sender>) { self.stop.push(response); - self.reset_open_batch(); - for pending in self.pending.drain(..) { - let _ = pending - .response - .send(Err(ObjectStoreWalStoppedSnafu.build())); - } - if let Some(index) = self - .sealed - .iter() - .position(|batch| matches!(batch.state, CreateState::Pending)) - { - for batch in self.sealed.drain(index..) { - batch.fail(|| ObjectStoreWalStoppedSnafu.build()); + match self.ack_mode { + AckMode::Durable => { + self.reset_open_batch(); + for pending in self.pending.drain(..) { + let _ = pending + .response + .send(Err(ObjectStoreWalStoppedSnafu.build())); + } + if let Some(index) = self + .sealed + .iter() + .position(|batch| matches!(batch.state, CreateState::Pending)) + { + for batch in self.sealed.drain(index..) { + batch.fail(|| ObjectStoreWalStoppedSnafu.build()); + } + } + } + AckMode::Enqueued => { + if let Some((_, response)) = self.stalled.take() { + let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); + } + self.flush_open_batch(); } } } @@ -1135,8 +1461,14 @@ impl Actor { if self.stop.is_empty() || !self.sealed.is_empty() || !self.creates.is_empty() { return false; } + for waiter in self.durable_waiters.drain(..) { + let _ = waiter + .response + .send(Err(ObjectStoreWalStoppedSnafu.build())); + } + let error = self.stop_error.take(); for response in self.stop.drain(..) { - let _ = response.send(Ok(())); + let _ = response.send(error.as_ref().map_or(Ok(()), |error| Err(shared(error)))); } true } @@ -1179,6 +1511,13 @@ fn set_terminal(terminal_error: &TerminalError, error: Error) -> Arc { .clone() } +/// Returns true for a storage error that a later attempt may not meet: one +/// the object store reports as temporary, or as persistent, which is how its +/// retry layer reports a temporary error that outlasted its retries. +fn is_transient(error: &Error) -> bool { + matches!(error, Error::WalObjectStore { error, .. } if !error.is_permanent()) +} + /// Wraps an error that several callers receive. fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) @@ -1623,6 +1962,14 @@ mod tests { } } + /// The same batching as `config`, acknowledging appends on admission. + fn enqueued(config: ObjectStoreWalConfig) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + ack_mode: AckMode::Enqueued, + ..config + } + } + /// Every append reaches the size limit, so it is persisted on its own. fn eager() -> ObjectStoreWalConfig { config(Duration::from_secs(3600), 1) @@ -1787,6 +2134,14 @@ mod tests { prefix: "/absolute".to_string(), ..config(Duration::from_secs(1), 1) }, + ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(0), + ..config(Duration::from_secs(1), 1) + }, + ObjectStoreWalConfig { + max_unpersisted_age: Duration::ZERO, + ..config(Duration::from_secs(1), 1) + }, ] { let error = ObjectStoreLogStore::try_new(memory_store(), &config, 1, 2) .await @@ -2393,6 +2748,7 @@ mod tests { let (append_tx, _) = mpsc::channel(APPEND_BUFFER); let store = ObjectStoreLogStore { prefix: PREFIX.to_string(), + ack_mode: AckMode::Durable, io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), catalog: Arc::default(), obsolete_entry_ids: ObsoleteEntryIds::default(), @@ -4268,6 +4624,561 @@ mod tests { ); } + #[tokio::test] + async fn test_store_enqueued_append_returns_before_the_object_exists() { + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + + let response = timeout(WAIT, append(&store, region_id, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + assert_eq!(0, latest(&store, region_id)); + assert_eq!(0, store.durable_entry_id(&provider(region_id)).unwrap()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 2))]), + response.last_entry_ids + ); + + store.seal_open_batch().await.unwrap(); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(id(1, 2), latest(&store, region_id)); + assert_eq!( + id(1, 2), + store.durable_entry_id(&provider(region_id)).unwrap() + ); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_durable_id_advances_after_create_and_indexing() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_one = region(1); + let region_two = region(2); + + let response = timeout(WAIT, append(&store, region_one, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_one, id(1, 1))]), + response.last_entry_ids + ); + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(1, object_seq); + assert_eq!(0, store.durable_entry_id(&provider(region_one)).unwrap()); + let wait = |entry_id| { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_one), entry_id).await }) + }; + // An id the store never handed out waits for the ids of the region + // that were handed out below it. + let waits = [wait(id(1, 1)), wait(id(1, 7))]; + // The other region has nothing to wait for. + timeout(WAIT, store.wait_durable(&provider(region_two), id(1, 1))) + .await + .unwrap() + .unwrap(); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(waits.iter().all(|wait| !wait.is_finished())); + + release.send(true).unwrap(); + for wait in waits { + timeout(WAIT, wait).await.unwrap().unwrap().unwrap(); + } + assert_eq!( + id(1, 1), + store.durable_entry_id(&provider(region_one)).unwrap() + ); + assert_eq!(0, store.durable_entry_id(&provider(region_two)).unwrap()); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + } + + #[tokio::test] + async fn test_store_enqueued_wait_ignores_later_pending_entries_of_the_region() { + let (store, _io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let region_a = region(1); + let region_b = region(2); + // Object 1 holds region A's id(1, 1) and object 2 region B's id(2, 1). + for (region_id, data) in [(region_a, "a1"), (region_b, "b1")] { + let response = append(&store, region_id, data).await.unwrap(); + let (_, release) = next_create(&mut parked).await; + release.send(true).unwrap(); + let entry_id = response.last_entry_ids[®ion_id]; + timeout(WAIT, store.wait_durable(&provider(region_id), entry_id)) + .await + .unwrap() + .unwrap(); + } + let response = append(&store, region_a, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_a, id(3, 1))]), + response.last_entry_ids + ); + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(3, object_seq); + + // Every entry of region A up to id(2, 1) is durable, so the wait does + // not depend on the parked create of object 3. + timeout(WAIT, store.wait_durable(&provider(region_a), id(2, 1))) + .await + .unwrap() + .unwrap(); + // A wait for id(3, 1) itself is answered only once object 3 is + // indexed; the round trip orders the check after the actor handled it. + let (response_tx, mut response_rx) = oneshot::channel(); + store + .command_tx + .send(Command::WaitDurable { + region_id: region_a, + entry_id: id(3, 1), + response: response_tx, + }) + .await + .unwrap(); + round_trip_actor(&store).await; + assert!(matches!( + response_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + release.send(true).unwrap(); + timeout(WAIT, response_rx).await.unwrap().unwrap().unwrap(); + } + + #[tokio::test] + async fn test_store_enqueued_obsolete_never_passes_the_durable_id() { + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + store.seal_open_batch().await.unwrap(); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap(); + assert_eq!( + Some(&id(1, 2)), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + } + + /// Appends one entry in the enqueued mode with `config`, whose create + /// parks, then appends a second one that the backlog threshold must + /// stall. Returns the store, the parked creates, the stalled append and + /// the release of the first create. + async fn stall_second_append( + config: ObjectStoreWalConfig, + ) -> ( + Arc, + Arc, + mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, + tokio::task::JoinHandle>, + oneshot::Sender, + ) { + let (store, io, mut parked) = open_parking_creates(memory_store(), &config).await; + let region_id = region(1); + let response = timeout(WAIT, append(&store, region_id, "a1")) + .await + .unwrap() + .unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + if config.max_unpersisted_age < Duration::from_secs(1) { + tokio::time::sleep(config.max_unpersisted_age * 2).await; + } + + let stalled = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); + // The stall seals the open batch so that an upload is in flight. + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(1, object_seq); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!stalled.is_finished()); + assert!(parked.try_recv().is_err()); + (store, io, parked, stalled, release) + } + + #[tokio::test] + async fn test_store_enqueued_backlog_bytes_stall_admission_until_an_upload_completes() { + let config = ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(1), + ..enqueued(manual()) + }; + let (store, io, mut parked, stalled, release) = stall_second_append(config).await; + let region_id = region(1); + // Two more appends queue behind the stalled one. + let third = spawn_append_batch(&store, vec![entry(&store, region_id, "a3")]); + let fourth = spawn_append_batch(&store, vec![entry(&store, region_id, "a4")]); + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!third.is_finished() && !fourth.is_finished()); + assert!(parked.try_recv().is_err()); + + // The upload releases the second append, whose entry reaches the + // threshold again; the stall seals it so that the next upload can + // release the third, and so on. + release.send(true).unwrap(); + let response = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(id(1, 1), latest(&store, region_id)); + for (append, object_seq) in [(third, 2), (fourth, 3)] { + let (parked_seq, release) = next_create(&mut parked).await; + assert_eq!(object_seq, parked_seq); + release.send(true).unwrap(); + let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(object_seq + 1, 1))]), + response.last_entry_ids + ); + } + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + assert_eq!( + expected_entries( + region_id, + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(3, 1), "a3")] + ), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_backlog_age_stalls_admission_until_an_upload_completes() { + let config = ObjectStoreWalConfig { + max_unpersisted_age: Duration::from_millis(50), + ..enqueued(manual()) + }; + let (store, io, _parked, stalled, release) = stall_second_append(config).await; + let region_id = region(1); + + release.send(true).unwrap(); + let response = timeout(WAIT, stalled).await.unwrap().unwrap().unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + } + + #[tokio::test] + async fn test_store_enqueued_stop_fails_a_stalled_append_and_uploads_the_backlog() { + let config = ObjectStoreWalConfig { + max_unpersisted_bytes: ReadableSize(1), + ..enqueued(manual()) + }; + let (store, io, mut parked, stalled, release) = stall_second_append(config).await; + + let stop = { + let store = store.clone(); + tokio::spawn(async move { store.stop().await }) + }; + assert_stopped(&timeout(WAIT, stalled).await.unwrap().unwrap().unwrap_err()); + release.send(true).unwrap(); + timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert!(parked.try_recv().is_err()); + } + + #[tokio::test] + async fn test_store_enqueued_repeats_a_create_that_failed_transiently() { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &enqueued(eager())).await; + let region_id = region(1); + + // The create stores the object but reports a failure; the repeat + // writes the same bytes under the same sequence. + io.fail_after_next_put.store(true, Ordering::Relaxed); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(1, 1))]), + response.last_entry_ids + ); + timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + + // An object store with a retry layer reports a temporary error that + // outlasted its retries as persistent; the create is repeated too. + io.fail_next_put_persistently(); + append(&store, region_id, "a2").await.unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(2, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + } + + #[tokio::test] + async fn test_store_enqueued_permanent_failure_poisons() { + // A conflict of any epoch poisons the store without reading the + // existing object, since the acknowledged entries cannot move to + // another sequence. So does a permanent storage error. + for foreign_epoch in [Some(0), Some(2), None] { + let object_store = memory_store(); + let (io, reads) = RecordingIo::over(object_store.clone()); + let store = open_over(io.clone(), &enqueued(eager())).await; + let region_id = region(1); + match foreign_epoch { + Some(epoch) => put_foreign(&object_store, 1, epoch).await, + None => io.fail_next_put_permanently(), + } + let assert_poisoned = |error: Error| { + let poisoned = match foreign_epoch { + Some(_) => matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + None => matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + }; + assert!(poisoned, "unexpected error: {error:?}"); + }; + + // The append was acknowledged; the failure surfaces afterwards. + let second = entry(&store, region_id, "a2"); + append(&store, region_id, "a1").await.unwrap(); + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert_poisoned(error); + assert_poisoned(store.append_batch(vec![second]).await.unwrap_err()); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + // The acknowledged entry was dropped, which stop reports. + assert_poisoned(store.stop().await.unwrap_err()); + assert!(reads.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn test_store_enqueued_stop_reports_a_backlog_lost_before_the_stop_command() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let seal = { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + let (object_seq, release) = next_create(&mut parked).await; + assert_eq!(1, object_seq); + + // Stop began, but the actor has not received the stop command when + // the create fails: the backlog is dropped and the failure is kept + // for the stop that follows. + store.begin_stop(); + release.send(false).unwrap(); + assert_stopped(&timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err()); + // The lost entry cannot be certified as durable before the stop + // command is handled; a durable id still is. + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + timeout(WAIT, store.wait_durable(&provider(region_id), 0)) + .await + .unwrap() + .unwrap(); + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert!(parked.try_recv().is_err()); + } + + #[tokio::test] + async fn test_store_enqueued_stop_uploads_the_backlog() { + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + + store.stop().await.unwrap(); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(id(1, 2), latest(&store, region_id)); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(1, 2), "a2")]), + read_entries(&store, region_id, 1).await + ); + + // A backlog that cannot be uploaded is reported by stop. A transient + // failure, temporary or persistent, drops it; a permanent one also + // poisons the store. + for failure in ["temporary", "persistent", "permanent"] { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &enqueued(manual())).await; + append(&store, region_id, "a1").await.unwrap(); + match failure { + "temporary" => store.fail_creates(), + "persistent" => io.fail_next_put_persistently(), + _ => io.fail_next_put_permanently(), + } + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_eq!( + failure == "permanent", + store.latest_entry_id(&provider(region_id)).is_err() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_enqueued_issued_id_needs_no_floor_before_it_is_durable() { + let store = open(memory_store(), &enqueued(manual())).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + + // The open batch holds id(1, 1) under the next sequence. The id was + // handed out, so it is accepted without a floor, with the watermark + // capped to what is durable; an id the store never handed out under + // that sequence is refused. + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + let error = store + .obsolete(&provider(region_id), region_id, id(1, 2)) + .await + .unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectSequenceUnsettled { object_seq: 1, .. } + ), + "unexpected error: {error:?}" + ); + + store.seal_open_batch().await.unwrap(); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + assert!(read_entries(&store, region_id, 0).await.is_empty()); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + } + + #[tokio::test] + async fn test_store_enqueued_inherited_watermark_raises_the_floor() { + let object_store = memory_store(); + let store = open(object_store.clone(), &enqueued(manual())).await; + let region_id = region(1); + + // The region has nothing durable here, so the recorded watermark + // stays at zero, but its ids must still start above the watermark. + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + assert_eq!( + Some(&0), + store.obsolete_entry_ids.lock().unwrap().get(®ion_id) + ); + let response = append(&store, region_id, "r1").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(6, 1))]), + response.last_entry_ids + ); + store.seal_open_batch().await.unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(6, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(id(6, 1), latest(&store, region_id)); + store.stop().await.unwrap(); + + // Replay from the watermark sees the entry. + let store = open(object_store, &enqueued(manual())).await; + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "r1")]), + read_entries(&store, region_id, id(5, 1) + 1).await + ); + } + + #[tokio::test] + async fn test_store_durability_wait_pending_at_stop_fails_with_stopped() { + let store = open(memory_store(), &manual()).await; + let region_id = region(1); + let append = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); + store.wait_for_admitted_appends(1).await.unwrap(); + let wait = { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_id), id(1, 1)).await }) + }; + for _ in 0..16 { + tokio::task::yield_now().await; + } + assert!(!wait.is_finished()); + + store.stop().await.unwrap(); + assert_stopped(&timeout(WAIT, append).await.unwrap().unwrap().unwrap_err()); + assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); + } + /// Object access whose next create stores the object but reports a /// transient failure, as a create whose response is lost. struct LostResponseIo { @@ -4446,11 +5357,13 @@ mod tests { /// Object access that records every range read as (sequence, offset, /// length) and, on request, reports the next conditional create as failed - /// after it wrote the object. + /// after it wrote the object, or fails it with a given error before it + /// writes. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, fail_after_next_put: AtomicBool, + fail_next_put: Mutex>, } impl RecordingIo { @@ -4460,14 +5373,38 @@ mod tests { inner: ObjectStoreIo::new(object_store, PREFIX).unwrap(), reads: reads.clone(), fail_after_next_put: AtomicBool::new(false), + fail_next_put: Mutex::new(None), }; (Arc::new(io), reads) } + + fn fail_next_put_permanently(&self) { + let error = object_store::Error::new( + object_store::ErrorKind::PermissionDenied, + "injected failure", + ); + *self.fail_next_put.lock().unwrap() = Some(error); + } + + /// Fails the next create as a retry layer reports a temporary error + /// that outlasted its retries. + fn fail_next_put_persistently(&self) { + let error = object_store::Error::new(object_store::ErrorKind::Unexpected, "injected") + .set_temporary() + .set_persistent(); + *self.fail_next_put.lock().unwrap() = Some(error); + } } #[async_trait::async_trait] impl WalObjectIo for RecordingIo { async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result { + if let Some(error) = self.fail_next_put.lock().unwrap().take() { + return Err(error).context(WalObjectStoreSnafu { + operation: "write", + path: self.object_path(object_seq), + }); + } let result = self.inner.put_if_absent(object_seq, content).await?; if self.fail_after_next_put.swap(false, Ordering::Relaxed) { return injected_failure("write", self.object_path(object_seq)); diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index a2c1e3ee803..6084f207641 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -105,6 +105,24 @@ pub trait LogStore: Send + Sync + 'static + std::fmt::Debug { /// Returns the latest entry id in the log store. fn latest_entry_id(&self, provider: &Provider) -> Result; + + /// Waits until every entry of the provider's region with an id at or below + /// `entry_id` has the persistence this log store guarantees for an + /// acknowledged append, so that a caller can persist `entry_id` as a + /// replay watermark. + /// + /// The default returns at once and keeps each log store's own + /// acknowledgement guarantee, whatever it is: Raft Engine with + /// `sync_write = false`, for example, acknowledges an append before its + /// periodic sync. A log store that acknowledges an append before its + /// entries have that persistence overrides this method. + async fn wait_durable( + &self, + _provider: &Provider, + _entry_id: EntryId, + ) -> Result<(), Self::Error> { + Ok(()) + } } /// The response of an `append` operation. From 28e415a10344bb539e7b778cf6dbf9cc95bfa995 Mon Sep 17 00:00:00 2001 From: discord9 Date: Mon, 28 Sep 2026 09:38:12 +0000 Subject: [PATCH 059/133] feat(query): implement PromQL @ modifier on vector and matrix selectors (#9224) * feat(query): implement PromQL @ modifier on vector and matrix selectors The planner previously dropped the `at` field of selectors (`at: _`), so `some_metric @ 300` silently returned step-following values instead of the samples anchored at the fixed timestamp. Anchoring follows Prometheus's `setOffsetForAtModifier` + `refetch` semantics: resolve the anchor (`@ `, `@ start()`, `@ end()`), apply `offset` to the anchor, rewrite the selector offset to `eval_start - anchor`, scan only the anchored window, then report the same window at every evaluation step via a grid-wide-replay `InstantManipulate`. `@ start()` / `@ end()` resolve against the statement's evaluation range (`stmt_start` / `stmt_end`), which subquery planning does not rewrite. Timestamps before the Unix epoch are accepted; unrepresentable ones are rejected with `AtModifierTimestampOutOfRange` instead of wrapping. Selectors without `@` are planned exactly as before. Report: .e-agent/greptimedb_promql_compatibility_report_2026-09-16.md P0-2 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): center predict_linear on the evaluation timestamp predict_linear_impl used the window's last sample time as the evaluation timestamp, and the UDF took only (ts_range, value_range, t) with no channel for the current instant. After a range selector is folded for @ (or with an offset) the same window is replayed at every step, so the prediction stayed constant at the anchor's answer; even a plain window ended before the step produced a stale value. Give prom_predict_linear a 4th argument carrying the step's evaluation instant (ms timestamp), derived in the planner from the row's time index plus the offset the window was folded with (at_offset for @, offset_ms otherwise, recorded on PromPlannerContext). The regression is centered on that instant, matching Prometheus' use of enh.Ts. The mixed float/native-histogram path forwards the same argument. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): tighten @ modifier planner and predict_linear - range_fold_offset is always a concrete millisecond offset, not an Option; the single-step replay helper no longer wraps an infallible plan in Result. - predict_linear's eval timestamp is always cast to Timestamp(ms) at the call site, so the UDF drops its dead Int64 branch and the extra func_name argument. - Drop two redundant comparison cases from the at_modifier sqlness test and fix the lookback window comment to the half-open (0s, 300s]. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): accept parse-time rejection of @ on Windows `@ 1e16` is 10^19 milliseconds, beyond i64::MAX. A Unix SystemTime holds it and the planner rejects the anchor it cannot represent, but a Windows SystemTime tops out near 1.8e12 seconds, so the parser's checked_add fails first and the same literal is rejected while parsing. Accept either rejection path so the test passes on both platforms. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): avoid replaying label_join over rewritten series keys Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): narrow anchored range call promotion Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): address @ modifier review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): satisfy super import format check Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(query): address at modifier review follow-ups Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/benches/bench_range_fn.rs | 4 +- src/promql/src/functions/predict_linear.rs | 156 +++- src/query/src/promql/error.rs | 8 + src/query/src/promql/planner.rs | 408 ++++++++-- src/query/src/promql/planner/at_modifier.rs | 326 ++++++++ src/query/src/promql/planner/test.rs | 675 +++++++++++++++- src/query/src/promql/planner/test/delta.rs | 2 +- .../common/promql/at_modifier.result | 743 ++++++++++++++++++ .../standalone/common/promql/at_modifier.sql | 301 +++++++ 9 files changed, 2529 insertions(+), 94 deletions(-) create mode 100644 src/query/src/promql/planner/at_modifier.rs create mode 100644 tests/cases/standalone/common/promql/at_modifier.result create mode 100644 tests/cases/standalone/common/promql/at_modifier.sql diff --git a/src/promql/benches/bench_range_fn.rs b/src/promql/benches/bench_range_fn.rs index 0f0a4867665..b4e026bae3d 100644 --- a/src/promql/benches/bench_range_fn.rs +++ b/src/promql/benches/bench_range_fn.rs @@ -270,7 +270,7 @@ fn make_quantile_input(num_points: usize, window_size: u32) -> Vec Vec { - let (ts_range, val_range, _) = build_sliding_ranges( + let (ts_range, val_range, eval_ts) = build_sliding_ranges( num_points, window_size, 1, @@ -282,6 +282,8 @@ fn make_predict_linear_input(num_points: usize, window_size: u32) -> Vec Result { error::ensure( - input.len() == 3, - DataFusionError::Plan("prom_predict_linear function should have 3 inputs".to_string()), + input.len() == 4, + DataFusionError::Plan("prom_predict_linear function should have 4 inputs".to_string()), )?; let t_col = &input[2]; + let eval_ts_col = &input[3]; let ts_range = extract_range_array(&input[0])?; let value_range = extract_range_array(&input[1])?; @@ -128,6 +131,51 @@ impl PredictLinear { Box::new(t_array.iter()) } }; + // The evaluation instant the window above was folded for. It is the instant the regression + // is centered on, which is *not* the last sample's timestamp: a window may end before the + // step being evaluated, and an `@`-anchored window is replayed at every step while keeping + // its own end. See `PromPlanner::create_range_eval_ts_expr`. + let eval_ts_iter: Box>> = match eval_ts_col { + ColumnarValue::Scalar(eval_ts_scalar) => { + let eval_ts = match eval_ts_scalar { + ScalarValue::TimestampMillisecond(Some(eval_ts), _) => *eval_ts, + // For a NULL or otherwise unusable evaluation timestamp, returns NULL array, + // which conforms to PromQL's behavior. + _ => { + let null_array = Float64Array::new_null(ts_range.len()); + return Ok(ColumnarValue::Array(Arc::new(null_array))); + } + }; + Box::new((0..ts_range.len()).map(move |_| Some(eval_ts))) + } + ColumnarValue::Array(eval_ts_array) => { + error::ensure( + eval_ts_array.len() == ts_range.len(), + DataFusionError::Execution(format!( + "{}: evaluation timestamp array should have the same length as other columns, found {} and {}", + Self::name(), + eval_ts_array.len(), + ts_range.len() + )), + )?; + match eval_ts_array.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, _) => Box::new( + eval_ts_array + .as_any() + .downcast_ref::() + .expect("checked by data type") + .iter(), + ), + other => { + return Err(DataFusionError::Execution(format!( + "{}: expect TimestampMillisecond as evaluation timestamp array's type, found {}", + Self::name(), + other + ))); + } + } + } + }; let all_timestamps = ts_range .values() .as_any() @@ -140,7 +188,12 @@ impl PredictLinear { .downcast_ref::() .unwrap(); let mut result_builder = Float64Builder::with_capacity(ts_range.len()); - for (index, t) in t_iter.enumerate() { + for (index, (t, eval_ts)) in t_iter.zip(eval_ts_iter).enumerate() { + // A step without an evaluation instant has nothing to extrapolate to. + let Some(eval_ts) = eval_ts else { + result_builder.append_null(); + continue; + }; match predict_linear_impl( &ts_range, &value_range, @@ -148,7 +201,7 @@ impl PredictLinear { all_values, index, t.unwrap(), - Self::name(), + eval_ts, )? { Some(value) => result_builder.append_value(value), None => result_builder.append_null(), @@ -167,7 +220,7 @@ fn predict_linear_impl( all_values: &Float64Array, index: usize, t: i64, - func_name: &str, + eval_ts: i64, ) -> Result, DataFusionError> { let (ts_offset, ts_len) = ts_range.get_offset_length(index).unwrap(); let (value_offset, value_len) = value_range.get_offset_length(index).unwrap(); @@ -175,22 +228,28 @@ fn predict_linear_impl( ts_len == value_len, DataFusionError::Execution(format!( "{}: time and value arrays in a group should have the same length, found {} and {}", - func_name, ts_len, value_len + PredictLinear::name(), + ts_len, + value_len )), )?; if ts_len < 2 { return Ok(None); } - // last timestamp is evaluation timestamp - let evaluate_ts = all_timestamps[ts_offset + ts_len - 1]; + // Like Prometheus, the regression is centered on the evaluation timestamp: the returned + // intercept is the value the window's trend predicts at the step being evaluated, and `t` is + // the horizon (in seconds) that is added on top of it. Centering on the last sample instead + // would drop the distance between the window's end and the step, which is non-zero both + // without `@` (the window can end before the step) and with it (the same anchored window is + // replayed at every step). let (slope, intercept) = linear_regression_slices( all_timestamps, ts_offset, all_values, value_offset, value_len, - evaluate_ts, + eval_ts, ); if slope.is_none() || intercept.is_none() { @@ -231,6 +290,11 @@ mod test { (ts_range_array, value_range_array) } + /// The evaluation timestamp at which the window built by [`build_test_range_arrays`] ends. + fn window_end_eval_ts() -> ScalarValue { + ScalarValue::TimestampMillisecond(Some(3000), None) + } + #[test] fn calculate_predict_linear_none() { let ts_array = Arc::new(TimestampMillisecondArray::from_iter( @@ -244,7 +308,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], vec![None, None], ); } @@ -256,7 +320,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], // value at t = 0 vec![Some(38.63636363636364)], ); @@ -269,7 +333,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(3000))], + vec![ScalarValue::Int64(Some(3000)), window_end_eval_ts()], // value at t = 3000 vec![Some(31856.818181818187)], ); @@ -282,7 +346,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(4200))], + vec![ScalarValue::Int64(Some(4200)), window_end_eval_ts()], // value at t = 4200 vec![Some(44584.09090909091)], ); @@ -295,7 +359,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(6600))], + vec![ScalarValue::Int64(Some(6600)), window_end_eval_ts()], // value at t = 6600 vec![Some(70038.63636363638)], ); @@ -308,7 +372,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(7800))], + vec![ScalarValue::Int64(Some(7800)), window_end_eval_ts()], // value at t = 7800 vec![Some(82765.9090909091)], ); @@ -327,7 +391,7 @@ mod test { PredictLinear::scalar_udf(), ts_array, value_array, - vec![ScalarValue::Int64(Some(0))], + vec![ScalarValue::Int64(Some(0)), window_end_eval_ts()], vec![Some(30.0)], ); } @@ -348,9 +412,69 @@ mod test { ColumnarValue::Array(Arc::new(ts_dict)), ColumnarValue::Array(Arc::new(value_dict)), ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Scalar(window_end_eval_ts()), ]) .unwrap_err(); assert!(err.to_string().contains("Empty range is not expected")); } + + /// The regression is centered on the evaluation timestamp, so the returned intercept follows + /// the step being evaluated even though the window stays the same. + #[test] + fn calculate_predict_linear_centers_on_eval_ts() { + // One sample per second with values 0, 1, 2: the trend is exactly one unit per second. + for (eval_ts, t, expected) in [ + // The window ends at 2s: evaluating there returns its last value. + (2000, 0, 2.0), + // One second past the window's end the trend is at 3. + (3000, 0, 3.0), + // One second past the window's end, predicting another two seconds ahead. + (3000, 2, 5.0), + ] { + let ts_values = Arc::new(TimestampMillisecondArray::from_iter( + [0i64, 1000, 2000].into_iter().map(Some), + )); + let value_values = Arc::new(Float64Array::from_iter([0.0, 1.0, 2.0])); + let ts_array = RangeArray::from_ranges(ts_values, [(0, 3)]).unwrap(); + let value_array = RangeArray::from_ranges(value_values, [(0, 3)]).unwrap(); + + simple_range_udf_runner( + PredictLinear::scalar_udf(), + ts_array, + value_array, + vec![ + ScalarValue::Int64(Some(t)), + ScalarValue::TimestampMillisecond(Some(eval_ts), None), + ], + vec![Some(expected)], + ); + } + } + + /// The evaluation timestamp arrives as a column, so the UDF accepts it as an array too. + #[test] + fn calculate_predict_linear_accepts_eval_ts_array() { + let (ts_array, value_array) = build_test_range_arrays(); + let eval_ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([3000])); + + let result = PredictLinear::predict_linear(&[ + ColumnarValue::Array(Arc::new(ts_array.into_dict())), + ColumnarValue::Array(Arc::new(value_array.into_dict())), + ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Array(eval_ts_array), + ]) + .unwrap(); + + let result = result + .into_array(1) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(result.len(), 1); + // The same value `calculate_predict_linear_test1` gets from a scalar evaluation timestamp. + assert!((result.value(0) - 38.63636363636364).abs() < 0.0001); + } } diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index 45e08349623..6cd45d9e57c 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -195,6 +195,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Timestamp out of range for the `@` modifier: {}", timestamp))] + AtModifierTimestampOutOfRange { + timestamp: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Timestamp out of range: {} of {:?}", timestamp, unit))] TimestampOutOfRange { timestamp: i64, @@ -252,6 +259,7 @@ impl ErrorExt for Error { | SameLabelSet { .. } | TimestampOutOfRange { .. } | SystemTimeOutOfRange { .. } + | AtModifierTimestampOutOfRange { .. } | InvalidRegularExpression { .. } | InvalidDestinationLabelName { .. } => StatusCode::InvalidArguments, diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 08e7005f9fb..be0c1364489 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -12,10 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod at_modifier; mod function_plans; - mod island; - mod matching_filters; mod set_operator; @@ -165,6 +164,12 @@ struct PromPlannerContext { end: Millisecond, interval: Millisecond, lookback_delta: Millisecond, + /// Evaluation range of the whole statement, which `@ start()` and `@ end()` refer to. + /// + /// Unlike [`Self::start`] and [`Self::end`], these are never rewritten while planning, so a + /// selector inside a subquery still resolves `@ start()` / `@ end()` against the statement. + stmt_start: Millisecond, + stmt_end: Millisecond, // planner states table_name: Option, @@ -193,6 +198,16 @@ struct PromPlannerContext { schema_name: Option, /// The range in millisecond of range selector. None if there is no range selector. range: Option, + /// The offset in milliseconds the window of the last planned range selector is folded with, + /// or `None` when no range selector has been planned since the last read. + /// + /// [`Self::start`] and the sample timestamps are compared on the shifted timeline: a range + /// payload carries `sample_timestamp + offset`, while the time index column of a folded row + /// stays the evaluation timestamp of its step. A function that reads it right after its input + /// plan is built, like `predict_linear`, consumes it instead of reading the state, so the + /// offset cannot leak from one input to another; see + /// [`PromPlanner::create_function_expr`]. + range_fold_offset: Option, } /// Result labels a vector-vector binary operation derives from its matching modifier, projected @@ -226,11 +241,15 @@ impl BinaryResultLabels { impl PromPlannerContext { fn from_eval_stmt(stmt: &EvalStmt) -> Self { + let start = stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as Millisecond; + let end = stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as Millisecond; Self { - start: stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as _, - end: stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as _, + start, + end, interval: stmt.interval.as_millis() as _, lookback_delta: stmt.lookback_delta.as_millis() as _, + stmt_start: start, + stmt_end: end, ..Default::default() } } @@ -246,6 +265,7 @@ impl PromPlannerContext { self.selector_matcher.clear(); self.schema_name = None; self.range = None; + self.range_fold_offset = None; } /// Reset table name and schema to empty @@ -325,6 +345,16 @@ impl PromPlanner { timestamp_fn: bool, query_engine_state: &QueryEngineState, ) -> Result { + // An anchored range call is step-invariant: evaluate it once, at the start of the + // evaluation, and report its result at every step; see + // [`Self::promote_anchored_range_call`]. + if let Some(plan) = self + .promote_anchored_range_call(prom_expr, timestamp_fn, query_engine_state) + .await? + { + return Ok(plan); + } + let res = match prom_expr { PromExpr::Aggregate(expr) => { self.prom_aggr_expr_to_plan(query_engine_state, expr) @@ -458,6 +488,9 @@ impl PromPlanner { divide_plan, ) .context(DataFusionPlanningSnafu)?; + // A subquery always folds with offset 0, so its payload timestamps are already on the + // evaluation timeline a function above it reads; see [`Self::create_range_eval_ts_expr`]. + self.ctx.range_fold_offset = Some(0); Ok(LogicalPlan::Extension(Extension { node: Arc::new(manipulate), @@ -1579,6 +1612,44 @@ impl PromPlanner { Ok(plan) } + /// The offset of a selector in milliseconds. A positive offset selects samples from an earlier + /// time and moves them forward into the evaluation timeline. + fn offset_millis(offset: &Option) -> Millisecond { + match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + } + } + + /// The columns that identify one series, which is the series key expected by the PromQL plan + /// nodes that hold exactly one series per input batch. + fn series_key_columns(&self) -> Vec { + if self.ctx.use_tsid { + vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] + } else { + self.ctx.tag_columns.clone() + } + } + + /// Keep replay and series division on the same effective keys when a call rewrites labels. + fn series_key_columns_for_schema(&self, schema: &DFSchemaRef) -> Vec { + let has_tsid = schema.fields().iter().any(|field| { + field.name() == DATA_SCHEMA_TSID_COLUMN_NAME + && field.data_type() == &ArrowDataType::UInt64 + }); + if has_tsid { + vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] + } else { + self.ctx + .tag_columns + .iter() + .filter(|name| schema.has_column_with_unqualified_name(name)) + .cloned() + .collect() + } + } + async fn prom_vector_selector_to_plan( &mut self, vector_selector: &VectorSelector, @@ -1588,20 +1659,37 @@ impl PromPlanner { name, offset, matchers, - at: _, + at, } = vector_selector; let matchers = self.preprocess_label_matchers(matchers, name)?; if let Some(empty_plan) = self.setup_context().await? { return Ok(empty_plan); } - let offset_ms = match offset { - Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, - Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), - None => 0, + let offset_ms = Self::offset_millis(offset); + // `@` anchors the sample window at a fixed timestamp: the selector selects its samples + // around the anchor once, instead of following the outer evaluation grid. See + // [`Self::at_modifier_offset`]. + let at_offset = self.at_modifier_offset(at, offset)?; + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; + let normalize = match at_offset { + Some(at_offset) => { + // Select the anchored samples at the start of the evaluation, with the offset that + // re-anchors the selector. + // The planner is single-use (one `EvalStmt` produces one plan), so an error + // below aborts the whole planning and `ctx.end` needs no restore-on-error. + self.ctx.end = grid_start; + let plan = self + .selector_to_series_normalize_plan(at_offset, matchers, false) + .await?; + self.ctx.end = grid_end; + plan + } + None => { + self.selector_to_series_normalize_plan(offset_ms, matchers, false) + .await? + } }; - let normalize = self - .selector_to_series_normalize_plan(offset, matchers, false) - .await?; let time_index_column = self.ctx .time_index_column @@ -1683,24 +1771,46 @@ impl PromPlanner { }; let field_column = self.ctx.field_columns.first().cloned(); - let manipulate = InstantManipulate::new( - self.ctx.start, - self.ctx.end, - self.ctx.lookback_delta, - self.ctx.interval, - offset_ms, - time_index_column, - if self.ctx.use_tsid { - vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] - } else { - self.ctx.tag_columns.clone() - }, - field_column, - normalize, - ); - let manipulate = LogicalPlan::Extension(Extension { - node: Arc::new(manipulate), - }); + let series_key_columns = self.series_key_columns(); + let manipulate = match at_offset { + Some(at_offset) => { + // Select the anchored sample once, then report it at every step of the outer + // grid. The samples keep their native anchor-time timestamps, so the manipulate + // must shift them onto the evaluation timeline with the rewritten offset. + let anchored = InstantManipulate::new( + grid_start, + grid_start, + self.ctx.lookback_delta, + self.ctx.interval, + at_offset, + time_index_column.clone(), + series_key_columns, + field_column, + normalize, + ); + self.replay_over_grid( + LogicalPlan::Extension(Extension { + node: Arc::new(anchored), + }), + grid_start, + grid_end, + time_index_column, + ) + } + None => LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + grid_start, + grid_end, + self.ctx.lookback_delta, + self.ctx.interval, + offset_ms, + time_index_column, + series_key_columns, + field_column, + normalize, + )), + }), + }; if let Some(timestamp_value_column) = timestamp_value_column { self.create_timestamp_func_plan(manipulate, ×tamp_value_column) } else { @@ -1758,46 +1868,108 @@ impl PromPlanner { name, offset, matchers, - .. + at, } = vs; let matchers = self.preprocess_label_matchers(matchers, name)?; ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); let range_ms = range.as_millis() as _; self.ctx.range = Some(range_ms); - let offset_ms = match offset { - Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, - Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), - None => 0, - }; + let offset_ms = Self::offset_millis(offset); + + // `@` anchors the range selector's window at a fixed timestamp, so the same window is fed + // to the enclosing function at every evaluation step. See [`Self::at_modifier_offset`]. + let at_offset = self.at_modifier_offset(at, offset)?; + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; // Some functions like rate may require special fields in the RangeManipulate plan // so we can't skip RangeManipulate. - let normalize = match self.setup_context().await? { - Some(empty_plan) => empty_plan, + let (normalize, at_offset) = match self.setup_context().await? { + // An empty metric does not contain any sample, so anchoring cannot change the result. + // The manipulate below folds the empty input with `offset_ms`, and the recorded fold + // offset agrees with it instead of with the anchor the window cannot use. + Some(empty_plan) => { + self.ctx.range_fold_offset = Some(offset_ms); + (empty_plan, None) + } None => { - self.selector_to_series_normalize_plan(offset, matchers, true) - .await? + let normalize = match at_offset { + Some(at_offset) => { + // Fold the anchored window once, at the start of the evaluation. + // Single-use planner: an error below aborts planning, so `ctx.end` + // needs no restore-on-error. + self.ctx.end = grid_start; + let plan = self + .selector_to_series_normalize_plan(at_offset, matchers, true) + .await?; + self.ctx.end = grid_end; + plan + } + None => { + self.selector_to_series_normalize_plan(offset_ms, matchers, true) + .await? + } + }; + // Samples are shifted onto the evaluation timeline with the very same offset while + // the window is folded. Record it so that the function above the selector can + // recover the evaluation instant of the folded window; see + // [`Self::create_range_eval_ts_expr`]. + self.ctx.range_fold_offset = Some(at_offset.unwrap_or(offset_ms)); + (normalize, at_offset) } }; - let manipulate = RangeManipulate::new( - self.ctx.start, - self.ctx.end, - self.ctx.interval, - offset_ms, - // TODO(ruihang): convert via Timestamp datatypes to support different time units - range_ms, - self.ctx - .time_index_column - .clone() - .expect("time index should be set in `setup_context`"), - self.ctx.field_columns.clone(), - normalize, - ) - .context(DataFusionPlanningSnafu)?; + let time_index_column = self + .ctx + .time_index_column + .clone() + .expect("time index should be set in `setup_context`"); + let manipulate = match at_offset { + Some(at_offset) => { + // Fold the anchored window once, then report it at every step of the outer + // grid. The samples keep their native anchor-time timestamps, so the manipulate + // must shift them onto the evaluation timeline with the rewritten offset. + let anchored = RangeManipulate::new( + grid_start, + grid_start, + self.ctx.interval, + at_offset, + // TODO(ruihang): convert via Timestamp datatypes to support different time units + range_ms, + time_index_column.clone(), + self.ctx.field_columns.clone(), + normalize, + ) + .context(DataFusionPlanningSnafu)?; + self.replay_over_grid( + LogicalPlan::Extension(Extension { + node: Arc::new(anchored), + }), + grid_start, + grid_end, + time_index_column, + ) + } + None => { + let manipulate = RangeManipulate::new( + grid_start, + grid_end, + self.ctx.interval, + offset_ms, + // TODO(ruihang): convert via Timestamp datatypes to support different time units + range_ms, + time_index_column, + self.ctx.field_columns.clone(), + normalize, + ) + .context(DataFusionPlanningSnafu)?; - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(manipulate), - })) + LogicalPlan::Extension(Extension { + node: Arc::new(manipulate), + }) + } + }; + + Ok(manipulate) } async fn prom_call_expr_to_plan( @@ -1846,11 +2018,17 @@ impl PromPlanner { ), }) }; + // The input plan records the fold offset of the range selector it is built from. Take it + // here, so that the offset of one input cannot leak into another call, and pass it to + // `create_function_expr`: the function that reads it (`predict_linear`) then depends on an + // argument instead of on planner state written by the selector below it. + let range_fold_offset = self.ctx.range_fold_offset.take(); let (mut func_exprs, new_tags) = self.create_function_expr( func, args.literals.clone(), input.schema(), query_engine_state, + range_fold_offset, )?; func_exprs.insert(0, self.create_time_index_column_expr()?); func_exprs.extend_from_slice(&self.create_tag_column_exprs()?); @@ -2034,7 +2212,7 @@ impl PromPlanner { async fn selector_to_series_normalize_plan( &mut self, - offset: &Option, + offset_duration: Millisecond, label_matchers: Matchers, is_range_selector: bool, ) -> Result { @@ -2044,11 +2222,6 @@ impl PromPlanner { let table_schema = table_scan.schema(); // make filter exprs - let offset_duration = match offset { - Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, - Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), - None => 0, - }; let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?; if let Some(time_index_filter) = self.build_time_index_filter(offset_duration, table_schema)? @@ -2147,11 +2320,7 @@ impl PromPlanner { } // make sort plan - let series_key_columns = if self.ctx.use_tsid { - vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] - } else { - self.ctx.tag_columns.clone() - }; + let series_key_columns = self.series_key_columns(); let sort_exprs = if self.ctx.use_tsid { vec![ @@ -2948,6 +3117,7 @@ impl PromPlanner { float_field: &str, histogram_field: &str, input_schema: &DFSchemaRef, + range_fold_offset: Option, ) -> Result>> { let returns_histogram = matches!( func.name, @@ -2987,6 +3157,14 @@ impl PromPlanner { Box::new(other_input_exprs[0].clone()), ArrowDataType::Int64, )); + // Same evaluation instant as in the non-mixed path; it follows the other inputs so the + // float UDF below is called as `predict_linear(ts_range, value_range, t, eval_ts)`. + // The regression is only defined for a folded window, so a missing fold offset is an + // error instead of a default. + other_input_exprs.push_back(self.create_range_eval_ts_expr( + range_fold_offset.context(ExpectRangeSelectorSnafu)?, + input_schema, + )?); } let timestamp_range = DfExpr::Column(Column::from_name( @@ -3045,7 +3223,13 @@ impl PromPlanner { .alias(histogram_field), ] } else { - let display_name = float_expr.schema_name().to_string(); + let display_name = if func.name == "predict_linear" { + // The evaluation instant is the private last argument of the mixed float UDF + // call; keep it out of the output column name like on the non-mixed path. + Self::name_without_last_arg(&float_expr) + } else { + float_expr.schema_name().to_string() + }; self.ctx.field_columns = vec![display_name.clone()]; vec![float_expr.alias(display_name)] }; @@ -3063,6 +3247,7 @@ impl PromPlanner { other_input_exprs: Vec, input_schema: &DFSchemaRef, query_engine_state: &QueryEngineState, + range_fold_offset: Option, ) -> Result<(Vec, Vec)> { // TODO(ruihang): check function args list let mut other_input_exprs: VecDeque = other_input_exprs.into(); @@ -3075,6 +3260,7 @@ impl PromPlanner { &float_field, &histogram_field, input_schema, + range_fold_offset, )? { return Ok((exprs, vec![])); @@ -3276,6 +3462,16 @@ impl PromPlanner { Box::new(other_input_exprs[0].clone()), ArrowDataType::Int64, )); + // The prediction starts at the evaluation instant of the step, which the + // window's fold offset recovers from the row's time index; see + // [`Self::create_range_eval_ts_expr`]. It is appended last, so the UDF is + // called as `predict_linear(ts_range, value_range, t, eval_ts)`. The + // regression is only defined for a folded window, so a missing fold offset + // is an error instead of a default. + other_input_exprs.push_back(self.create_range_eval_ts_expr( + range_fold_offset.context(ExpectRangeSelectorSnafu)?, + input_schema, + )?); ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf())) } } @@ -3648,7 +3844,17 @@ impl PromPlanner { exprs = exprs .into_iter() .map(|expr| { - let display_name = expr.schema_name().to_string(); + // `predict_linear` appends its private evaluation instant as the last argument + // of the UDF call; the output column is named after the call without it, so + // the injected expression stays out of the user-visible schema. The + // native-histogram drop UDF takes no such argument. + let display_name = if func.name == "predict_linear" + && !all_field_columns_are_native_histogram_ranges + { + Self::name_without_last_arg(&expr) + } else { + expr.schema_name().to_string() + }; new_field_columns.push(display_name.clone()); Ok(expr.alias(display_name)) }) @@ -3917,6 +4123,66 @@ impl PromPlanner { ))) } + /// Builds the evaluation instant the window of the last planned range selector is folded for, + /// as a `Timestamp(Millisecond)` expression. + /// + /// The timestamp payload of a folded window is shifted onto the evaluation timeline by the + /// offset the window is folded with (`fold_offset`), while the time index column of a folded + /// row keeps the evaluation timestamp of its step. Adding the offset back yields the + /// evaluation instant on the payload timeline, which is where the regression of + /// `predict_linear` is centered: neither a plain window (which may end before the step, and + /// is additionally shifted by `offset` on the payload timeline) nor an `@`-anchored one + /// (whose end is the anchor, while the payload is shifted by `at_offset`) ends at the step it + /// is evaluated at. + /// + /// The sum is computed on the millisecond representation and cast back, so that the result + /// keeps the `Timestamp(Millisecond)` type the range functions declare for it. + fn create_range_eval_ts_expr( + &self, + fold_offset: Millisecond, + input_schema: &DFSchemaRef, + ) -> Result { + let eval_ts = self + .create_time_index_column_expr()? + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + input_schema, + ) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Int64, input_schema) + .context(DataFusionPlanningSnafu)?; + DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(eval_ts), + op: Operator::Plus, + right: Box::new(lit(fold_offset)), + }) + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + input_schema, + ) + .context(DataFusionPlanningSnafu) + } + + /// The name of `expr` without its last argument. + /// + /// A `predict_linear` call appends the evaluation instant of its window as a private last + /// argument ([`Self::create_range_eval_ts_expr`]). Naming the output column after the call + /// the user wrote, without that argument, keeps the injected expression out of the + /// user-visible schema; the expression itself keeps every argument it needs. + fn name_without_last_arg(expr: &DfExpr) -> String { + if let DfExpr::ScalarFunction(ScalarFunction { func, args }) = expr + && let Some((_, visible_args)) = args.split_last() + { + let visible = ScalarFunction { + func: func.clone(), + args: visible_args.to_vec(), + }; + return DfExpr::ScalarFunction(visible).schema_name().to_string(); + } + + expr.schema_name().to_string() + } + fn create_tag_column_exprs(&self) -> Result> { let mut result = Vec::with_capacity(self.ctx.tag_columns.len()); for tag in &self.ctx.tag_columns { diff --git a/src/query/src/promql/planner/at_modifier.rs b/src/query/src/promql/planner/at_modifier.rs new file mode 100644 index 00000000000..bde7c51bcb1 --- /dev/null +++ b/src/query/src/promql/planner/at_modifier.rs @@ -0,0 +1,326 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Anchored selector planning and replay for the PromQL `@` modifier. + +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use datafusion::prelude::{Column, Expr as DfExpr}; +use promql::extension_plan::{InstantManipulate, Millisecond, SeriesDivide}; +use promql_parser::parser::{ + AtModifier, Call, Expr as PromExpr, MatrixSelector, Offset, ParenExpr, +}; +use snafu::{OptionExt, ResultExt, ensure}; + +use crate::promql::error::{ + AtModifierTimestampOutOfRangeSnafu, DataFusionPlanningSnafu, Result, TimeIndexNotFoundSnafu, +}; +use crate::promql::planner::PromPlanner; +use crate::query_engine::QueryEngineState; + +impl PromPlanner { + /// Resolve the `@` modifier of a vector or matrix selector into the timestamp its sample + /// window is anchored at, in milliseconds since the Unix epoch. + /// + /// Prometheus semantics: + /// - `@ ` anchors at the given timestamp, + /// - `@ start()` / `@ end()` anchor at the evaluation range of the whole statement, + /// - `offset` shifts the anchor backwards: the window ends at `anchor - offset`. + /// + /// Returns `None` when the selector has no `@` modifier. + fn at_ref_time( + &self, + at: &Option, + offset: &Option, + ) -> Result> { + let anchor = match at { + None => return Ok(None), + Some(AtModifier::Start) => self.ctx.stmt_start, + Some(AtModifier::End) => self.ctx.stmt_end, + Some(AtModifier::At(time)) => Self::system_time_to_millis(time)?, + }; + Ok(Some(Self::anchor_sub(anchor, Self::offset_millis(offset))?)) + } + + /// Subtracts `rhs` from `lhs` on the millisecond timeline of an `@` anchor. + /// + /// A negative result is valid: `@` and `offset` accept timestamps before the Unix epoch. A + /// result outside the representable millisecond range is rejected like an unrepresentable + /// anchor ([`Self::system_time_to_millis`]) instead of clamping it, so the same class of + /// input always gets the same answer. + pub(crate) fn anchor_sub(lhs: Millisecond, rhs: Millisecond) -> Result { + lhs.checked_sub(rhs) + .with_context(|| AtModifierTimestampOutOfRangeSnafu { + timestamp: format!("{}ms - {}ms", lhs, rhs), + }) + } + + /// The offset a selector with an `@` modifier is evaluated with. + /// + /// Prometheus anchors such a selector by rewriting its offset to `eval_time - anchor` + /// (`setOffsetForAtModifier`), so that the selector always selects its samples around `anchor` + /// regardless of the step being evaluated. `eval_time` is the start of the evaluation the + /// selector belongs to, which is `ctx.start`. + /// + /// Returns `None` when the selector has no `@` modifier. + pub(crate) fn at_modifier_offset( + &self, + at: &Option, + offset: &Option, + ) -> Result> { + let Some(anchor) = self.at_ref_time(at, offset)? else { + return Ok(None); + }; + Ok(Some(Self::anchor_sub(self.ctx.start, anchor)?)) + } + + /// Whether `expr` is a call that has to be evaluated once for the whole grid, because it folds + /// a range selector anchored by `@` — the range argument of the call's parser signature, which + /// is a [`MatrixSelector`] here; only a call with such an argument can take this path, so no + /// function-name registry is involved. + /// + /// This is the shape that needs Prometheus' `StepInvariantExpr` the most: a range function such + /// as `rate` derives its result from the evaluation instant it is called at, so folding the + /// window once per step would let the outer evaluation grid change the result of a window that + /// `@` fixed. Evaluated once, at the start of the grid, the rewritten offset of + /// [`Self::at_modifier_offset`] places the anchor at that instant, and [`Self::replay_over_grid`] + /// reports the result at every step. + /// + /// Unlike Prometheus, which wraps the whole step-invariant subtree (`preprocessExprHelper`), + /// only the call itself is promoted here. The operators above it are not: they are still planned + /// at every step over the replayed result, which is safe for the row-wise ones and keeps the + /// promotion root narrow. The promotion root has to stay a call over one range selector, because + /// the replay needs one series per batch ([`Self::series_divide_plan`]) and only such a call + /// guarantees that the rows it emits still describe the series it was divided by. The operators + /// left out — an aggregation, a join, or a label rewriting call such as `label_join` — mix or + /// re-label the rows of different series, so they are unsafe as promotion roots even though + /// evaluating them after the replay is fine. A call whose input is an anchored *instant* + /// selector (`abs(some_metric @ 300)`) needs no promotion either: the selector anchors and + /// replays its sample per series on its own, and the call above it is row-wise. + /// + /// `predict_linear` is the exception among the range functions: it predicts from the evaluation + /// instant of each step ([`Self::create_range_eval_ts_expr`]), so it has to stay outside the + /// promoted subtree and follow the grid. The remaining arguments of the call have to be + /// literals, since the replay of the promoted result has no second vector input to divide. + /// Parentheses around the range argument are transparent (`rate((m[5m] @ 300))`), so they are + /// looked through and the call is promoted as if they were absent. Nothing else of the subtree + /// is unwrapped, so an outer parenthesis promotes no operator above the call. + fn promotes_anchored_range_call(expr: &PromExpr) -> bool { + let PromExpr::Call(Call { func, args }) = expr else { + return false; + }; + // See the doc comment: the regression of `predict_linear` follows the evaluation step. + if func.name == "predict_linear" { + return false; + } + let mut anchored_range = false; + for arg in &args.args { + // Parentheses around the range argument are transparent, so the call is promoted the + // same way for `rate((m[5m] @ 300))` as for `rate(m[5m] @ 300)`. Only the parentheses + // directly around this one argument are looked through here: the promotion stays + // confined to a call over one anchored range selector instead of descending into an + // arbitrary parenthesized subtree. + let mut arg = arg.as_ref(); + while let PromExpr::Paren(ParenExpr { expr }) = arg { + arg = expr; + } + match arg { + // The window is pinned by `@`, so every step folds the same samples. + PromExpr::MatrixSelector(MatrixSelector { vs, .. }) if vs.at.is_some() => { + if anchored_range { + return false; + } + anchored_range = true; + } + // A literal argument is the same value at every step. + arg if Self::try_build_literal_expr(arg).is_some() => {} + _ => return false, + } + } + anchored_range + } + + /// Plans the anchored range call `prom_expr` ([`Self::promotes_anchored_range_call`]) as a + /// step-invariant subtree: the call is evaluated on a single evaluation instant (`grid_start`, + /// the start of the outer evaluation) and its result is then reported at every step of + /// `[grid_start, ctx.end]` by [`Self::replay_over_grid`]. + /// + /// This is the planner's counterpart of Prometheus' `StepInvariantExpr` for the one shape it + /// promotes. Evaluating the call once matters for the functions that derive their result from + /// the step being evaluated: `rate(m[5m] @ 300)` folds its window around the anchor once, and + /// the extrapolation boundaries of `rate` must be derived from that same window at every step + /// instead of following the outer evaluation timestamp. + /// + /// Only the call itself is promoted; the operators above it are planned as usual over the + /// replayed result. The result of the promoted call is split into one series per batch before it + /// is replayed ([`Self::series_divide_plan`]), because the row-wise projection of the call does + /// not preserve the batch layout of the selector. + /// + /// Returns `None` when `prom_expr` is not such a call, so that the caller plans it as usual. + /// The selector inside the promoted call keeps its own `@` anchoring (see + /// [`Self::at_modifier_offset`]), and planning it with `ctx.end == ctx.start` folds its window + /// once for that single instant instead of expanding it over the grid, which the replay of the + /// call result above already does. + pub(crate) async fn promote_anchored_range_call( + &mut self, + prom_expr: &PromExpr, + timestamp_fn: bool, + query_engine_state: &QueryEngineState, + ) -> Result> { + let grid_start = self.ctx.start; + let grid_end = self.ctx.end; + // An instant query evaluates a single step, so there is nothing to promote. + if grid_start == grid_end || !Self::promotes_anchored_range_call(prom_expr) { + return Ok(None); + } + + // Plan the subtree on a single evaluation instant: every selector below still anchors its + // window through `@`, and the functions above them derive their result from that one + // instant. The planner is single-use, so `ctx.end` needs no restore-on-error. + self.ctx.end = grid_start; + let anchored = self + .prom_expr_to_plan_inner(prom_expr, timestamp_fn, query_engine_state) + .await?; + self.ctx.end = grid_end; + + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + // The replay reads one series per batch; see [`Self::series_divide_plan`] for why the + // layout of the selector below the promoted call does not survive it. + let anchored = self.series_divide_plan(anchored, &time_index_column)?; + Ok(Some(self.replay_over_grid( + anchored, + grid_start, + grid_end, + time_index_column, + ))) + } + + /// Convert the timestamp of an `@` modifier into milliseconds since the Unix epoch. + fn system_time_to_millis(time: &SystemTime) -> Result { + let (millis, negative) = match time.duration_since(UNIX_EPOCH) { + Ok(duration) => (duration.as_millis(), false), + // The `@` modifier accepts timestamps before the Unix epoch, e.g. `@ -1`. + Err(err) => (err.duration().as_millis(), true), + }; + ensure!( + millis <= i64::MAX as u128, + AtModifierTimestampOutOfRangeSnafu { + timestamp: time + .duration_since(UNIX_EPOCH) + .map(|duration| format!("+{}ms", duration.as_millis())) + .unwrap_or_else(|err| format!("-{}ms", err.duration().as_millis())), + } + ); + let millis = millis as Millisecond; + Ok(if negative { -millis } else { millis }) + } + + /// Report the samples of `anchored` at every step of the evaluation grid + /// `[grid_start, grid_end]`. + /// + /// A selector with an `@` modifier is anchored: the sample window is selected once, around the + /// anchor timestamp, and every evaluation step reports that same window. Prometheus does this + /// by rewriting the selector's offset to `eval_time - anchor` and only fetching the samples on + /// the first step (`setOffsetForAtModifier` plus the `refetch` shortcut in `rangeEval`). + /// + /// The expansion reuses [`InstantManipulate`] with a lookback that spans the whole grid: every + /// step then picks the same sample (or the same already computed value, when `anchored` ends + /// with a function call such as `rate`) and stamps it with the step's timestamp. + /// + /// Every input batch of `anchored` must hold exactly one series, because [`InstantManipulate`] + /// takes a batch as one timeline. A leaf-level replay (`m @ 300`) consumes the [`SeriesDivide`] + /// of its selector directly. A promoted call is guaranteed that layout by + /// [`Self::series_divide_plan`], which is why [`Self::promote_anchored_range_call`] splits + /// its result before calling this method. + pub(crate) fn replay_over_grid( + &self, + anchored: LogicalPlan, + grid_start: Millisecond, + grid_end: Millisecond, + time_index_column: String, + ) -> LogicalPlan { + if grid_start == grid_end { + // A single evaluation step: `anchored` is already stamped with that timestamp. + return anchored; + } + + let series_key_columns = self.series_key_columns_for_schema(anchored.schema()); + let replayed = InstantManipulate::new( + grid_start, + grid_end, + // The lookback must keep the single anchored sample eligible for every step. + grid_end - grid_start + 1, + self.ctx.interval, + 0, + time_index_column, + series_key_columns, + self.ctx.field_columns.first().cloned(), + anchored, + ); + LogicalPlan::Extension(Extension { + node: Arc::new(replayed), + }) + } + + /// Sorts `input` by its series key and time index and splits it into one batch per series. + /// + /// [`InstantManipulate`] reads every input batch as one series (it takes the timeline of the + /// batch and reports the row selected at every step), so a batch holding several series would + /// lose all but one of them. A selector establishes that layout with its own [`SeriesDivide`], + /// but the per-series distribution requirement does not reach a promoted subtree above it: + /// the row-wise projection of a call sits in between, so the batch boundaries of the selector + /// are not preserved — in a distributed plan the promoted result can be delivered as one batch + /// holding every series. Sorting and dividing here restores the layout, exactly like + /// [`Self::prom_matrix_selector_to_plan`] does for the input of a range function. + /// + /// Series keys that are not present in `input` are dropped, since `ctx.tag_columns` may have + /// drifted from the actual output schema. A plan without any series key column is returned + /// unchanged: there is nothing to divide by. + fn series_divide_plan( + &self, + input: LogicalPlan, + time_index_column: &str, + ) -> Result { + let series_key_columns = self.series_key_columns_for_schema(input.schema()); + if series_key_columns.is_empty() { + return Ok(input); + } + + let mut sort_exprs = series_key_columns + .iter() + .map(|name| DfExpr::Column(Column::from_name(name)).sort(true, true)) + .collect::>(); + sort_exprs.push(DfExpr::Column(Column::from_name(time_index_column)).sort(true, true)); + let sort_plan = LogicalPlanBuilder::from(input) + .sort(sort_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + series_key_columns, + time_index_column.to_string(), + sort_plan, + )), + })) + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 1634e6d9e1d..a83a42bae53 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1929,6 +1929,600 @@ async fn tsid_is_used_for_series_divide_when_available() { assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); } +async fn build_at_modifier_plan(query: &str, start_secs: u64, end_secs: u64) -> LogicalPlan { + let eval_stmt = build_at_modifier_eval_stmt(query, start_secs, end_secs); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap() +} + +fn build_at_modifier_eval_stmt(query: &str, start_secs: u64, end_secs: u64) -> EvalStmt { + EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH + .checked_add(Duration::from_secs(start_secs)) + .unwrap(), + end: UNIX_EPOCH + .checked_add(Duration::from_secs(end_secs)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } +} + +/// Every selector of an `@` anchored query must scan around the anchor only, instead of the +/// whole evaluation range. +#[tokio::test] +async fn at_modifier_anchors_selector_scan_window() { + // `@ 100` anchors at t=100s; the lookback delta is 1s, so the scan covers (99s, 100s]. + let plan = build_at_modifier_plan("some_metric @ 100", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(99001, None) AND some_metric.timestamp <= TimestampMillisecond(100000, None)" + ), + "{plan_str}" + ); + // The result is reported at the evaluation timestamps, not at the anchor. + assert!(plan_str.contains("range=[0..1000000]"), "{plan_str}"); + + // `@ start()` / `@ end()` resolve to the evaluation range of the statement. + let plan = build_at_modifier_plan("some_metric @ start()", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(0, None)" + ), + "{plan_str}" + ); + + let plan = build_at_modifier_plan("some_metric @ end()", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(999001, None) AND some_metric.timestamp <= TimestampMillisecond(1000000, None)" + ), + "{plan_str}" + ); + + // `offset` moves the anchor backwards and is not applied twice. + let plan = build_at_modifier_plan("some_metric @ 200 offset 50s", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(149001, None) AND some_metric.timestamp <= TimestampMillisecond(150000, None)" + ), + "{plan_str}" + ); + + // A timestamp before the Unix epoch is accepted, as in Prometheus. + let plan = build_at_modifier_plan("some_metric @ -1", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(-1999, None) AND some_metric.timestamp <= TimestampMillisecond(-1000, None)" + ), + "{plan_str}" + ); +} + +/// A call over one range selector anchored by `@` is evaluated once, at the start of the +/// evaluation, and its result is reported at every step: the window is folded around the anchor +/// instead of following the outer evaluation grid. This is the planner's counterpart of +/// Prometheus' `StepInvariantExpr` wrapper; see [`PromPlanner::promotes_anchored_range_call`]. +#[tokio::test] +async fn at_modifier_promotes_anchored_range_call() { + let plan = build_at_modifier_plan("rate(some_metric[5m] @ 300)", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The scan is limited to the anchored window (offset by `eval_start - anchor`). + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(1, None) AND some_metric.timestamp <= TimestampMillisecond(300000, None)" + ), + "{plan_str}" + ); + // A single fold, at the anchor... + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + // ... never a fold per step of the outer grid. + assert!( + !plan_str.contains("PromRangeManipulate: req range=[0..1000000]"), + "{plan_str}" + ); + // A single replay of the function result over the whole grid... + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{plan_str}" + ); + // ... so `rate` itself is evaluated below that replay node, on the single evaluation + // instant of the anchored subtree, instead of once per step. + let replay = plan_str + .find("PromInstantManipulate") + .expect("instant manipulate node"); + let rate = plan_str.find("prom_rate(").expect("rate projection"); + assert!( + replay < rate, + "`rate` must be evaluated below the replay node:\n{plan_str}" + ); +} + +/// Parentheses around the range argument are transparent: `rate((some_metric[5m] @ 300))` gets +/// the same fixed-window promotion as `rate(some_metric[5m] @ 300)`, with the anchored window +/// folded once and its result replayed at every step of the grid. Only that one argument is +/// looked through, so a parenthesis above the call promotes no operator of its own and a +/// parenthesized subtree is planned exactly like the bare one. +#[tokio::test] +async fn at_modifier_promotes_parenthesized_range_argument() { + // Each form is planned exactly like its unparenthesized counterpart: parentheses below the + // call are transparent, a parenthesis around the call adds nothing, and a parenthesis above + // it does not widen the promotion. + for (query, plain) in [ + ( + "rate((some_metric[5m] @ 300))", + "rate(some_metric[5m] @ 300)", + ), + ( + "rate(((some_metric[5m] @ 300)))", + "rate(some_metric[5m] @ 300)", + ), + ( + "(rate(some_metric[5m] @ 300))", + "rate(some_metric[5m] @ 300)", + ), + ( + "abs((rate(some_metric[5m] @ 300)))", + "abs(rate(some_metric[5m] @ 300))", + ), + ] { + assert_eq!( + build_at_modifier_plan(query, 0, 1000) + .await + .display_indent_schema() + .to_string(), + build_at_modifier_plan(plain, 0, 1000) + .await + .display_indent_schema() + .to_string(), + "`{query}` must be planned like `{plain}`" + ); + } + + // The parentheses do not push the enclosing operator into the promotion either: `abs` stays + // above the replay of the promoted call, exactly as it does without them. The promotion + // itself (one anchored fold, one replay, `rate` below it) is asserted for the + // unparenthesized form by `at_modifier_promotes_anchored_range_call`, and the form above is + // planned identically to it. + let plan_str = build_at_modifier_plan("abs((rate(some_metric[5m] @ 300)))", 0, 1000) + .await + .display_indent_schema() + .to_string(); + let replay = plan_str + .find("PromInstantManipulate") + .expect("instant manipulate node"); + assert!( + plan_str.find("abs(").expect("`abs` projection") < replay, + "`abs` must be evaluated above the replay of the promoted call:\n{plan_str}" + ); +} + +/// `@ start()` and `@ end()` are fixed anchors for the whole statement, so a call using them is +/// promoted as well. +#[tokio::test] +async fn at_modifier_promotes_start_and_end_anchored_call() { + for query in [ + "rate(some_metric[5m] @ start())", + "rate(some_metric[5m] @ end())", + "max_over_time(some_metric[5m] @ end())", + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{query}:\n{plan_str}" + ); + } + + // Only the call itself is promoted: a call above it (`abs`) is planned as usual and + // evaluated at every step over the replayed result of the promoted call. The window is + // still folded once, around the anchor. + let plan = build_at_modifier_plan("abs(max_over_time(some_metric[5m] @ end()))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + let abs = plan_str.find("abs(").expect("`abs` projection"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + abs < replay, + "`abs` must be evaluated above the replay of the promoted call:\n{plan_str}" + ); + // The window is the only one folded once, and it sits below the replay: the promoted call + // feeds the grid from there. + let fold = plan_str.find("PromRangeManipulate").expect("range fold"); + assert!( + replay < fold, + "the anchored window must be folded below the replay:\n{plan_str}" + ); + + // An aggregation above the promoted call stays above it as well: `sum` aggregates the + // replayed per-series rows at every step, instead of aggregating the single anchored + // instant and replaying the aggregation — which would also have to replay rows of several + // groups through the one-series-per-batch `InstantManipulate`. + let plan = build_at_modifier_plan("sum(rate(some_metric[5m] @ start()))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + // A single replay, and it sits below the aggregation node: `sum` aggregates the replayed + // per-series rows at every step. + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{plan_str}" + ); + let aggregate = plan_str.find("Aggregate:").expect("aggregate node"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + aggregate < replay, + "the aggregation must stay above the replay of the promoted call:\n{plan_str}" + ); + + // Without `@` nothing is promoted: the function keeps folding one window per step. + let plan = build_at_modifier_plan("rate(some_metric[5m])", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!( + plan_str.contains("PromRangeManipulate: req range=[0..1000000]"), + "{plan_str}" + ); + assert!(!plan_str.contains("lookback=[1000001]"), "{plan_str}"); +} + +/// A call or a unary operator above the anchored range call is planned as usual: the inner range +/// call is promoted on its own (it is the direct call over the anchored range selector), and the +/// operator above it is evaluated at every step over the replayed result. +#[tokio::test] +async fn at_modifier_promotes_inner_range_call_below_wrappers() { + for (query, wrapper) in [ + ("abs(rate(some_metric[5m] @ 300))", "abs(prom_rate("), + ("-rate(some_metric[5m] @ 300)", "(- prom_rate("), + ( + "abs(max_over_time(some_metric[5m] @ 300))", + "abs(prom_max_over_time(", + ), + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The anchored window is folded once, and the wrapper sits above its replay. + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 1, + "{query}:\n{plan_str}" + ); + let wrapper = plan_str + .find(wrapper) + .unwrap_or_else(|| panic!("no `{wrapper}` projection in:\n{plan_str}")); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + wrapper < replay, + "the wrapper must be evaluated above the replay of the range call:\n{plan_str}" + ); + } +} + +/// `anchored + plain`: only the binary operand that is a call over the anchored range selector +/// is promoted, and the plain side keeps following the evaluation step. +#[tokio::test] +async fn at_modifier_promotes_only_anchored_binary_operand() { + let plan = build_at_modifier_plan("rate(some_metric[5m] @ 300) + some_metric", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The anchored operand is folded once and its result replayed over the whole grid. + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000001]"), + "{plan_str}" + ); + // The plain operand still selects one sample per step with the default lookback. + assert!( + plan_str.contains("PromInstantManipulate: range=[0..1000000], lookback=[1000]"), + "{plan_str}" + ); + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 2, + "{plan_str}" + ); + + // Both operands anchored: the binary expression is not promoted as a whole, because a join + // emits the rows of several series in shared batches and the replay needs one series per + // batch. Each operand anchors and replays on its own instead, and the join runs at every + // step over those per-series results. + let plan = build_at_modifier_plan("some_metric @ 300 + some_metric @ 0", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // One anchoring node and one replay per operand... + assert_eq!( + plan_str.matches("PromInstantManipulate").count(), + 4, + "{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 2, + "{plan_str}" + ); + // ... and the two operands are anchored at different timestamps, so both select their own + // sample. + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..0], lookback=[1000]") + .count(), + 2, + "{plan_str}" + ); +} + +/// Only a direct call over an anchored range selector is promoted. A value function or a unary +/// operator over an anchored *instant* selector needs no promotion: the selector anchors and +/// replays its sample per series on its own, and the operator above it is row-wise, so it can be +/// evaluated at every step over that replay. +#[tokio::test] +async fn at_modifier_does_not_promote_value_calls_over_anchored_selectors() { + for (query, value_expr) in [ + ("abs(some_metric @ 300)", "abs(some_metric.field_0)"), + ("-some_metric @ 300", "(- some_metric.field_0)"), + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The scan is limited to the anchored sample (the lookback delta of this test is 1s). + assert!( + plan_str.contains( + "some_metric.timestamp >= TimestampMillisecond(299001, None) AND some_metric.timestamp <= TimestampMillisecond(300000, None)" + ), + "{query}:\n{plan_str}" + ); + // ... and the result is replayed at every step by the selector itself: the anchored + // selection, then the grid replay, with no promoted subtree on top of the operator. + assert!( + !plan_str.starts_with("PromInstantManipulate"), + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..0], lookback=[1000]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .unwrap_or_else(|| panic!("replay node:\n{plan_str}")); + let value = plan_str + .find(value_expr) + .unwrap_or_else(|| panic!("no `{value_expr}` projection in:\n{plan_str}")); + assert!( + value < replay, + "`{value_expr}` must be evaluated above the per-series replay:\n{plan_str}" + ); + } +} + +/// A call whose argument merges series (an aggregation or a join) or whose own output reorders +/// the whole vector (`sort*`, the histogram folds) is never the promoted root either: it is +/// planned as usual over the leaf-level anchoring of its selectors, which replays every selector +/// per series and keeps the one-series-per-batch layout the replay needs. +#[tokio::test] +async fn at_modifier_keeps_multi_series_roots_out_of_promoted_subtree() { + // An aggregation below the call emits one row per group in shared batches, so the call is + // not promoted: the replay of the anchored selector stays below the aggregate. + let plan = build_at_modifier_plan("abs(sum(some_metric @ 300))", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + let aggregate = plan_str.find("Aggregate:").expect("aggregate node"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + aggregate < replay, + "the aggregation must stay above the replay:\n{plan_str}" + ); + + // A join of two anchored selectors below the call: neither the join nor the call is + // promoted, so each operand is replayed on its own. + let plan = build_at_modifier_plan("abs(some_metric @ 300 + some_metric @ 0)", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!(!plan_str.starts_with("PromInstantManipulate"), "{plan_str}"); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 2, + "{plan_str}" + ); + + // A call that reorders the whole vector keeps its sort above the per-series replay. + let plan = build_at_modifier_plan("sort_by_label(some_metric @ 300, \"tag_0\")", 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + assert!(plan_str.starts_with("Sort:"), "{plan_str}"); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 1, + "{plan_str}" + ); +} + +/// `label_join` rewrites the labels of its input series, so it is never the promoted root: the +/// anchored selector keeps replaying one series per batch, and the join runs at every step above +/// that replay (see [`Self::promotes_anchored_range_call`]). Promoting it would replay the +/// joined rows through the labels the join just rewrote, which merges the distinct input series +/// into one timeline. +#[tokio::test] +async fn at_modifier_does_not_promote_label_join() { + for query in [ + // Directly above the anchored instant selector... + "label_join(some_metric @ 300, \"tag_0\", \"-\", \"\")", + // ... and below another call, which is planned as usual over the join. + "abs(label_join(some_metric @ 300, \"tag_0\", \"-\", \"\"))", + ] { + let plan = build_at_modifier_plan(query, 0, 1000).await; + let plan_str = plan.display_indent_schema().to_string(); + // The join is not wrapped in a replay of its own; the only replay over the grid is the + // one of the anchored selector... + assert!( + !plan_str.starts_with("PromInstantManipulate"), + "{query}:\n{plan_str}" + ); + assert_eq!( + plan_str + .matches("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .count(), + 1, + "{query}:\n{plan_str}" + ); + // ... and the projected join stays above it, evaluated at every step. + let join = plan_str + .find("concat_ws(") + .unwrap_or_else(|| panic!("no `label_join` projection in:\n{plan_str}")); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + join < replay, + "`label_join` must be evaluated above the per-series replay:\n{plan_str}" + ); + } + + // A range call below the join is still promoted on its own: the anchored window is folded + // once per series, and the join above it is evaluated at every step over that replay. + let plan = build_at_modifier_plan( + "label_join(rate(some_metric[5m] @ 300), \"tag_0\", \"-\", \"\")", + 0, + 1000, + ) + .await; + let plan_str = plan.display_indent_schema().to_string(); + assert_eq!( + plan_str + .matches("PromRangeManipulate: req range=[0..0]") + .count(), + 1, + "{plan_str}" + ); + assert!(!plan_str.starts_with("PromInstantManipulate"), "{plan_str}"); + let join = plan_str.find("concat_ws(").expect("join projection"); + let replay = plan_str + .find("PromInstantManipulate: range=[0..1000000], lookback=[1000001]") + .expect("replay node"); + assert!( + join < replay, + "the join must be evaluated above the replay of the range call:\n{plan_str}" + ); +} + +#[test] +fn at_modifier_rejects_subtraction_overflow() { + for (anchor, offset) in [(i64::MAX, -1), (i64::MIN, 1)] { + let err = PromPlanner::anchor_sub(anchor, offset).unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + assert!( + err.to_string() + .contains("Timestamp out of range for the `@` modifier"), + "{err}" + ); + } + assert_eq!(PromPlanner::anchor_sub(-1, 1).unwrap(), -2); +} + +/// `@` beyond the representable millisecond range is rejected instead of silently wrapping. +/// +/// `@ 1e16` is 10^19 milliseconds, beyond `i64::MAX`. A Unix `SystemTime` can hold it, so +/// the planner rejects the anchor it cannot represent. A Windows `SystemTime` tops out +/// below `i64::MAX` milliseconds, so the same literal is already rejected while parsing. +#[tokio::test] +async fn at_modifier_rejects_unrepresentable_timestamp() { + #[cfg(windows)] + { + let err = parser::parse("some_metric @ 1e16").unwrap_err(); + assert!( + err.to_string() + .contains("timestamp out of bounds for @ modifier"), + "{err}" + ); + } + + #[cfg(not(windows))] + { + let eval_stmt = build_eval_stmt("some_metric @ 1e16"); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let err = + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Timestamp out of range for the `@` modifier"), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + } +} + #[tokio::test] async fn default_binary_join_uses_tsid_when_available() { let eval_stmt = build_eval_stmt("some_metric / some_alt_metric"); @@ -3421,7 +4015,7 @@ async fn binary_op_column_column() { assert_eq!(plan.display_indent_schema().to_string(), expected); } -async fn indie_query_plan_compare>(query: &str, expected: T) { +async fn indie_query_plan(query: &str) -> String { let prom_expr = parser::parse(query).unwrap(); let eval_stmt = EvalStmt { expr: prom_expr, @@ -3449,7 +4043,12 @@ async fn indie_query_plan_compare>(query: &str, expected: T) { .await .unwrap(); - assert_eq!(plan.display_indent_schema().to_string(), expected.as_ref()); + plan.display_indent_schema().to_string() +} + +async fn indie_query_plan_compare>(query: &str, expected: T) { + let plan = indie_query_plan(query).await; + assert_eq!(plan, expected.as_ref()); } #[tokio::test] @@ -3537,6 +4136,51 @@ async fn increase_aggr() { indie_query_plan_compare(query, expected).await; } +#[tokio::test] +async fn predict_linear_injects_the_eval_timestamp() { + // The fourth argument is the evaluation instant of each row, which the fold offset + // recovers from the row's time index (here: no `@` and no `offset`, so the step itself). + let query = "predict_linear(some_metric[5m], 60)"; + let expected = String::from( + "Filter: prom_predict_linear(timestamp_range,field_0,Float64(60)) IS NOT NULL [timestamp:Timestamp(ms), prom_predict_linear(timestamp_range,field_0,Float64(60)):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_predict_linear(timestamp_range, field_0, CAST(Float64(60) AS Int64), CAST(CAST(some_metric.timestamp AS Int64) + Int64(0) AS Timestamp(ms))) AS prom_predict_linear(timestamp_range,field_0,Float64(60)), some_metric.tag_0 [timestamp:Timestamp(ms), prom_predict_linear(timestamp_range,field_0,Float64(60)):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[300000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-299999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + + indie_query_plan_compare(query, expected).await; +} + +/// The evaluation instant follows the selector's fold offset, not the projection's: an +/// `@`-anchored selector is folded with `at_offset`, which is what the window's timestamps were +/// shifted by. The evaluation starts at 0s here, so `@ 100` anchors 100s in the future and +/// yields `at_offset` = -100000ms. +#[tokio::test] +async fn predict_linear_eval_ts_follows_the_fold_offset() { + for (query, expected_offset) in [ + ( + "predict_linear(some_metric[5m] @ 100, 60)", + "Int64(-100000)", + ), + ( + "predict_linear(some_metric[5m] offset 2m, 60)", + "Int64(120000)", + ), + ] { + let plan = indie_query_plan(query).await; + assert!( + plan.contains(&format!( + "some_metric.timestamp AS Int64) + {expected_offset}" + )), + "{query}\n{plan}" + ); + } +} + async fn native_histogram_plan(query: &str) -> String { let table_provider = build_test_native_histogram_table_provider("some_metric").await; let plan = PromPlanner::stmt_to_plan( @@ -4214,6 +4858,27 @@ async fn mixed_native_histogram_ranges_use_coordinated_udfs() { assert_eq!(plan, expected); } +#[tokio::test] +async fn mixed_native_histogram_predict_linear_forwards_the_eval_timestamp() { + let query = "predict_linear(some_metric[5m], 60)"; + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + assert!( + plan.contains( + "prom_mixed_range_float(Utf8(\"predict_linear\"), timestamp_range, greptime_value, greptime_native_histogram, CAST(Float64(60) AS Int64), CAST(CAST(some_metric.timestamp AS Int64) + Int64(0) AS Timestamp(ms)))" + ), + "{query}\n{plan}" + ); +} + #[tokio::test] async fn mixed_native_histogram_rate_executes_real_ranges() { let schema = Arc::new(ArrowSchema::new(vec![ @@ -4288,7 +4953,7 @@ async fn mixed_native_histogram_rate_executes_real_ranges() { ); let state = build_query_engine_state(); let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) + .create_function_expr(&call.func, vec![], input.schema(), &state, None) .unwrap(); exprs.insert(0, planner.create_time_index_column_expr().unwrap()); let plan = LogicalPlanBuilder::from(input) @@ -7124,7 +7789,7 @@ async fn test_mixed_or_routes_float_histogram_and_label_functions() { }; let state = build_query_engine_state(); let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) + .create_function_expr(&call.func, vec![], input.schema(), &state, None) .unwrap(); exprs.insert(0, planner.create_time_index_column_expr().unwrap()); exprs.extend(planner.create_tag_column_exprs().unwrap()); @@ -7177,7 +7842,7 @@ async fn test_mixed_or_routes_float_histogram_and_label_functions() { let args = planner.create_function_args(&call.args.args).unwrap(); let state = build_query_engine_state(); let (mut exprs, _) = planner - .create_function_expr(&call.func, args.literals, input.schema(), &state) + .create_function_expr(&call.func, args.literals, input.schema(), &state, None) .unwrap(); exprs.insert(0, planner.create_time_index_column_expr().unwrap()); exprs.extend(planner.create_tag_column_exprs().unwrap()); diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 3727ea9a30c..a647821e0b3 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -475,7 +475,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { ); let state = build_query_engine_state(); let (mut exprs, _) = planner - .create_function_expr(&call.func, vec![], input.schema(), &state) + .create_function_expr(&call.func, vec![], input.schema(), &state, None) .unwrap(); exprs.insert(0, planner.create_time_index_column_expr().unwrap()); let plan = LogicalPlanBuilder::from(input) diff --git a/tests/cases/standalone/common/promql/at_modifier.result b/tests/cases/standalone/common/promql/at_modifier.result new file mode 100644 index 00000000000..0a19eac692c --- /dev/null +++ b/tests/cases/standalone/common/promql/at_modifier.result @@ -0,0 +1,743 @@ +-- Tests for the PromQL `@` modifier on vector and matrix selectors. +-- +-- `@` anchors the sample selection window at a fixed timestamp instead of the timestamp of each +-- evaluation step: `@ ` uses the given timestamp, `@ start()` / `@ end()` use the +-- start/end of the statement's evaluation range, and `offset` shifts the anchor backwards before +-- the window is selected. The output timestamps still follow the evaluation grid. +-- +-- Every metric below carries two series, host 'a' and host 'b', with different values: the +-- anchored results must therefore report both of them, per step. A series silently dropped, or a +-- batch holding several series handled as one timeline, does not look like a correct single-series +-- answer here but shows up as a missing or shifted row. +-- +-- Sample timestamps below are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds, +-- as in Prometheus. +CREATE TABLE at_modifier_gauge ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +CREATE TABLE at_modifier_counter_total ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +Affected Rows: 0 + +-- One sample per minute, so the anchored sample is easy to identify. The default lookback is 5m. +-- Host 'a' reports 0.0 to 6.0, host 'b' reports 7.0 to 13.0: the two series are told apart by their +-- values alone, which stays true after any anchoring. +INSERT INTO at_modifier_gauge VALUES + (0, 0.0, 'a'), + (60000, 1.0, 'a'), + (120000, 2.0, 'a'), + (180000, 3.0, 'a'), + (240000, 4.0, 'a'), + (300000, 5.0, 'a'), + (360000, 6.0, 'a'), + (0, 7.0, 'b'), + (60000, 8.0, 'b'), + (120000, 9.0, 'b'), + (180000, 10.0, 'b'), + (240000, 11.0, 'b'), + (300000, 12.0, 'b'), + (360000, 13.0, 'b'); + +Affected Rows: 14 + +-- Host 'a': a counter increasing by 60 per minute, i.e. by 1 per second. +-- Host 'b': a counter starting at 100 and increasing by 30 per minute, i.e. by 0.5 per second. Its +-- rate differs from host 'a', so a window folded for the wrong series cannot pass as the right one. +INSERT INTO at_modifier_counter_total VALUES + (0, 0.0, 'a'), + (60000, 60.0, 'a'), + (120000, 120.0, 'a'), + (180000, 180.0, 'a'), + (240000, 240.0, 'a'), + (300000, 300.0, 'a'), + (360000, 360.0, 'a'), + (0, 100.0, 'b'), + (60000, 130.0, 'b'), + (120000, 160.0, 'b'), + (180000, 190.0, 'b'), + (240000, 220.0, 'b'), + (300000, 250.0, 'b'), + (360000, 280.0, 'b'); + +Affected Rows: 14 + +-- 1. Instant query anchored at an absolute timestamp: the sample window is `(anchor - lookback, +-- anchor]` = `(0s, 300s]`, so its newest samples are 5.0 for host 'a' and 12.0 for host 'b', both +-- at 300s. The output timestamps stay the evaluation timestamps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 12.0 | b | +| 1970-01-01T00:01:40 | 5.0 | a | ++---------------------+------+------+ + +-- Control: the same query without `@` looks back from the evaluation timestamp instead: +-- `(-200s, 100s]` selects the samples 1.0 (a) and 8.0 (b) at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge; + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | ++---------------------+-----+------+ + +-- The anchor, not the evaluation timestamp, decides which sample is reported. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 360; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 13.0 | b | +| 1970-01-01T00:01:40 | 6.0 | a | ++---------------------+------+------+ + +-- Sub-second precision is preserved in the anchor: `240.999` selects the samples at 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 240.999; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 11.0 | b | +| 1970-01-01T00:01:40 | 4.0 | a | ++---------------------+------+------+ + +-- The window `(-300s, 0s]` includes the samples exactly at the anchor. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 0; + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 0.0 | a | +| 1970-01-01T00:01:40 | 7.0 | b | ++---------------------+-----+------+ + +-- 2. `@ start()`: every step selects its samples around the start of the evaluation range +-- (`(-200s, 100s]`), so every step reports the same value per series (1.0 for 'a', 8.0 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ start(); + ++---------------------+-----+------+ +| ts | val | host | ++---------------------+-----+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | +| 1970-01-01T00:03:20 | 1.0 | a | +| 1970-01-01T00:03:20 | 8.0 | b | +| 1970-01-01T00:05:00 | 1.0 | a | +| 1970-01-01T00:05:00 | 8.0 | b | +| 1970-01-01T00:06:40 | 1.0 | a | +| 1970-01-01T00:06:40 | 8.0 | b | ++---------------------+-----+------+ + +-- Control: without `@` every step looks back from its own timestamp, so the steps differ. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 8.0 | b | +| 1970-01-01T00:03:20 | 10.0 | b | +| 1970-01-01T00:03:20 | 3.0 | a | +| 1970-01-01T00:05:00 | 12.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:06:40 | 13.0 | b | +| 1970-01-01T00:06:40 | 6.0 | a | ++---------------------+------+------+ + +-- 3. `@ end()`: every step selects its samples around the end of the evaluation range +-- (`(100s, 400s]`), whose newest samples are 6.0 (a) and 13.0 (b) at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ end(); + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 13.0 | b | +| 1970-01-01T00:01:40 | 6.0 | a | +| 1970-01-01T00:03:20 | 13.0 | b | +| 1970-01-01T00:03:20 | 6.0 | a | +| 1970-01-01T00:05:00 | 13.0 | b | +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:40 | 13.0 | b | +| 1970-01-01T00:06:40 | 6.0 | a | ++---------------------+------+------+ + +-- 4. Range selector anchored by `@`: the call directly above the anchored range selector is +-- evaluated once, at the start of the evaluation, and its result is reported at every step — the +-- planner's counterpart of Prometheus' `StepInvariantExpr`. The window `(0s, 300s]` is left-open, +-- so it holds 60 -> 300 for host 'a' and 130 -> 250 for host 'b', i.e. 240 and 120 units. The 60s +-- gap to the window start is below the extrapolation threshold (1.1 x the 60s sampling interval), +-- so `rate` adds it and divides by the 300s window: 240 * 1.25 / 300 = 1.0 and +-- 120 * 1.25 / 300 = 0.5 per second. Every step must report those same per-series values, and both +-- series must be there at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 1.0 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 1.0 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 1.0 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- Consistency: the instant query below shows the anchored values over the same left-open +-- `(0s, 300s]` window (1.0 per second for 'a', 0.5 for 'b'). The `!= 1.0` comparison filters +-- host 'a', but host 'b' remains at every step with its rate of 0.5 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 300, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:05:00 | 0.5 | b | +| 1970-01-01T00:05:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300) != 1.0; + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.5 | b | ++---------------------+------------------------------------------+------+ + +-- Parentheses around the range argument are transparent: the call must be promoted exactly like +-- the unparenthesized one above, so every step reports the same anchored values for both series +-- (1.0 for host 'a', 0.5 for host 'b') instead of folding a window per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate((at_modifier_counter_total[5m] @ 300)); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 1.0 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 1.0 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 1.0 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 1.0 | a | ++---------------------+------------------------------------------+------+ + +-- The same holds for `@ start()` and `@ end()`, which are also fixed anchors: they resolve to the +-- statement's evaluation range, so a range call using them is evaluated once as well. +-- The statement starts at 60s: the window `(-240s, 60s]` holds only the samples at 0s and 60s, so +-- host 'a' goes 0 -> 60 and host 'b' 100 -> 130. The far 240s gap to the window start is beyond the +-- extrapolation threshold (1.1 x the 60s interval), so only half an interval (30s) is added there; +-- host 'a' starts its counter at 0, which clamps that back: 60 / 300 = 0.2 and +-- 30 * 1.5 / 300 = 0.15 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 300, '60s') rate(at_modifier_counter_total[5m] @ start()); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | +| 1970-01-01T00:02:00 | 0.15 | b | +| 1970-01-01T00:02:00 | 0.2 | a | +| 1970-01-01T00:03:00 | 0.15 | b | +| 1970-01-01T00:03:00 | 0.2 | a | +| 1970-01-01T00:04:00 | 0.15 | b | +| 1970-01-01T00:04:00 | 0.2 | a | +| 1970-01-01T00:05:00 | 0.15 | b | +| 1970-01-01T00:05:00 | 0.2 | a | ++---------------------+------------------------------------------+------+ + +-- The instant query at 60s folds the same window as every step of the query above. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 60, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | ++---------------------+------------------------------------------+------+ + +-- The statement ends at 240s: the window `(-60s, 240s]` holds the samples 0s to 240s, so host 'a' +-- goes 0 -> 240 and host 'b' 100 -> 220. The 60s gap to the window start is added for both, but +-- host 'a' starts at zero, which clamps it back (240 / 300 = 0.8 per second) while host 'b' keeps +-- it (120 * 1.25 / 300 = 0.5 per second). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ end()); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:00:00 | 0.5 | b | +| 1970-01-01T00:00:00 | 0.8 | a | +| 1970-01-01T00:01:00 | 0.5 | b | +| 1970-01-01T00:01:00 | 0.8 | a | +| 1970-01-01T00:02:00 | 0.5 | b | +| 1970-01-01T00:02:00 | 0.8 | a | +| 1970-01-01T00:03:00 | 0.5 | b | +| 1970-01-01T00:03:00 | 0.8 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (240, 240, '1s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- Control: without `@` each step folds its own window, so the values keep changing. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:01:00 | 0.15 | b | +| 1970-01-01T00:01:00 | 0.2 | a | +| 1970-01-01T00:02:00 | 0.25 | b | +| 1970-01-01T00:02:00 | 0.4 | a | +| 1970-01-01T00:03:00 | 0.35000000000000003 | b | +| 1970-01-01T00:03:00 | 0.6000000000000001 | a | +| 1970-01-01T00:04:00 | 0.5 | b | +| 1970-01-01T00:04:00 | 0.8 | a | ++---------------------+------------------------------------------+------+ + +-- `@ start()` on a range selector: the window `(-300s, 0s]` is left-open, so it holds a single +-- sample per series — the one exactly at the anchor — and `count_over_time` returns 1 at every step +-- for both of them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m] @ start()); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 1.0 | b | +| 1970-01-01T00:01:40 | 1.0 | a | +| 1970-01-01T00:01:40 | 1.0 | b | +| 1970-01-01T00:03:20 | 1.0 | a | +| 1970-01-01T00:03:20 | 1.0 | b | +| 1970-01-01T00:05:00 | 1.0 | a | +| 1970-01-01T00:05:00 | 1.0 | b | ++---------------------+------------------------------------+------+ + +-- Control: without `@` the window keeps moving, so the count grows with the step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 1.0 | b | +| 1970-01-01T00:01:40 | 2.0 | a | +| 1970-01-01T00:01:40 | 2.0 | b | +| 1970-01-01T00:03:20 | 4.0 | a | +| 1970-01-01T00:03:20 | 4.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:05:00 | 5.0 | b | ++---------------------+------------------------------------+------+ + +-- 5. `@` combined with `offset`: the offset moves the anchor backwards before the window is +-- selected, so `@ 300 offset 2m` reports the samples at 180s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 2m; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 10.0 | b | +| 1970-01-01T00:01:40 | 3.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 1m; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:01:40 | 11.0 | b | +| 1970-01-01T00:01:40 | 4.0 | a | ++---------------------+------+------+ + +-- Control: the offset applies to the evaluation timestamp when there is no `@`. +TQL EVAL (100, 100, '1s') at_modifier_gauge offset 2m; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +-- 6. `@` inside a larger expression. The anchored operand contributes the same value at every +-- step, while the plain operand still follows the evaluation step. A binary expression is never +-- promoted: when one side is anchored and the other is not, the anchored side replays on its own, +-- and when both sides are anchored, each side reports its own anchored samples. The join then runs +-- at every step over the replayed per-series inputs. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge @ 0; + ++------+---------------------+-------------------+ +| host | ts | lhs.val + rhs.val | ++------+---------------------+-------------------+ +| a | 1970-01-01T00:01:40 | 5.0 | +| a | 1970-01-01T00:03:20 | 5.0 | +| a | 1970-01-01T00:05:00 | 5.0 | +| a | 1970-01-01T00:06:40 | 5.0 | +| b | 1970-01-01T00:01:40 | 19.0 | +| b | 1970-01-01T00:03:20 | 19.0 | +| b | 1970-01-01T00:05:00 | 19.0 | +| b | 1970-01-01T00:06:40 | 19.0 | ++------+---------------------+-------------------+ + +-- The aggregation above the anchored selector runs at every step over the replayed (still +-- per-series) samples, so it sums both series at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') sum(at_modifier_gauge @ 0); + ++---------------------+----------------------------+ +| ts | sum(at_modifier_gauge.val) | ++---------------------+----------------------------+ +| 1970-01-01T00:01:40 | 7.0 | +| 1970-01-01T00:03:20 | 7.0 | +| 1970-01-01T00:05:00 | 7.0 | +| 1970-01-01T00:06:40 | 7.0 | ++---------------------+----------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 0); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:01:40 | 0.0 | a | +| 1970-01-01T00:01:40 | 7.0 | b | +| 1970-01-01T00:03:20 | 0.0 | a | +| 1970-01-01T00:03:20 | 7.0 | b | +| 1970-01-01T00:05:00 | 0.0 | a | +| 1970-01-01T00:05:00 | 7.0 | b | +| 1970-01-01T00:06:40 | 0.0 | a | +| 1970-01-01T00:06:40 | 7.0 | b | ++---------------------+----------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge; + ++------+---------------------+-------------------+ +| host | ts | lhs.val + rhs.val | ++------+---------------------+-------------------+ +| a | 1970-01-01T00:01:40 | 6.0 | +| a | 1970-01-01T00:03:20 | 8.0 | +| a | 1970-01-01T00:05:00 | 10.0 | +| a | 1970-01-01T00:06:40 | 11.0 | +| b | 1970-01-01T00:01:40 | 20.0 | +| b | 1970-01-01T00:03:20 | 22.0 | +| b | 1970-01-01T00:05:00 | 24.0 | +| b | 1970-01-01T00:06:40 | 25.0 | ++------+---------------------+-------------------+ + +-- 7. `@ -1`, an anchor before the Unix epoch, is accepted. Its window `(-301s, -1s]` holds no +-- sample, so the result is empty. +TQL EVAL (0, 0, '1s') at_modifier_gauge @ -1; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +-- 8. `@` on a metric that does not exist yields an empty result. +TQL EVAL (100, 100, '1s') at_modifier_missing @ 300; + ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ + +-- The same for a range selector over a missing metric: the anchored window is folded once by the +-- promoted call, but the scan has no row to fold, so the whole grid stays empty. +TQL EVAL (0, 240, '60s') rate(at_modifier_missing[5m] @ 300); + ++------+------------------------------------------------+ +| time | prom_rate(time_range,value,time,Int64(300000)) | ++------+------------------------------------------------+ ++------+------------------------------------------------+ + +-- 9. Multi-series coverage of the promoted call. The call above the anchored range selector is +-- evaluated once per series and its result is replayed per series, so grouping by `host` must +-- report both groups at every step: 1.0 for host 'a' and 0.5 for host 'b'. A promoted result that +-- treated the two series as one timeline would report the first step from the first row of the +-- batch and the later steps from its last row, dropping or mixing a host here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)); + ++------+---------------------+-----------------------------------------------+ +| host | ts | sum(prom_rate(ts_range,val,ts,Int64(300000))) | ++------+---------------------+-----------------------------------------------+ +| a | 1970-01-01T00:00:00 | 1.0 | +| a | 1970-01-01T00:01:00 | 1.0 | +| a | 1970-01-01T00:02:00 | 1.0 | +| a | 1970-01-01T00:03:00 | 1.0 | +| a | 1970-01-01T00:04:00 | 1.0 | +| b | 1970-01-01T00:00:00 | 0.5 | +| b | 1970-01-01T00:01:00 | 0.5 | +| b | 1970-01-01T00:02:00 | 0.5 | +| b | 1970-01-01T00:03:00 | 0.5 | +| b | 1970-01-01T00:04:00 | 0.5 | ++------+---------------------+-----------------------------------------------+ + +-- Consistency: comparing the per-host sequence above against 1.0 keeps exactly host 'b' at +-- every step, confirming it stays at 0.5 for the whole grid. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)) != 1.0; + ++------+---------------------+-----------------------------------------------+ +| host | ts | sum(prom_rate(ts_range,val,ts,Int64(300000))) | ++------+---------------------+-----------------------------------------------+ +| b | 1970-01-01T00:00:00 | 0.5 | +| b | 1970-01-01T00:01:00 | 0.5 | +| b | 1970-01-01T00:02:00 | 0.5 | +| b | 1970-01-01T00:03:00 | 0.5 | +| b | 1970-01-01T00:04:00 | 0.5 | ++------+---------------------+-----------------------------------------------+ + +-- A value function directly over an anchored instant selector is not promoted: the selector +-- anchors and replays its sample per series, and `abs` is evaluated at every step over that replay. +-- It must keep one row per series per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 300); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:01:40 | 12.0 | b | +| 1970-01-01T00:01:40 | 5.0 | a | +| 1970-01-01T00:03:20 | 12.0 | b | +| 1970-01-01T00:03:20 | 5.0 | a | +| 1970-01-01T00:05:00 | 12.0 | b | +| 1970-01-01T00:05:00 | 5.0 | a | +| 1970-01-01T00:06:40 | 12.0 | b | +| 1970-01-01T00:06:40 | 5.0 | a | ++---------------------+----------+------+ + +-- 10. `predict_linear` over an anchored range selector: the regression is centered on the +-- evaluation instant of the step, not on the last sample of the window. The window `(0s, 300s]` of +-- host 'a' is left-open (the sample at 0s is outside it) and rises by 1.0 per minute, so each step +-- predicts 60s ahead of the trend read at its own instant: 6.0 at 300s and 7.0 at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 360, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 300, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:00 | 7.0 | a | ++---------------------+-----------------------------------------------+------+ + +-- The anchor and the start of the evaluation differ here, so the window `(-60s, 240s]` — samples 0s +-- to 240s — is shifted by `at_offset` = 300s - 240s = 60s. Every step must follow the evaluation +-- grid (6.0 to 11.0); a regression read at the wrong instant predicts 60s ahead of the window's +-- last sample instead, i.e. 5.0 at every step — the sample at 240s is 4.0 and the trend rises by +-- 1.0 per minute — and shows up as a constant here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 600, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 240, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:00 | 6.0 | a | +| 1970-01-01T00:06:00 | 7.0 | a | +| 1970-01-01T00:07:00 | 8.0 | a | +| 1970-01-01T00:08:00 | 9.0 | a | +| 1970-01-01T00:09:00 | 10.0 | a | +| 1970-01-01T00:10:00 | 11.0 | a | ++---------------------+-----------------------------------------------+------+ + +-- Control: without `@` every step folds its own window, whose last sample is up to a step older +-- than the evaluation instant. Reading the regression off that sample reports 6.0 / 7.0 instead +-- of 6.5 / 7.5 here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m], 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:30 | 6.5 | a | +| 1970-01-01T00:06:30 | 7.5 | a | ++---------------------+-----------------------------------------------+------+ + +-- The same holds with an `offset`, which shifts the whole window backwards without moving the +-- evaluation instants: the window of step `T` is the left-open `(T - 1m - 5m, T - 1m]`, and the +-- trend is read at `T`, i.e. one minute past its newest sample. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset 1m, 60); + ++---------------------+-----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(60)) | host | ++---------------------+-----------------------------------------------+------+ +| 1970-01-01T00:05:30 | 6.5 | a | +| 1970-01-01T00:06:30 | 7.5 | a | ++---------------------+-----------------------------------------------+------+ + +-- `timestamp()` of an anchored instant selector is not promoted either: its argument is an +-- instant selector, not a range selector, so the selector anchors and replays its sample per series +-- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every +-- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects +-- the timestamp above it. +-- +-- Only an `@` without an `offset` is asserted here: the value of `timestamp()` for an anchored +-- selector that also carries an `offset` is a pre-existing question of its own, out of scope here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + ++---------------------+-------+------+ +| ts | value | host | ++---------------------+-------+------+ +| 1970-01-01T00:00:00 | 300.0 | a | +| 1970-01-01T00:00:00 | 300.0 | b | +| 1970-01-01T00:01:00 | 300.0 | a | +| 1970-01-01T00:01:00 | 300.0 | b | +| 1970-01-01T00:02:00 | 300.0 | a | +| 1970-01-01T00:02:00 | 300.0 | b | +| 1970-01-01T00:03:00 | 300.0 | a | +| 1970-01-01T00:03:00 | 300.0 | b | +| 1970-01-01T00:04:00 | 300.0 | a | +| 1970-01-01T00:04:00 | 300.0 | b | ++---------------------+-------+------+ + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: at_modifier_gauge.ts, value AS value, at_modifier_gauge.host | +| | Projection: at_modifier_gauge.ts, __promql_timestamp_value_ AS value, at_modifier_gauge.host | +| | PromInstantManipulate: range=[0..240000], lookback=[240001], interval=[60000], time index=[ts] | +| | PromInstantManipulate: range=[0..0], lookback=[300000], interval=[60000], time index=[ts] | +| | Projection: at_modifier_gauge.ts, at_modifier_gauge.val, at_modifier_gauge.host, CAST(CAST(CAST(CAST(CAST(at_modifier_gauge.ts AS Int64) AS Decimal128(19, 0)) AS Decimal128(21, 0)) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ | +| | PromSeriesNormalize: offset=[-300000], time index=[ts], filter NaN: [false] | +| | PromSeriesDivide: tags=["host"] | +| | Sort: at_modifier_gauge.host ASC NULLS FIRST, at_modifier_gauge.ts ASC NULLS FIRST | +| | Filter: at_modifier_gauge.ts >= TimestampMillisecond(1, None) AND at_modifier_gauge.ts <= TimestampMillisecond(300000, None) | +| | TableScan: at_modifier_gauge, partial_filters=[at_modifier_gauge.ts >= TimestampMillisecond(1, None), at_modifier_gauge.ts <= TimestampMillisecond(300000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- 11. `label_join` above an anchored selector. The call rewrites the label the input series are +-- told apart by (`host` becomes the empty string), so it is never promoted: it must stay above the +-- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, +-- with their own value, at every step. A promoted `label_join` would instead be replayed through +-- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row +-- per step. +-- +-- The invariant asserted here is scoped to the replay: it must not drop or mix the input rows. +-- It is not a claim about the final PromQL semantics of this query: joining `host` to one value +-- leaves two samples with the same label set at the same timestamp, which Prometheus rejects, +-- while `label_join` does not validate that yet (`label_replace` errors on such a rewrite +-- instead). That duplicate-labelset validation gap is pre-existing and out of scope here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "", ""); + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:05:00 | 12.0 | | +| 1970-01-01T00:05:00 | 5.0 | | +| 1970-01-01T00:06:00 | 12.0 | | +| 1970-01-01T00:06:00 | 5.0 | | +| 1970-01-01T00:07:00 | 12.0 | | +| 1970-01-01T00:07:00 | 5.0 | | +| 1970-01-01T00:08:00 | 12.0 | | +| 1970-01-01T00:08:00 | 5.0 | | ++---------------------+------+------+ + +-- The same join below another call: neither is promoted, and both are evaluated at every step over +-- the per-series replay of the anchored selector. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "", "")); + ++---------------------+----------+------+ +| ts | abs(val) | host | ++---------------------+----------+------+ +| 1970-01-01T00:05:00 | 12.0 | | +| 1970-01-01T00:05:00 | 5.0 | | +| 1970-01-01T00:06:00 | 12.0 | | +| 1970-01-01T00:06:00 | 5.0 | | +| 1970-01-01T00:07:00 | 12.0 | | +| 1970-01-01T00:07:00 | 5.0 | | +| 1970-01-01T00:08:00 | 12.0 | | +| 1970-01-01T00:08:00 | 5.0 | | ++---------------------+----------+------+ + +-- A range call below the join is still promoted on its own (it is the direct call over the +-- anchored range selector): the anchored window is folded once per series, and the join above that +-- replay reports the rate of both hosts at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "", ""); + ++---------------------+------------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | ++---------------------+------------------------------------------+------+ +| 1970-01-01T00:05:00 | 0.5 | | +| 1970-01-01T00:05:00 | 1.0 | | +| 1970-01-01T00:06:00 | 0.5 | | +| 1970-01-01T00:06:00 | 1.0 | | +| 1970-01-01T00:07:00 | 0.5 | | +| 1970-01-01T00:07:00 | 1.0 | | +| 1970-01-01T00:08:00 | 0.5 | | +| 1970-01-01T00:08:00 | 1.0 | | ++---------------------+------------------------------------------+------+ + +DROP TABLE at_modifier_gauge; + +Affected Rows: 0 + +DROP TABLE at_modifier_counter_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/at_modifier.sql b/tests/cases/standalone/common/promql/at_modifier.sql new file mode 100644 index 00000000000..9e9bc21a9c1 --- /dev/null +++ b/tests/cases/standalone/common/promql/at_modifier.sql @@ -0,0 +1,301 @@ +-- Tests for the PromQL `@` modifier on vector and matrix selectors. +-- +-- `@` anchors the sample selection window at a fixed timestamp instead of the timestamp of each +-- evaluation step: `@ ` uses the given timestamp, `@ start()` / `@ end()` use the +-- start/end of the statement's evaluation range, and `offset` shifts the anchor backwards before +-- the window is selected. The output timestamps still follow the evaluation grid. +-- +-- Every metric below carries two series, host 'a' and host 'b', with different values: the +-- anchored results must therefore report both of them, per step. A series silently dropped, or a +-- batch holding several series handled as one timeline, does not look like a correct single-series +-- answer here but shows up as a missing or shifted row. +-- +-- Sample timestamps below are in milliseconds, while `@` and `TQL EVAL` timestamps are in seconds, +-- as in Prometheus. + +CREATE TABLE at_modifier_gauge ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +CREATE TABLE at_modifier_counter_total ( + ts TIMESTAMP TIME INDEX, + val DOUBLE, + host STRING, + PRIMARY KEY(host) +); + +-- One sample per minute, so the anchored sample is easy to identify. The default lookback is 5m. +-- Host 'a' reports 0.0 to 6.0, host 'b' reports 7.0 to 13.0: the two series are told apart by their +-- values alone, which stays true after any anchoring. +INSERT INTO at_modifier_gauge VALUES + (0, 0.0, 'a'), + (60000, 1.0, 'a'), + (120000, 2.0, 'a'), + (180000, 3.0, 'a'), + (240000, 4.0, 'a'), + (300000, 5.0, 'a'), + (360000, 6.0, 'a'), + (0, 7.0, 'b'), + (60000, 8.0, 'b'), + (120000, 9.0, 'b'), + (180000, 10.0, 'b'), + (240000, 11.0, 'b'), + (300000, 12.0, 'b'), + (360000, 13.0, 'b'); + +-- Host 'a': a counter increasing by 60 per minute, i.e. by 1 per second. +-- Host 'b': a counter starting at 100 and increasing by 30 per minute, i.e. by 0.5 per second. Its +-- rate differs from host 'a', so a window folded for the wrong series cannot pass as the right one. +INSERT INTO at_modifier_counter_total VALUES + (0, 0.0, 'a'), + (60000, 60.0, 'a'), + (120000, 120.0, 'a'), + (180000, 180.0, 'a'), + (240000, 240.0, 'a'), + (300000, 300.0, 'a'), + (360000, 360.0, 'a'), + (0, 100.0, 'b'), + (60000, 130.0, 'b'), + (120000, 160.0, 'b'), + (180000, 190.0, 'b'), + (240000, 220.0, 'b'), + (300000, 250.0, 'b'), + (360000, 280.0, 'b'); + +-- 1. Instant query anchored at an absolute timestamp: the sample window is `(anchor - lookback, +-- anchor]` = `(0s, 300s]`, so its newest samples are 5.0 for host 'a' and 12.0 for host 'b', both +-- at 300s. The output timestamps stay the evaluation timestamps. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300; + +-- Control: the same query without `@` looks back from the evaluation timestamp instead: +-- `(-200s, 100s]` selects the samples 1.0 (a) and 8.0 (b) at 60s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge; + +-- The anchor, not the evaluation timestamp, decides which sample is reported. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 360; + +-- Sub-second precision is preserved in the anchor: `240.999` selects the samples at 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 240.999; + +-- The window `(-300s, 0s]` includes the samples exactly at the anchor. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 0; + +-- 2. `@ start()`: every step selects its samples around the start of the evaluation range +-- (`(-200s, 100s]`), so every step reports the same value per series (1.0 for 'a', 8.0 for 'b'). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ start(); + +-- Control: without `@` every step looks back from its own timestamp, so the steps differ. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge; + +-- 3. `@ end()`: every step selects its samples around the end of the evaluation range +-- (`(100s, 400s]`), whose newest samples are 6.0 (a) and 13.0 (b) at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ end(); + +-- 4. Range selector anchored by `@`: the call directly above the anchored range selector is +-- evaluated once, at the start of the evaluation, and its result is reported at every step — the +-- planner's counterpart of Prometheus' `StepInvariantExpr`. The window `(0s, 300s]` is left-open, +-- so it holds 60 -> 300 for host 'a' and 130 -> 250 for host 'b', i.e. 240 and 120 units. The 60s +-- gap to the window start is below the extrapolation threshold (1.1 x the 60s sampling interval), +-- so `rate` adds it and divides by the 300s window: 240 * 1.25 / 300 = 1.0 and +-- 120 * 1.25 / 300 = 0.5 per second. Every step must report those same per-series values, and both +-- series must be there at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300); + +-- Consistency: the instant query below shows the anchored values over the same left-open +-- `(0s, 300s]` window (1.0 per second for 'a', 0.5 for 'b'). The `!= 1.0` comparison filters +-- host 'a', but host 'b' remains at every step with its rate of 0.5 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 300, '1s') rate(at_modifier_counter_total[5m]); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ 300) != 1.0; + +-- Parentheses around the range argument are transparent: the call must be promoted exactly like +-- the unparenthesized one above, so every step reports the same anchored values for both series +-- (1.0 for host 'a', 0.5 for host 'b') instead of folding a window per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate((at_modifier_counter_total[5m] @ 300)); + +-- The same holds for `@ start()` and `@ end()`, which are also fixed anchors: they resolve to the +-- statement's evaluation range, so a range call using them is evaluated once as well. +-- The statement starts at 60s: the window `(-240s, 60s]` holds only the samples at 0s and 60s, so +-- host 'a' goes 0 -> 60 and host 'b' 100 -> 130. The far 240s gap to the window start is beyond the +-- extrapolation threshold (1.1 x the 60s interval), so only half an interval (30s) is added there; +-- host 'a' starts its counter at 0, which clamps that back: 60 / 300 = 0.2 and +-- 30 * 1.5 / 300 = 0.15 per second. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 300, '60s') rate(at_modifier_counter_total[5m] @ start()); + +-- The instant query at 60s folds the same window as every step of the query above. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (60, 60, '1s') rate(at_modifier_counter_total[5m]); + +-- The statement ends at 240s: the window `(-60s, 240s]` holds the samples 0s to 240s, so host 'a' +-- goes 0 -> 240 and host 'b' 100 -> 220. The 60s gap to the window start is added for both, but +-- host 'a' starts at zero, which clamps it back (240 / 300 = 0.8 per second) while host 'b' keeps +-- it (120 * 1.25 / 300 = 0.5 per second). +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m] @ end()); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (240, 240, '1s') rate(at_modifier_counter_total[5m]); + +-- Control: without `@` each step folds its own window, so the values keep changing. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') rate(at_modifier_counter_total[5m]); + +-- `@ start()` on a range selector: the window `(-300s, 0s]` is left-open, so it holds a single +-- sample per series — the one exactly at the anchor — and `count_over_time` returns 1 at every step +-- for both of them. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m] @ start()); + +-- Control: without `@` the window keeps moving, so the count grows with the step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 300, '100s') count_over_time(at_modifier_counter_total[5m]); + +-- 5. `@` combined with `offset`: the offset moves the anchor backwards before the window is +-- selected, so `@ 300 offset 2m` reports the samples at 180s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 2m; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 100, '1s') at_modifier_gauge @ 300 offset 1m; + +-- Control: the offset applies to the evaluation timestamp when there is no `@`. +TQL EVAL (100, 100, '1s') at_modifier_gauge offset 2m; + +-- 6. `@` inside a larger expression. The anchored operand contributes the same value at every +-- step, while the plain operand still follows the evaluation step. A binary expression is never +-- promoted: when one side is anchored and the other is not, the anchored side replays on its own, +-- and when both sides are anchored, each side reports its own anchored samples. The join then runs +-- at every step over the replayed per-series inputs. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge @ 0; + +-- The aggregation above the anchored selector runs at every step over the replayed (still +-- per-series) samples, so it sums both series at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') sum(at_modifier_gauge @ 0); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 0); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') at_modifier_gauge @ 300 + at_modifier_gauge; + +-- 7. `@ -1`, an anchor before the Unix epoch, is accepted. Its window `(-301s, -1s]` holds no +-- sample, so the result is empty. +TQL EVAL (0, 0, '1s') at_modifier_gauge @ -1; + +-- 8. `@` on a metric that does not exist yields an empty result. +TQL EVAL (100, 100, '1s') at_modifier_missing @ 300; + +-- The same for a range selector over a missing metric: the anchored window is folded once by the +-- promoted call, but the scan has no row to fold, so the whole grid stays empty. +TQL EVAL (0, 240, '60s') rate(at_modifier_missing[5m] @ 300); + +-- 9. Multi-series coverage of the promoted call. The call above the anchored range selector is +-- evaluated once per series and its result is replayed per series, so grouping by `host` must +-- report both groups at every step: 1.0 for host 'a' and 0.5 for host 'b'. A promoted result that +-- treated the two series as one timeline would report the first step from the first row of the +-- batch and the later steps from its last row, dropping or mixing a host here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)); + +-- Consistency: comparing the per-host sequence above against 1.0 keeps exactly host 'b' at +-- every step, confirming it stays at 0.5 for the whole grid. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') sum by (host) (rate(at_modifier_counter_total[5m] @ 300)) != 1.0; + +-- A value function directly over an anchored instant selector is not promoted: the selector +-- anchors and replays its sample per series, and `abs` is evaluated at every step over that replay. +-- It must keep one row per series per step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (100, 400, '100s') abs(at_modifier_gauge @ 300); + +-- 10. `predict_linear` over an anchored range selector: the regression is centered on the +-- evaluation instant of the step, not on the last sample of the window. The window `(0s, 300s]` of +-- host 'a' is left-open (the sample at 0s is outside it) and rises by 1.0 per minute, so each step +-- predicts 60s ahead of the trend read at its own instant: 6.0 at 300s and 7.0 at 360s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 360, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 300, 60); + +-- The anchor and the start of the evaluation differ here, so the window `(-60s, 240s]` — samples 0s +-- to 240s — is shifted by `at_offset` = 300s - 240s = 60s. Every step must follow the evaluation +-- grid (6.0 to 11.0); a regression read at the wrong instant predicts 60s ahead of the window's +-- last sample instead, i.e. 5.0 at every step — the sample at 240s is 4.0 and the trend rises by +-- 1.0 per minute — and shows up as a constant here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 600, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] @ 240, 60); + +-- Control: without `@` every step folds its own window, whose last sample is up to a step older +-- than the evaluation instant. Reading the regression off that sample reports 6.0 / 7.0 instead +-- of 6.5 / 7.5 here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m], 60); + +-- The same holds with an `offset`, which shifts the whole window backwards without moving the +-- evaluation instants: the window of step `T` is the left-open `(T - 1m - 5m, T - 1m]`, and the +-- trend is read at `T`, i.e. one minute past its newest sample. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset 1m, 60); + +-- `timestamp()` of an anchored instant selector is not promoted either: its argument is an +-- instant selector, not a range selector, so the selector anchors and replays its sample per series +-- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every +-- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects +-- the timestamp above it. +-- +-- Only an `@` without an `offset` is asserted here: the value of `timestamp()` for an anchored +-- selector that also carries an `offset` is a pre-existing question of its own, out of scope here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); + +-- 11. `label_join` above an anchored selector. The call rewrites the label the input series are +-- told apart by (`host` becomes the empty string), so it is never promoted: it must stay above the +-- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, +-- with their own value, at every step. A promoted `label_join` would instead be replayed through +-- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row +-- per step. +-- +-- The invariant asserted here is scoped to the replay: it must not drop or mix the input rows. +-- It is not a claim about the final PromQL semantics of this query: joining `host` to one value +-- leaves two samples with the same label set at the same timestamp, which Prometheus rejects, +-- while `label_join` does not validate that yet (`label_replace` errors on such a rewrite +-- instead). That duplicate-labelset validation gap is pre-existing and out of scope here. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "", ""); + +-- The same join below another call: neither is promoted, and both are evaluated at every step over +-- the per-series replay of the anchored selector. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "", "")); + +-- A range call below the join is still promoted on its own (it is the direct call over the +-- anchored range selector): the anchored window is folded once per series, and the join above that +-- replay reports the rate of both hosts at every step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "", ""); + +DROP TABLE at_modifier_gauge; + +DROP TABLE at_modifier_counter_total; From fedce5c5ec1e458483901e2dc75dab6d29b741cf Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 28 Sep 2026 12:27:36 +0000 Subject: [PATCH 060/133] feat: support non-millisecond time index units in the logical batcher (#9346) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: allow customized time index unit for metric engine table * test: provide query tests * refactor: revert unnecessary change * refactor: share timestamp unit conversions in api helper Address review feedback on the time index unit changeset: - Add shared timestamp_unit/timestamp_datatype helpers to api::helper (the only crate that sees both proto ColumnDataType and TimeUnit due to layering; common-time and datatypes have no greptime-proto dep). This removes the ColumnDataType -> TimeUnit match duplicated between operator's insert path and the OTLP logs path. - Collapse the two TimeUnit <-> ValueData matches in convert_timestamp_value_data by reusing api::helper::to_grpc_value for the construction side. - Note that convert_rows_time_unit rewrites the schema before the values, so an overflow mid-batch leaves the request half-converted; harmless because the error aborts the whole insert request. Signed-off-by: Ning Sun * fix: align time units per destination table and floor remote-read timestamps Address review feedback on PR #9236: - Align each metric insert request to the unit of the table it actually targets: an existing logical table keeps its own unit (it may be bound to a different physical table than the one selected by the request), and only new tables use the selected physical table's unit. The previous blanket conversion rewrote valid millisecond samples to the selected physical table's unit and the engine rejected them. Regression test: writing an existing millisecond logical table and a new table in one request that selects a microsecond physical table. - Remote read now floors narrowing timestamp conversions towards negative infinity (div_euclid), consistent with Timestamp::convert_to on the ingestion path; arrow's cast truncates towards zero and returned -1ms for a stored -1001us. Widening (second -> millisecond) keeps the exact arrow cast. Regression test: a negative, non-aligned timestamp round-trips as -2ms. Signed-off-by: Ning Sun * perf: fold time unit alignment into existing table lookups Address review feedback on PR #9236: - The per-destination unit alignment no longer runs its own pass of table lookups: create_or_alter_tables_on_demand gains an align_time_index_unit parameter (metric engine path only) and converts each request inside the lookups it already performs — existing tables to their own unit, new tables to the selected physical table's. Default ingest paths now issue zero additional catalog lookups compared to main; the separate alignment pass remains only in the opt-in logical batcher pre-gate, next to the eligibility check that already looks up the same tables. - convert_rows_time_unit indexes the time index position directly (validate_column_count_match guarantees row widths) instead of Optional get_mut; the gate-side alignment validates widths itself. Signed-off-by: Ning Sun * perf: resolve the batcher time index guard once per write target All batches of one remote write request share the same write target (catalog, schema, physical table), so the batcher time index guard now resolves each distinct target once instead of once per batch. Signed-off-by: Ning Sun * feat: support non-millisecond time index units in the logical batcher Make the logical table batcher's bulk encode path unit-aware so physical metric tables with a non-millisecond time index (e.g. TIMESTAMP(6)) can use logical batching instead of falling back to the ordinary insert path. - rows_to_aligned_record_batch builds the time index column in the TARGET schema's unit, converting any timestamp encoding via Timestamp::convert_to (flooring on narrowing, consistent with the ordinary insert path). - New tables created by the batcher use the selected physical table's time index unit (resolved once per submit; a missing physical table keeps the millisecond auto-create default). - columns_taxonomy and the can_batch_metric_rows schema whitelist accept any timestamp unit; the prometheus remote write v1/v2 batcher gates and the OTLP pre-gate alignment are removed together with Inserter::align_metric_row_inserts_time_unit, as the batcher now converts internally. Closes #9342 Signed-off-by: Ning Sun * fix: address review comments * fix: address review issue * refactor: drop the OTLP pre-gate unit alignment made redundant by the bulk path The main merge of #9236 (squash) resurrected the OTLP pre-gate alignment and Inserter::align_metric_row_inserts_time_unit, which this branch had removed. Drop them again: - The pre-gate existed because the #9236-era bulk eligibility gate only accepted millisecond schemas, so nanosecond-encoded OTLP requests had to be converted before the check. This branch makes the bulk path unit-aware (the gate accepts all time index units and batch alignment converts each request to its destination's unit), so the pre-gate is redundant and only added N+1 catalog lookups per batched request — the very lookup-count overhead raised in the #9236 review. - The per-destination unit semantics it implemented remain enforced in the two paths that need them: the ordinary insert path (create_or_alter_tables_on_demand converts inside its existing table lookups) and the batched path (batch alignment resolves each destination schema and converts to it). test_otlp_logical_batcher_alignment (the test the pre-gate originally fixed) and the mixed-physical-table regression both pass without it. Signed-off-by: Ning Sun * test: cover OTLP batcher cross-physical fallback and nanosecond physical Extend the logical batcher integration coverage for the cases previously guarded by the removed OTLP pre-gate alignment: - test_otlp_logical_batcher_fallback_for_cross_physical_destination: with the batcher enabled, an OTLP request targeting an existing logical table bound to another physical table must NOT enter the batcher (the bulk eligibility check rejects the destination binding) and the ordinary insert path must convert it to the destination's unit (60s -> 60_000_000us). - test_otlp_logical_batcher_non_millisecond_physical_table now covers both microsecond and nanosecond physical tables (parameterized), asserting batcher submissions and unit-precise stored values. Signed-off-by: Ning Sun * fix: validate per-batch physical bindings and avoid intermediate timestamp buffers Address review feedback on PR #9346: - accepts_bulk_destinations dedupes on (schema, table, selected physical) instead of (schema, table): one request can select different physical tables per batch (per-series x_greptime_physical_table labels), and the old key let a second selection skip validation and flush rows through the wrong physical's regions. Missing tables additionally reject conflicting physical selections within the same request. Regression test covers an existing destination, a missing destination, and a consistent selection (which must still batch). - The timestamp column builder appends each value directly into the target-unit Arrow builder; values already in the target unit (the unchanged millisecond fast path) are appended without conversion, so the default millisecond physical pays no Timestamp construction or intermediate Vec allocation. - The non-millisecond batching tests assert the submit_build_and_align counter, which increments on every batcher submission in both acknowledgement modes, so a silent fallback to ordinary insertion fails the tests instead of passing on stored values alone. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/frontend/src/instance/otlp.rs | 28 +- src/operator/src/insert.rs | 141 +-------- src/servers/src/batcher/logical_table.rs | 148 +++++----- .../batcher/logical_table/batch_convert.rs | 4 +- .../src/batcher/logical_table/tables.rs | 25 +- src/servers/src/http/prom_store.rs | 20 +- src/servers/src/prom_row_builder.rs | 274 +++++++++++++++--- tests-integration/src/otlp.rs | 263 +++++++++++++++++ tests-integration/tests/http.rs | 234 ++++++++++++++- 9 files changed, 855 insertions(+), 282 deletions(-) diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index e82aa343870..5132b9f8a8d 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -124,7 +124,7 @@ impl OpenTelemetryProtocolHandler for Instance { metric_ctx.resource_info = self.otlp_resource_info; let otlp::metrics::MetricsConversion { - mut requests, + requests, rows, semantic_index, resource_info, @@ -170,33 +170,21 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); + // The bulk path converts each request's time index unit to the + // destination table's unit during batch alignment, so no pre-gate + // alignment is needed here. let batcher = self.logical_batcher().filter(|_| { ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine }); - let batcher = if batcher.is_some() { - // Align the requests' time index unit with the physical table's - // before the bulk eligibility check: the OTLP encoder keeps - // nanosecond precision on the metric engine path, while the bulk - // path assumes the physical table's unit (today millisecond). - // Without this, nanosecond requests would silently skip the - // batcher, and a non-millisecond physical table must not enter - // the bulk path either (the eligibility check rejects its unit). - self.inserter - .align_metric_row_inserts_time_unit(&ctx, &physical_table, &mut requests) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu)?; - if self + 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 - } + { + batcher } else { None }; diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 24acdeaba28..a0ebe236767 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -172,14 +172,19 @@ impl Inserter { return Ok(false); } for request in &requests.inserts { - // The logical bulk encoder only supports scalar metric schemas. + // The logical bulk encoder only supports scalar metric schemas; + // any time index unit is accepted — requests are converted to the + // destination table's unit during batch alignment. // Check new tables too, before catalog lookup or schema changes. if request.rows.as_ref().is_some_and(|rows| { rows.schema.iter().any(|column| { column.datatype_extension.is_some() || !matches!( ColumnDataType::try_from(column.datatype), - Ok(ColumnDataType::TimestampMillisecond + Ok(ColumnDataType::TimestampSecond + | ColumnDataType::TimestampMillisecond + | ColumnDataType::TimestampMicrosecond + | ColumnDataType::TimestampNanosecond | ColumnDataType::Float64 | ColumnDataType::String) ) @@ -1376,75 +1381,6 @@ impl Inserter { } } - /// Aligns each request's time index unit with the unit of the table the - /// request targets, so that write-path gates (e.g. the logical batcher - /// eligibility check) observe the destination table's unit instead of the - /// ingestion endpoint's encoding unit (Prometheus remote write always - /// uses millisecond; OTLP keeps nanosecond precision on the metric engine - /// path). Existing destination tables keep their own unit — they may be - /// bound to a different physical table than `physical_table` — while new - /// tables use `physical_table`'s unit. A missing physical table is - /// treated as the millisecond unit of the auto-created default; the - /// actual creation (if any) happens later in - /// [`Inserter::handle_metric_row_inserts`]. - /// - /// This is only needed on paths that choose a write path before - /// `handle_metric_row_inserts` (whose own table lookups perform the - /// alignment again, as a no-op after this). - pub async fn align_metric_row_inserts_time_unit( - &self, - ctx: &QueryContextRef, - physical_table: &str, - requests: &mut RowInsertRequests, - ) -> Result<()> { - // The unit conversion indexes rows by the time index position, which - // requires well-formed requests. - validate_column_count_match(requests)?; - let physical_unit = match self - .get_table(ctx.current_catalog(), &ctx.current_schema(), physical_table) - .await? - { - Some(table) => table_time_index_unit(&table), - // A missing physical table is auto-created with the millisecond - // unit later. - None => Some(TimeUnit::Millisecond), - }; - self.align_metric_rows_per_destination(ctx, physical_unit, requests) - .await - } - - /// Converts each request's time index column to the unit of the table the - /// request targets: the existing table's own unit when the table exists, - /// and `physical_unit` (of the request's selected physical table) for new - /// tables. - async fn align_metric_rows_per_destination( - &self, - ctx: &QueryContextRef, - physical_unit: Option, - requests: &mut RowInsertRequests, - ) -> Result<()> { - for request in &mut requests.inserts { - let Some(rows) = request.rows.as_mut() else { - continue; - }; - let target_unit = match self - .get_table( - ctx.current_catalog(), - &ctx.current_schema(), - &request.table_name, - ) - .await? - { - Some(table) => table_time_index_unit(&table), - None => physical_unit, - }; - if let Some(target_unit) = target_unit { - convert_rows_time_unit(rows, target_unit)?; - } - } - Ok(()) - } - async fn get_table( &self, catalog: &str, @@ -2373,69 +2309,6 @@ mod tests { ); } - #[tokio::test] - async fn test_align_metric_rows_per_destination() { - use catalog::RegisterTableRequest; - use catalog::memory::MemoryCatalogManager; - - // An existing millisecond logical table `existing`, plus a - // microsecond physical table selected by the request: the existing - // table's request keeps the millisecond unit (it is bound to another - // physical table), while the new table's request is converted to the - // selected physical table's microsecond unit. - let mut inserter = batcher_test_inserter().await; - let existing = make_metric_physical_table_ref_with_time_unit(TimeUnit::Millisecond); - let phy_us = make_metric_physical_table_ref_with_time_unit(TimeUnit::Microsecond); - let catalog = MemoryCatalogManager::with_default_setup(); - for (table_name, table_id, table) in [("existing", 1, existing), ("phy_us", 2, phy_us)] { - catalog - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - table_id, - table, - }) - .unwrap(); - } - inserter.catalog_manager = catalog; - - let ctx = Arc::new(QueryContext::with( - DEFAULT_CATALOG_NAME, - DEFAULT_SCHEMA_NAME, - )); - let mut requests = RowInsertRequests { - inserts: vec![ - ms_row_insert_request_named("existing", 123), - ms_row_insert_request_named("fresh", 123), - ], - }; - inserter - .align_metric_row_inserts_time_unit(&ctx, "phy_us", &mut requests) - .await - .unwrap(); - - let existing_rows = requests.inserts[0].rows.as_ref().unwrap(); - assert_eq!( - existing_rows.schema[0].datatype, - ColumnDataType::TimestampMillisecond as i32 - ); - assert!(matches!( - existing_rows.rows[0].values[0].value_data, - Some(ValueData::TimestampMillisecondValue(123)) - )); - - let fresh_rows = requests.inserts[1].rows.as_ref().unwrap(); - assert_eq!( - fresh_rows.schema[0].datatype, - ColumnDataType::TimestampMicrosecond as i32 - ); - assert!(matches!( - fresh_rows.rows[0].values[0].value_data, - Some(ValueData::TimestampMicrosecondValue(123_000)) - )); - } - #[tokio::test] async fn test_accommodate_existing_schema_and_reject_kind_changes() { let ts_name = "my_ts"; diff --git a/src/servers/src/batcher/logical_table.rs b/src/servers/src/batcher/logical_table.rs index fcd4e1ea38f..7e0cea44b6a 100644 --- a/src/servers/src/batcher/logical_table.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -21,7 +21,7 @@ mod tables; #[cfg(test)] mod test_util; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::{Future, ready}; use std::num::NonZeroUsize; use std::sync::Arc; @@ -42,6 +42,7 @@ use meter_macros::write_meter; use partition::manager::PartitionRuleManagerRef; use session::context::QueryContextRef; use snafu::ResultExt; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; use tokio::sync::{Semaphore, broadcast, mpsc, oneshot}; use crate::batcher::flow_notifier::{FlowNotifier, start_flow_notification_worker}; @@ -173,79 +174,17 @@ impl LogicalTablePendingRowsBatcher { .map(|(rows, ())| rows) } - /// Returns whether the bulk path can accept `batches`: - /// - the physical metric table selected by each context must have a - /// millisecond time index (or not exist yet — it is auto-created with - /// the millisecond unit): new tables are created on it and the bulk - /// encode produces millisecond batches; - /// - every existing destination table must itself have a millisecond time - /// index: it may be bound to another physical table with another unit, - /// and the bulk encoder would build millisecond arrays against its - /// schema. - /// - /// Incompatible requests must stay on the ordinary insert path, which - /// converts the requests to each destination table's unit. Write targets - /// and destination tables are resolved once per distinct name. - pub(crate) async fn accepts_bulk_time_indexes( - &self, - batches: impl Iterator, - ) -> bool { - let mut checked_targets = HashSet::new(); - let mut checked_tables = HashSet::new(); - for (ctx, requests) in batches { - let key = batch_key_from_ctx(ctx); - if checked_targets.insert((key.catalog, key.schema, key.physical_table)) - && !self.accepts_physical_table_time_index(ctx).await - { - return false; - } - for request in &requests.inserts { - if !checked_tables.insert((ctx.current_schema(), request.table_name.clone())) { - continue; - } - let Ok(Some(table)) = self - .catalog_manager - .table( - ctx.current_catalog(), - &ctx.current_schema(), - &request.table_name, - None, - ) - .await - else { - // New table: governed by the selected physical table's - // unit, checked above. - continue; - }; - if table - .table_info() - .meta - .schema - .timestamp_column() - .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) - .is_none_or(|unit| unit != TimeUnit::Millisecond) - { - return false; - } - } - } - true - } - - /// Returns whether the physical metric table resolved from `ctx` can use - /// the bulk path. The bulk encode produces millisecond timestamp batches - /// only, so a physical table with another time index unit (e.g. created - /// as `TIMESTAMP(6)`) must stay on the ordinary insert path, which - /// converts the requests to the physical table's unit. A missing physical - /// table is accepted: it is auto-created with the millisecond unit. - async fn accepts_physical_table_time_index(&self, ctx: &QueryContextRef) -> bool { + /// Returns the physical metric table's time index unit resolved from + /// `ctx`, defaulting to millisecond when the table does not exist yet + /// (the schema alterer auto-creates it as millisecond). + async fn physical_time_index_unit_or_default(&self, ctx: &QueryContextRef) -> TimeUnit { let key = batch_key_from_ctx(ctx); let Ok(Some(table)) = self .catalog_manager .table(&key.catalog, &key.schema, &key.physical_table, None) .await else { - return true; + return TimeUnit::Millisecond; }; table .table_info() @@ -253,8 +192,77 @@ impl LogicalTablePendingRowsBatcher { .schema .timestamp_column() .and_then(|col| col.data_type.as_timestamp().map(|ts| ts.unit())) - .map(|unit| unit == TimeUnit::Millisecond) - .unwrap_or(true) + .unwrap_or(TimeUnit::Millisecond) + } + + /// Returns whether the bulk path can accept `batches`: every existing + /// destination table must be a metric logical table bound to the + /// physical table selected by its context. A destination bound to + /// another physical table would be flushed through the selected + /// physical's regions, silently misplacing its rows, so such requests + /// must stay on the ordinary insert path (which routes per destination). + /// New tables are always fine: they are created on the selected physical + /// table. Time index units need no check here — the bulk encode converts + /// each request to its destination's unit. Destinations are resolved + /// once per distinct (schema, table). + pub(crate) async fn accepts_bulk_destinations( + &self, + batches: impl Iterator, + ) -> bool { + // One request can select different physical tables per batch (e.g. + // per-series physical-table labels), so the dedupe key includes the + // selected physical table: every distinct (schema, table, physical) + // triple is validated against the table's actual binding. + let mut checked = HashSet::new(); + // For missing tables, one request must not select two different + // physical tables: the batcher would create the table through one + // selection and flush its rows through the other's regions. + let mut missing_selections: HashMap<(String, String), String> = HashMap::new(); + for (ctx, requests) in batches { + let physical_table = batch_key_from_ctx(ctx).physical_table; + let schema = ctx.current_schema(); + for request in &requests.inserts { + if !checked.insert(( + schema.clone(), + request.table_name.clone(), + physical_table.clone(), + )) { + continue; + } + let Ok(Some(table)) = self + .catalog_manager + .table(ctx.current_catalog(), &schema, &request.table_name, None) + .await + else { + // New table: created on the selected physical table, but + // a conflicting selection within the same request cannot + // be batched. + if missing_selections + .insert( + (schema.clone(), request.table_name.clone()), + physical_table.clone(), + ) + .is_some_and(|previous| previous != physical_table) + { + return false; + } + continue; + }; + let info = table.table_info(); + if info.meta.engine != METRIC_ENGINE_NAME + || info + .meta + .options + .extra_options + .get(LOGICAL_TABLE_METADATA_KEY) + .map(String::as_str) + != Some(physical_table.as_str()) + { + return false; + } + } + } + true } /// Submits with request-level accounting after schema preparation and before diff --git a/src/servers/src/batcher/logical_table/batch_convert.rs b/src/servers/src/batcher/logical_table/batch_convert.rs index 28464eb2d3b..37080dbe258 100644 --- a/src/servers/src/batcher/logical_table/batch_convert.rs +++ b/src/servers/src/batcher/logical_table/batch_convert.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use arrow::compute::concat_batches; -use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use common_query::prelude::{greptime_timestamp, greptime_value}; use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; @@ -130,7 +130,7 @@ pub(in crate::batcher::logical_table) fn columns_taxonomy( essential_column_indices.push(index); } } - ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => { + ArrowDataType::Timestamp(_, _) => { ensure!( timestamp_index.replace(index).is_none(), error::InvalidPromRemoteRequestSnafu { diff --git a/src/servers/src/batcher/logical_table/tables.rs b/src/servers/src/batcher/logical_table/tables.rs index 87110366a8c..140dd2d4f23 100644 --- a/src/servers/src/batcher/logical_table/tables.rs +++ b/src/servers/src/batcher/logical_table/tables.rs @@ -15,9 +15,10 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use api::v1::{ColumnSchema, RowInsertRequests, Rows}; +use api::v1::{ColumnSchema, RowInsertRequests, Rows, SemanticType}; use arrow::datatypes::Schema as ArrowSchema; use async_trait::async_trait; +use common_time::timestamp::TimeUnit; use session::context::QueryContextRef; use snafu::OptionExt; @@ -91,6 +92,17 @@ impl LogicalTablePendingRowsBatcher { .plan_table_resolution(&catalog, &schema, ctx, &unique_tables) .await?; + // New tables are created on the request's selected physical table; + // their time index must use the physical table's unit (a missing + // physical table is auto-created as millisecond by the schema + // alterer, matching the default here). + if !plan.tables_to_create.is_empty() { + let physical_unit = self.physical_time_index_unit_or_default(ctx).await; + for (_, request_schema) in &mut plan.tables_to_create { + align_create_schema_time_index(request_schema, physical_unit); + } + } + self.create_missing_tables_and_refresh_schemas( &catalog, &schema, @@ -109,6 +121,17 @@ impl LogicalTablePendingRowsBatcher { } } +/// Rewrites the time index column of a create-table schema to `unit`, if the +/// schema carries a timestamp column in another unit. +fn align_create_schema_time_index(request_schema: &mut [ColumnSchema], unit: TimeUnit) { + for column in request_schema { + if column.semantic_type == SemanticType::Timestamp as i32 { + column.datatype = api::helper::timestamp_datatype(unit) as i32; + column.datatype_extension = None; + } + } +} + impl LogicalTablePendingRowsBatcher { /// Extracts non-empty `(table_name, rows)` pairs and computes total row /// count across the retained entries. diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 44abb18a5be..68542a6f642 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -408,11 +408,11 @@ async fn write_prometheus_rows_with_progress( error, rows_written: 0, })?; - // The bulk encode produces millisecond batches only; write targets - // and existing destination tables with another time index unit stay - // on the ordinary insert path, which converts the requests to each - // destination table's unit. - if batcher.accepts_bulk_time_indexes(batches.iter()).await { + // Destinations bound to another physical table must stay on the + // ordinary insert path, which routes per destination; time index + // units need no check — the bulk encode converts each request to + // its destination's unit. + if batcher.accepts_bulk_destinations(batches.iter()).await { let mut rows_written = 0; for (temp_ctx, reqs) in batches { let rows = @@ -543,11 +543,11 @@ async fn write_prometheus_v2_rows_with_progress( samples_written: 0, histograms_written: 0, })?; - // The bulk encode produces millisecond batches only; write targets - // and existing destination tables with another time index unit stay - // on the ordinary insert path, which converts the requests to each - // destination table's unit. - if batcher.accepts_bulk_time_indexes(batches.iter()).await { + // Destinations bound to another physical table must stay on the + // ordinary insert path, which routes per destination; time index + // units need no check — the bulk encode converts each request to + // its destination's unit. + if batcher.accepts_bulk_destinations(batches.iter()).await { return write_batched_prometheus_v2_rows_with_progress( prom_store_handler, batcher.as_ref(), diff --git a/src/servers/src/prom_row_builder.rs b/src/servers/src/prom_row_builder.rs index 8b4e35bcd32..43b3a464857 100644 --- a/src/servers/src/prom_row_builder.rs +++ b/src/servers/src/prom_row_builder.rs @@ -23,14 +23,17 @@ use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Rows, SemanticType}; use arrow::array::{ - ArrayRef, Float64Builder, StringBuilder, TimestampMicrosecondBuilder, - TimestampMillisecondBuilder, TimestampNanosecondBuilder, TimestampSecondBuilder, - new_null_array, + ArrayRef, ArrowPrimitiveType, Float64Builder, PrimitiveBuilder, StringBuilder, new_null_array, +}; +use arrow::datatypes::{ + DataType as ArrowDataType, Schema as ArrowSchema, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; use arrow_schema::TimeUnit; use common_query::prelude::{greptime_timestamp, greptime_value}; +use common_time::Timestamp; +use common_time::timestamp::TimeUnit as CommonTimeUnit; use datatypes::data_type::DataType; use datatypes::prelude::ConcreteDataType; use snafu::{OptionExt, ResultExt, ensure}; @@ -129,16 +132,7 @@ pub(crate) fn rows_to_aligned_record_batch( ArrowDataType::Float64 => { source_map.insert(&target_field_name, (src_idx, src_arrow_type)); } - ArrowDataType::Timestamp(unit, _) => { - ensure!( - unit == &TimeUnit::Millisecond, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Unexpected remote write batch timestamp unit, expect millisecond, got: {}", - unit - ) - } - ); + ArrowDataType::Timestamp(_, _) => { source_map.insert(&target_ts_name, (src_idx, src_arrow_type)); } ArrowDataType::Utf8 => { @@ -156,15 +150,28 @@ pub(crate) fn rows_to_aligned_record_batch( } } - // Build columns in target schema order + // Build columns in target schema order. The timestamp column is built in + // the TARGET schema's unit; `build_arrow_array` converts the request's + // encoding unit (flooring on narrowing) into it. let mut columns = Vec::with_capacity(target_schema.fields().len()); for target_field in target_schema.fields() { if let Some((src_idx, src_arrow_type)) = source_map.get(target_field.name().as_str()) { + let target_type = if matches!( + (src_arrow_type, target_field.data_type()), + ( + ArrowDataType::Timestamp(_, _), + ArrowDataType::Timestamp(_, _) + ) + ) { + target_field.data_type().clone() + } else { + src_arrow_type.clone() + }; let array = build_arrow_array( rows, *src_idx, &rows.schema[*src_idx].column_name, - src_arrow_type.clone(), + target_type, row_count, )?; columns.push(array); @@ -199,6 +206,65 @@ pub(crate) fn identify_missing_columns_from_proto( Ok(missing) } +/// Builds a timestamp array of `T`'s unit from a proto column, appending +/// each value directly into the builder. Values already in the target unit +/// are appended without conversion (the unchanged-unit fast path); others +/// are converted via `Timestamp::convert_to`, flooring on narrowing. +fn build_timestamp_array>( + rows: &Rows, + col_idx: usize, + column_name: &str, + row_count: usize, + target_unit: CommonTimeUnit, +) -> Result { + let mut builder = PrimitiveBuilder::::with_capacity(row_count); + for row in &rows.rows { + let Some(value) = row.values[col_idx].value_data.as_ref() else { + builder.append_null(); + continue; + }; + let (source_unit, raw) = match value { + ValueData::TimestampSecondValue(v) => (CommonTimeUnit::Second, *v), + ValueData::TimestampMillisecondValue(v) => (CommonTimeUnit::Millisecond, *v), + ValueData::DatetimeValue(v) | ValueData::TimestampMicrosecondValue(v) => { + (CommonTimeUnit::Microsecond, *v) + } + ValueData::TimestampNanosecondValue(v) => (CommonTimeUnit::Nanosecond, *v), + v => { + return error::InvalidPromRemoteRequestSnafu { + msg: format!("Unexpected value: {:?}", v), + } + .fail(); + } + }; + if source_unit == target_unit { + builder.append_value(raw); + } else { + let timestamp = Timestamp::new(raw, source_unit); + let Some(converted) = timestamp.convert_to(target_unit) else { + return error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Timestamp value in column '{column_name}' overflows when converting to unit {target_unit:?}" + ), + } + .fail(); + }; + builder.append_value(converted.value()); + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) +} + +/// Converts an arrow time unit to the common time unit. +fn arrow_time_unit(unit: TimeUnit) -> CommonTimeUnit { + match unit { + TimeUnit::Second => CommonTimeUnit::Second, + TimeUnit::Millisecond => CommonTimeUnit::Millisecond, + TimeUnit::Microsecond => CommonTimeUnit::Microsecond, + TimeUnit::Nanosecond => CommonTimeUnit::Nanosecond, + } +} + /// Build a `Vec` suitable for creating a new Prometheus logical table /// directly from the proto `rows.schema`, avoiding the round-trip through Arrow schema. pub fn build_prom_create_table_schema_from_proto( @@ -207,13 +273,22 @@ pub fn build_prom_create_table_schema_from_proto( rows_schema .iter() .map(|col| { - let semantic_type = if col.datatype == api::v1::ColumnDataType::TimestampMillisecond as i32 { + let datatype = api::v1::ColumnDataType::try_from(col.datatype).map_err(|_| { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Failed to build create table schema, column '{}' has unknown datatype {}", + col.column_name, col.datatype + ), + } + .build() + })?; + let semantic_type = if api::helper::timestamp_unit(datatype).is_some() { SemanticType::Timestamp - } else if col.datatype == api::v1::ColumnDataType::Float64 as i32 { + } else if datatype == api::v1::ColumnDataType::Float64 { SemanticType::Field } else { // tag columns must be String type - ensure!(col.datatype == api::v1::ColumnDataType::String as i32, error::InvalidPromRemoteRequestSnafu{ + ensure!(datatype == api::v1::ColumnDataType::String, error::InvalidPromRemoteRequestSnafu{ msg: format!( "Failed to build create table schema, tag column '{}' must be String but got datatype {}", col.column_name, col.datatype @@ -268,25 +343,45 @@ fn build_arrow_array( StringBuilder::with_capacity(row_count, 0), ValueData::StringValue(v) => v ), - arrow::datatypes::DataType::Timestamp(u, _) => match u { - TimeUnit::Second => build_array!( - TimestampSecondBuilder::with_capacity(row_count), - ValueData::TimestampSecondValue(v) => *v - ), - TimeUnit::Millisecond => build_array!( - TimestampMillisecondBuilder::with_capacity(row_count), - ValueData::TimestampMillisecondValue(v) => *v - ), - TimeUnit::Microsecond => build_array!( - TimestampMicrosecondBuilder::with_capacity(row_count), - ValueData::DatetimeValue(v) => *v, - ValueData::TimestampMicrosecondValue(v) => *v - ), - TimeUnit::Nanosecond => build_array!( - TimestampNanosecondBuilder::with_capacity(row_count), - ValueData::TimestampNanosecondValue(v) => *v - ), - }, + arrow::datatypes::DataType::Timestamp(u, _) => { + // Accept any timestamp encoding and append directly into the + // target-unit builder — no intermediate column allocation. + // Values already in the target unit (the common unchanged + // millisecond case) are appended as-is; others are converted + // first, flooring on narrowing (same semantics as + // `Timestamp::convert_to` on the ordinary insert path). + let target_unit = arrow_time_unit(u); + match u { + TimeUnit::Second => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Millisecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Microsecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + TimeUnit::Nanosecond => build_timestamp_array::( + rows, + col_idx, + column_name, + row_count, + target_unit, + )?, + } + } ty => { return error::InvalidPromRemoteRequestSnafu { msg: format!( @@ -305,7 +400,9 @@ fn build_arrow_array( mod tests { use api::v1::value::ValueData; use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value}; - use arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; + use arrow::array::{ + Array, Float64Array, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, + }; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use super::{ @@ -313,6 +410,105 @@ mod tests { rows_to_aligned_record_batch, }; + #[test] + fn test_rows_to_aligned_record_batch_converts_time_units() { + // Microsecond-encoded rows aligned to a millisecond schema: narrowing + // floors towards negative infinity (-1001us -> -2ms), matching + // `Timestamp::convert_to` on the ordinary insert path. + let rows = Rows { + schema: vec![ + ColumnSchema { + column_name: "greptime_timestamp".to_string(), + datatype: ColumnDataType::TimestampMicrosecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }, + ColumnSchema { + column_name: "greptime_value".to_string(), + datatype: ColumnDataType::Float64 as i32, + semantic_type: SemanticType::Field as i32, + ..Default::default() + }, + ], + rows: vec![ + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMicrosecondValue(1000)), + }, + Value { + value_data: Some(ValueData::F64Value(1.0)), + }, + ], + }, + Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMicrosecondValue(-1001)), + }, + Value { + value_data: Some(ValueData::F64Value(2.0)), + }, + ], + }, + ], + }; + let target = ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("greptime_value", DataType::Float64, true), + ]); + + let (batch, _) = rows_to_aligned_record_batch(&rows, &target) + .unwrap() + .into_parts(); + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ts.value(0), 1); + assert_eq!(ts.value(1), -2); + + // Millisecond-encoded rows aligned to a microsecond schema: widening + // is lossless. + let mut rows = Rows { + schema: vec![rows.schema[0].clone(), rows.schema[1].clone()], + rows: vec![Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(123)), + }, + Value { + value_data: Some(ValueData::F64Value(4.0)), + }, + ], + }], + }; + rows.schema[0].datatype = ColumnDataType::TimestampMillisecond as i32; + let target = ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("greptime_value", DataType::Float64, true), + ]); + + let (batch, _) = rows_to_aligned_record_batch(&rows, &target) + .unwrap() + .into_parts(); + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ts.value(0), 123_000); + } + #[test] fn test_rows_to_aligned_record_batch_renames_and_reorders() { let rows = Rows { diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 59ce38919f7..29c8891a695 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -16,6 +16,7 @@ mod test { use std::sync::Arc; + use TimeUnit as ArrowTimeUnit; use client::{DEFAULT_CATALOG_NAME, OutputData}; use common_recordbatch::RecordBatches; use datatypes::arrow::array::AsArray; @@ -508,6 +509,268 @@ WITH( Ok(()) } + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_non_millisecond_physical_table() { + // Microsecond and nanosecond physical tables must both use the + // logical batcher: requests are converted to the physical table's + // unit during batch alignment + // (). + run_otlp_logical_batcher_non_millisecond_physical_table( + "us", + "TIMESTAMP(6)", + ArrowTimeUnit::Microsecond, + [60_000_000, 120_000_000], + ) + .await; + run_otlp_logical_batcher_non_millisecond_physical_table( + "ns", + "TIMESTAMP(9)", + ArrowTimeUnit::Nanosecond, + [60_000_000_000, 120_000_000_000], + ) + .await; + } + + async fn run_otlp_logical_batcher_non_millisecond_physical_table( + suite: &str, + sql_ts_type: &str, + unit: ArrowTimeUnit, + expected: [i64; 2], + ) { + use std::time::Duration; + + use common_base::Plugins; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::{TimestampMicrosecondType, TimestampNanosecondType}; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use prost::Message; + use servers::batcher::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + let standalone = + GreptimeDbStandaloneBuilder::new(&format!("otlp_logical_{suite}_physical")) + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let options = standalone.opts.clone(); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + + // Pre-create the physical metric table with a non-millisecond time + // index BEFORE any ingestion, so the batcher's bulk path must + // handle it. + let mut output = instance + .do_query( + &format!( + "CREATE TABLE greptime_physical_table (\ + greptime_timestamp {sql_ts_type} NOT NULL, \ + greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')" + ), + ctx.clone(), + ) + .await; + assert!(output.remove(0).is_ok()); + + // submit_build_and_align counts every batcher submission in both + // acknowledgement modes, so the assertion detects a regression that + // disables non-millisecond batching. + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + for (ts, value) in [(60, 10), (120, 20)] { + let request = build_sum_request( + "non.ms.alignment", + AggregationTemporality::Cumulative, + &[(ts, value)], + ); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + } + assert_eq!( + submissions.get_sample_count() - before, + 2, + "non-millisecond physical tables must use the logical batcher" + ); + + // The rows land on the non-millisecond physical table, converted + // from the nanosecond encoding. + let sql = "SELECT greptime_timestamp, greptime_value FROM non_ms_alignment_total ORDER BY greptime_timestamp"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if batches.iter().map(|batch| batch.num_rows()).sum::() == 2 { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + let batch = &batches.take()[0]; + let stored = |row: usize| match unit { + ArrowTimeUnit::Microsecond => batch + .column(0) + .as_primitive::() + .value(row), + ArrowTimeUnit::Nanosecond => batch + .column(0) + .as_primitive::() + .value(row), + _ => unreachable!("covered units only"), + }; + assert_eq!(stored(0), expected[0]); + assert_eq!(stored(1), expected[1]); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("expected f64 values")); + assert_eq!(values.value(0), 10.0); + assert_eq!(values.value(1), 20.0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_logical_batcher_fallback_for_cross_physical_destination() { + use std::time::Duration; + + use common_base::Plugins; + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::TimestampMicrosecondType; + use frontend::server::Services; + use frontend::service_config::pending_rows_batcher::BatcherOptions; + use prost::Message; + use servers::batcher::BatchingProtocol; + use servers::http::test_helpers::TestClient; + use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + + // With the logical batcher enabled, an OTLP request targeting an + // existing logical table bound to ANOTHER physical table must fall + // back to the ordinary insert path (the bulk eligibility check + // rejects the destination binding), which converts the request to + // the destination's unit. This covers the case previously guarded + // by the removed OTLP pre-gate alignment. + let standalone = GreptimeDbStandaloneBuilder::new("otlp_logical_cross_physical") + .with_logical_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_millis(5), + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let options = standalone.opts.clone(); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); + let server = services + .http_server_builder( + &options.frontend_options(), + services.server_memory_limiter.clone(), + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + ctx.set_logical_batching_enabled(true); + ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx { + with_metric_engine: true, + ..Default::default() + })); + let ctx = Arc::new(ctx); + + // A physical table with a microsecond time index and a logical table + // bound to it; OTLP requests always select the default + // (millisecond) physical table, so the destination binding differs. + for sql in [ + "CREATE TABLE phy_us (\ + greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + "CREATE TABLE cross_alignment_total (\ + greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + \"stream\" STRING NULL, TIME INDEX (greptime_timestamp), PRIMARY KEY (\"stream\")) \ + ENGINE = metric WITH ('on_physical_table' = 'phy_us')", + ] { + let mut output = instance.do_query(sql, ctx.clone()).await; + let result = output.remove(0); + assert!(result.is_ok(), "setup ddl failed: {result:?} — {sql}"); + } + + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + let request = build_sum_request( + "cross.alignment", + AggregationTemporality::Cumulative, + &[(60, 10)], + ); + let response = client + .post("/v1/otlp/v1/metrics") + .header("content-type", "application/x-protobuf") + .body(request.encode_to_vec()) + .send() + .await; + assert_eq!(response.status().as_u16(), 200); + // The request must NOT enter the logical batcher. + assert_eq!( + submissions.get_sample_count() - before, + 0, + "cross-physical destinations must fall back to the ordinary insert path" + ); + + // The row lands on the microsecond logical table, converted from the + // nanosecond encoding: 60s -> 60_000_000us. + let sql = "SELECT greptime_timestamp, greptime_value FROM cross_alignment_total"; + let batches = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let output = instance.do_query(sql, ctx.clone()).await.remove(0).unwrap(); + let OutputData::Stream(stream) = output.data else { + panic!("expected stream") + }; + let batches = RecordBatches::try_collect(stream).await.unwrap(); + if batches.iter().map(|batch| batch.num_rows()).sum::() == 1 { + break batches; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + let batch = &batches.take()[0]; + let timestamps = batch.column(0).as_primitive::(); + assert_eq!(timestamps.value(0), 60_000_000); + } + #[tokio::test(flavor = "multi_thread")] async fn test_otlp_logical_batcher_alignment() { use std::time::Duration; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1b68d3720a7..780140e3500 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -175,6 +175,8 @@ macro_rules! http_tests { test_prometheus_remote_write_batched_mixed_time_index_units, test_prometheus_remote_write_batched_interceptor_time_index_units, test_prometheus_remote_write_v2_batched_interceptor_time_index_units, + test_prometheus_remote_write_batched_microsecond_physical_table, + test_prometheus_remote_write_batched_conflicting_physical_selections, test_prometheus_remote_special_labels, test_prometheus_remote_schema_labels, test_prometheus_remote_write_with_pipeline, @@ -3832,8 +3834,9 @@ pub async fn test_prometheus_remote_write_batched_mixed_time_index_units(store_t }; // Create the logical table on the microsecond physical table. The bulk - // guard rejects the microsecond selected physical table, so this write - // takes the ordinary insert path. + // path handles the microsecond selected physical table: the new table + // is created on it and the samples are widened to its unit during + // batch alignment. let res = client .post("/v1/prometheus/write?physical_table=phy_us") .header("Content-Encoding", "snappy") @@ -3843,10 +3846,11 @@ pub async fn test_prometheus_remote_write_batched_mixed_time_index_units(store_t assert_eq!(res.status(), StatusCode::NO_CONTENT); // Write the same existing table again while selecting the default - // (millisecond) physical table: the guard must reject the existing - // microsecond destination and fall back to the ordinary insert path — - // without the destination check the bulk path would build millisecond - // arrays against the microsecond schema and fail the write. + // (millisecond) physical table: the destination is bound to another + // physical table, so the bulk eligibility check must reject it and fall + // back to the ordinary insert path — without the destination binding + // check the bulk flush would write the rows through the selected + // physical's regions, silently misplacing them. let res = client .post("/v1/prometheus/write") .header("Content-Encoding", "snappy") @@ -4069,6 +4073,224 @@ pub async fn test_prometheus_remote_write_v2_batched_interceptor_time_index_unit /// bypasses `PromStoreProtocolHandler::write`. Verifies the metric table is created /// asynchronously and still carries the Prometheus semantic identity stamped on the /// shared request context. +/// Regression test for : +/// prometheus remote write uses the logical batcher against a physical +/// metric table pre-created with a microsecond time index; the millisecond +/// samples are widened to the physical table's unit during batch alignment. +pub async fn test_prometheus_remote_write_batched_microsecond_physical_table( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend_batched( + store_type, + "prometheus_remote_write_batched_us_physical", + ) + .await; + let client = TestClient::new(app).await; + + // Pre-create the default physical metric table with a microsecond time + // index before any remote write, so the batched bulk path must handle it. + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE greptime_physical_table \ + (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \ + TIME INDEX (greptime_timestamp)) \ + ENGINE = metric WITH ('physical_metric_table' = 'true')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write_request = WriteRequest { + timeseries: vec![prom_store::mock_timeseries()[0].clone()], + ..Default::default() + }; + let serialized_request = write_request.encode_to_vec(); + let compressed_request = + prom_store::snappy_compress(&serialized_request).expect("failed to encode snappy"); + + // submit_build_and_align counts every batcher submission in both + // acknowledgement modes, so the stored-value assertions below cannot be + // satisfied by a silent fallback to ordinary insertion. + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + let before = submissions.get_sample_count(); + + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 1, + "non-millisecond physical tables must use the logical batcher" + ); + + // metric1 samples are 1.0@1000ms and 2.0@2000ms; on the microsecond + // physical table they must be stored as 1_000_000us and 2_000_000us. + wait_for_data( + &client, + "select greptime_timestamp, greptime_value from metric1 order by greptime_timestamp", + "[[1000000,1.0],[2000000,2.0]]", + ) + .await; + + guard.remove_all().await; +} + +/// Regression test for the per-series physical-table selection conflict: +/// one remote-write request containing the same metric under two different +/// physical-table selections must fall back to the ordinary insert path — +/// both when the destination table already exists (bound to one of the two +/// physicals) and when it is missing (conflicting creations). A consistent +/// selection must still use the batcher. +pub async fn test_prometheus_remote_write_batched_conflicting_physical_selections( + store_type: StorageType, +) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend_batched( + store_type, + "prom_rw_batched_conflicting_physicals", + ) + .await; + let client = TestClient::new(app).await; + + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE p1 (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, TIME INDEX (greptime_timestamp)) ENGINE = metric WITH ('physical_metric_table' = 'true')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + // An existing logical table bound to p1. + let res = client + .get( + "/v1/sql?db=public&sql=CREATE TABLE conflict_existing (greptime_timestamp TIMESTAMP(6) NOT NULL, greptime_value DOUBLE NULL, \"job\" STRING NULL, TIME INDEX (greptime_timestamp), PRIMARY KEY (\"job\")) ENGINE = metric WITH ('on_physical_table' = 'p1')", + ) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let write = |metric: &str, job: &str, physical: Option<&str>, value: f64| { + let mut labels = vec![ + Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: metric.to_string(), + }, + Label { + name: "job".to_string(), + value: job.to_string(), + }, + ]; + if let Some(physical) = physical { + labels.push(Label { + name: "x_greptime_physical_table".to_string(), + value: physical.to_string(), + }); + } + TimeSeries { + labels, + samples: vec![Sample { + value, + timestamp: 1000, + }], + ..Default::default() + } + }; + let send = |timeseries: Vec| { + let write_request = WriteRequest { + timeseries, + ..Default::default() + }; + prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap() + }; + + let submissions = servers::metrics::PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED + .with_label_values(&["submit_build_and_align"]); + + // Existing destination under two different selections: the batch must + // fall back (the batch for the unbound selection would otherwise write + // through the default physical's regions), and both series land on the + // bound physical. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("conflict_existing", "a", Some("p1"), 1.0), + write("conflict_existing", "b", None, 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 0, + "conflicting selections for an existing table must fall back" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from conflict_existing", + "[[2,3.0]]", + ) + .await; + + // Missing destination under two different selections: conflicting + // creations must also fall back, and both series must be visible + // through the created table. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("conflict_fresh", "a", Some("p1"), 1.0), + write("conflict_fresh", "b", None, 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 0, + "conflicting selections for a missing table must fall back" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from conflict_fresh", + "[[2,3.0]]", + ) + .await; + + // A consistent selection still uses the batcher. + let before = submissions.get_sample_count(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(send(vec![ + write("consistent_fresh", "a", Some("p1"), 1.0), + write("consistent_fresh", "b", Some("p1"), 2.0), + ])) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + submissions.get_sample_count() - before, + 1, + "a consistent selection must use the logical batcher" + ); + wait_for_data( + &client, + "select count(*), sum(greptime_value) from consistent_fresh", + "[[2,3.0]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_prometheus_remote_write_batched(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = From 408536af18a329f5b935ee928f95393eac022200 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 12:34:32 +0000 Subject: [PATCH 061/133] chore: remove leftover flow worker config docs and unused flow metrics (#9375) Signed-off-by: Dennis Zhuang --- config/config.md | 3 - config/flownode.example.toml | 3 - config/standalone.example.toml | 6 - .../dashboards/metrics/cluster/dashboard.json | 340 +----------------- .../dashboards/metrics/cluster/dashboard.md | 5 +- .../dashboards/metrics/cluster/dashboard.yaml | 42 --- .../metrics/standalone/dashboard.json | 340 +----------------- .../metrics/standalone/dashboard.md | 5 +- .../metrics/standalone/dashboard.yaml | 42 --- src/flow/src/metrics.rs | 16 - src/pipeline/src/etl/transform/index.rs | 2 +- 11 files changed, 19 insertions(+), 785 deletions(-) diff --git a/config/config.md b/config/config.md index 89b3e13dcdd..116cf8a4f82 100644 --- a/config/config.md +++ b/config/config.md @@ -138,8 +138,6 @@ | `procedure.max_retry_times` | Integer | `3` | Procedure max retry time. | | `procedure.retry_delay` | String | `500ms` | Initial retry delay of procedures, increases exponentially | | `procedure.max_running_procedures` | Integer | `128` | Max running procedures.
The maximum number of procedures that can be running at the same time.
If the number of running procedures exceeds this limit, the procedure will be rejected. | -| `flow` | -- | -- | flow engine options. | -| `flow.num_workers` | Integer | `0` | The number of flow worker in flownode.
Not setting(or set to 0) this value will use the number of CPU cores divided by 2. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | | `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | @@ -742,7 +740,6 @@ | --- | -----| ------- | ----------- | | `node_id` | Integer | Unset | The flownode identifier and should be unique in the cluster. | | `flow` | -- | -- | flow engine options. | -| `flow.num_workers` | Integer | `0` | The number of flow worker in flownode.
Not setting(or set to 0) this value will use the number of CPU cores divided by 2. | | `flow.batching_mode` | -- | -- | -- | | `flow.batching_mode.query_timeout` | String | `600s` | The default batching engine query timeout is 10 minutes. | | `flow.batching_mode.slow_query_threshold` | String | `60s` | will output a warn log for any query that runs for more that this threshold | diff --git a/config/flownode.example.toml b/config/flownode.example.toml index 192c6af6977..b0a53bee54a 100644 --- a/config/flownode.example.toml +++ b/config/flownode.example.toml @@ -4,9 +4,6 @@ node_id = 14 ## flow engine options. [flow] -## The number of flow worker in flownode. -## Not setting(or set to 0) this value will use the number of CPU cores divided by 2. -#+num_workers=0 [flow.batching_mode] ## The default batching engine query timeout is 10 minutes. #+query_timeout="600s" diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e5eaa86c4c5..e3e28563582 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -477,12 +477,6 @@ retry_delay = "500ms" ## If the number of running procedures exceeds this limit, the procedure will be rejected. max_running_procedures = 128 -## flow engine options. -[flow] -## The number of flow worker in flownode. -## Not setting(or set to 0) this value will use the number of CPU cores divided by 2. -#+num_workers=0 - # Example of using S3 as the storage. # [storage] # type = "S3" diff --git a/grafana/dashboards/metrics/cluster/dashboard.json b/grafana/dashboards/metrics/cluster/dashboard.json index ddbcc8e190c..daa0f018836 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.json +++ b/grafana/dashboards/metrics/cluster/dashboard.json @@ -3585,19 +3585,6 @@ }, "pluginVersion": "11.6.0", "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by (code) (rate(greptime_flow_errors[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "flow-{{code}}", - "range": true, - "refId": "A" - }, { "datasource": { "type": "prometheus", @@ -16622,127 +16609,6 @@ "title": "Flow Ingest / Output Rate", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Ingest Latency.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "points", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 170 - }, - "id": 521, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p95", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", - "range": true, - "refId": "B" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", - "range": true, - "refId": "C" - } - ], - "title": "Flow Ingest Latency", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -16805,7 +16671,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, + "x": 8, "y": 170 }, "id": 522, @@ -16851,196 +16717,6 @@ "title": "Flow Operation Latency", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Buffer Size per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 0, - "y": 178 - }, - "id": 523, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "greptime_flow_input_buf_size", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Buffer Size per Instance", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Processing Error per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 178 - }, - "id": 524, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-[{{code}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Processing Error per Instance", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -17105,7 +16781,7 @@ "h": 8, "w": 8, "x": 16, - "y": 178 + "y": 170 }, "id": 558, "options": { @@ -17229,7 +16905,7 @@ "h": 8, "w": 8, "x": 0, - "y": 186 + "y": 178 }, "id": 559, "options": { @@ -17353,7 +17029,7 @@ "h": 8, "w": 8, "x": 8, - "y": 186 + "y": 178 }, "id": 560, "options": { @@ -17464,7 +17140,7 @@ "h": 8, "w": 8, "x": 16, - "y": 186 + "y": 178 }, "id": 561, "options": { @@ -17575,7 +17251,7 @@ "h": 8, "w": 8, "x": 0, - "y": 194 + "y": 186 }, "id": 562, "options": { @@ -17673,7 +17349,7 @@ "h": 8, "w": 8, "x": 8, - "y": 194 + "y": 186 }, "id": 563, "options": { @@ -17771,7 +17447,7 @@ "h": 8, "w": 8, "x": 16, - "y": 194 + "y": 186 }, "id": 564, "options": { diff --git a/grafana/dashboards/metrics/cluster/dashboard.md b/grafana/dashboards/metrics/cluster/dashboard.md index cc34ddad12a..3278db52bfd 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.md +++ b/grafana/dashboards/metrics/cluster/dashboard.md @@ -36,7 +36,7 @@ | Scan and Compaction Memory Rejects | `sum(rate(greptime_mito_scan_requests_rejected_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_scan_memory_exhausted_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_compaction_memory_rejected_total{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Datanode scan and compaction memory rejection/exhaustion counters. | `prometheus` | `rps` | `scan-rejected` | | OpenDAL Errors | `sum by (scheme, operation, error) (rate(opendal_operation_errors_total{instance=~"$datanode",error!="NotFound"}[$__rate_interval]))` | `timeseries` | Object-store errors by scheme, operation, and error, excluding NotFound noise. | `prometheus` | `eps` | `{{scheme}}-{{operation}}-{{error}}` | | Metasrv Failures | `sum(rate(greptime_meta_region_migration_fail[$__rate_interval]))`
`sum(rate(greptime_meta_reconciliation_procedure_error[$__rate_interval]))` | `timeseries` | Region migration and reconciliation failures in metasrv. | `prometheus` | `eps` | `migration-fail` | -| Flow and Trigger Failures | `sum by (code) (rate(greptime_flow_errors[$__rate_interval]))`
`sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `flow-{{code}}` | +| Flow and Trigger Failures | `sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `trigger-evaluate` | | Mito GC Failures | `sum(rate(greptime_mito_gc_errors_total{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_gc_orphaned_index_files{instance=~"$datanode"}[$__rate_interval]))`
`sum(rate(greptime_mito_gc_skipped_unparsable_files{instance=~"$datanode"}[$__rate_interval]))` | `timeseries` | Mito garbage-collection errors and skipped/orphaned files on datanodes. | `prometheus` | `short` | `gc-errors` | # Capacity | Title | Query | Type | Description | Datasource | Unit | Legend Format | @@ -278,10 +278,7 @@ ORDER BY data_size DESC;` | `piechart` | Distribution of leader regions and data | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | | Flow Ingest / Output Rate | `sum by(instance, pod, direction) (rate(greptime_flow_processed_rows[$__rate_interval]))` | `timeseries` | Flow Ingest / Output Rate. | `prometheus` | -- | `[{{pod}}]-[{{instance}}]-[{{direction}}]` | -| Flow Ingest Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))` | `timeseries` | Flow Ingest Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-p95` | | Flow Operation Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))`
`histogram_quantile(0.99, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))` | `timeseries` | Flow Operation Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | -| Flow Buffer Size per Instance | `greptime_flow_input_buf_size` | `timeseries` | Flow Buffer Size per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]` | -| Flow Processing Error per Instance | `sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))` | `timeseries` | Flow Processing Error per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{code}}]` | | Flow Batching Attempt / Error Rate | `sum by(instance,pod,flow_id) (rate(greptime_flow_batching_start_query_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_error_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_slow_query_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode task attempt, error, and slow-query rates by flow. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-attempt` | | Flow Batching Query Latency | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode query latency by flow. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | | Flow Batching Query Window Count | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_count[$__rate_interval]))` | `timeseries` | Flow batching mode query window count by flow. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | diff --git a/grafana/dashboards/metrics/cluster/dashboard.yaml b/grafana/dashboards/metrics/cluster/dashboard.yaml index 9664db35f69..bc659b401b5 100644 --- a/grafana/dashboards/metrics/cluster/dashboard.yaml +++ b/grafana/dashboards/metrics/cluster/dashboard.yaml @@ -515,11 +515,6 @@ groups: description: Derived-data and alerting pipeline failures. unit: eps queries: - - expr: sum by (code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: flow-{{code}} - expr: sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval])) datasource: type: prometheus @@ -2199,25 +2194,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{pod}}]-[{{instance}}]-[{{direction}}]' - - title: Flow Ingest Latency - type: timeseries - description: Flow Ingest Latency. - queries: - - expr: histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p95' - - expr: histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' - title: Flow Operation Latency type: timeseries description: Flow Operation Latency. @@ -2232,24 +2208,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' - - title: Flow Buffer Size per Instance - type: timeseries - description: Flow Buffer Size per Instance. - queries: - - expr: greptime_flow_input_buf_size - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]' - - title: Flow Processing Error per Instance - type: timeseries - description: Flow Processing Error per Instance. - queries: - - expr: sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-[{{code}}]' - title: Flow Batching Attempt / Error Rate type: timeseries description: Flow batching mode task attempt, error, and slow-query rates by flow. diff --git a/grafana/dashboards/metrics/standalone/dashboard.json b/grafana/dashboards/metrics/standalone/dashboard.json index 9d181b05ffe..c1fda955d90 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.json +++ b/grafana/dashboards/metrics/standalone/dashboard.json @@ -3585,19 +3585,6 @@ }, "pluginVersion": "11.6.0", "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by (code) (rate(greptime_flow_errors[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "flow-{{code}}", - "range": true, - "refId": "A" - }, { "datasource": { "type": "prometheus", @@ -16622,127 +16609,6 @@ "title": "Flow Ingest / Output Rate", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Ingest Latency.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "points", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 170 - }, - "id": 521, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p95", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-p99", - "range": true, - "refId": "B" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))", - "hide": false, - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-avg", - "range": true, - "refId": "C" - } - ], - "title": "Flow Ingest Latency", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -16805,7 +16671,7 @@ "gridPos": { "h": 8, "w": 8, - "x": 16, + "x": 8, "y": 170 }, "id": 522, @@ -16851,196 +16717,6 @@ "title": "Flow Operation Latency", "type": "timeseries" }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Buffer Size per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 0, - "y": 178 - }, - "id": 523, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "greptime_flow_input_buf_size", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Buffer Size per Instance", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "description": "Flow Processing Error per Instance.", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisBorderShow": false, - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "insertNulls": false, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green" - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 8, - "x": 8, - "y": 178 - }, - "id": 524, - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "multi", - "sort": "desc" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${metrics}" - }, - "editorMode": "code", - "expr": "sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))", - "instant": false, - "legendFormat": "[{{instance}}]-[{{pod}}]-[{{code}}]", - "range": true, - "refId": "A" - } - ], - "title": "Flow Processing Error per Instance", - "type": "timeseries" - }, { "datasource": { "type": "prometheus", @@ -17105,7 +16781,7 @@ "h": 8, "w": 8, "x": 16, - "y": 178 + "y": 170 }, "id": 558, "options": { @@ -17229,7 +16905,7 @@ "h": 8, "w": 8, "x": 0, - "y": 186 + "y": 178 }, "id": 559, "options": { @@ -17353,7 +17029,7 @@ "h": 8, "w": 8, "x": 8, - "y": 186 + "y": 178 }, "id": 560, "options": { @@ -17464,7 +17140,7 @@ "h": 8, "w": 8, "x": 16, - "y": 186 + "y": 178 }, "id": 561, "options": { @@ -17575,7 +17251,7 @@ "h": 8, "w": 8, "x": 0, - "y": 194 + "y": 186 }, "id": 562, "options": { @@ -17673,7 +17349,7 @@ "h": 8, "w": 8, "x": 8, - "y": 194 + "y": 186 }, "id": 563, "options": { @@ -17771,7 +17447,7 @@ "h": 8, "w": 8, "x": 16, - "y": 194 + "y": 186 }, "id": 564, "options": { diff --git a/grafana/dashboards/metrics/standalone/dashboard.md b/grafana/dashboards/metrics/standalone/dashboard.md index dd00fb93d49..5dd6f1600e1 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.md +++ b/grafana/dashboards/metrics/standalone/dashboard.md @@ -36,7 +36,7 @@ | Scan and Compaction Memory Rejects | `sum(rate(greptime_mito_scan_requests_rejected_total[$__rate_interval]))`
`sum(rate(greptime_mito_scan_memory_exhausted_total[$__rate_interval]))`
`sum(rate(greptime_mito_compaction_memory_rejected_total[$__rate_interval]))` | `timeseries` | Datanode scan and compaction memory rejection/exhaustion counters. | `prometheus` | `rps` | `scan-rejected` | | OpenDAL Errors | `sum by (scheme, operation, error) (rate(opendal_operation_errors_total{error!="NotFound"}[$__rate_interval]))` | `timeseries` | Object-store errors by scheme, operation, and error, excluding NotFound noise. | `prometheus` | `eps` | `{{scheme}}-{{operation}}-{{error}}` | | Metasrv Failures | `sum(rate(greptime_meta_region_migration_fail[$__rate_interval]))`
`sum(rate(greptime_meta_reconciliation_procedure_error[$__rate_interval]))` | `timeseries` | Region migration and reconciliation failures in metasrv. | `prometheus` | `eps` | `migration-fail` | -| Flow and Trigger Failures | `sum by (code) (rate(greptime_flow_errors[$__rate_interval]))`
`sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `flow-{{code}}` | +| Flow and Trigger Failures | `sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_send_alert_failure_count[$__rate_interval]))`
`sum(rate(greptime_trigger_save_alert_record_failure_count[$__rate_interval]))`
`sum(rate(greptime_flow_batching_error_count[$__rate_interval]))` | `timeseries` | Derived-data and alerting pipeline failures. | `prometheus` | `eps` | `trigger-evaluate` | | Mito GC Failures | `sum(rate(greptime_mito_gc_errors_total[$__rate_interval]))`
`sum(rate(greptime_mito_gc_orphaned_index_files[$__rate_interval]))`
`sum(rate(greptime_mito_gc_skipped_unparsable_files[$__rate_interval]))` | `timeseries` | Mito garbage-collection errors and skipped/orphaned files on datanodes. | `prometheus` | `short` | `gc-errors` | # Capacity | Title | Query | Type | Description | Datasource | Unit | Legend Format | @@ -278,10 +278,7 @@ ORDER BY data_size DESC;` | `piechart` | Distribution of leader regions and data | Title | Query | Type | Description | Datasource | Unit | Legend Format | | --- | --- | --- | --- | --- | --- | --- | | Flow Ingest / Output Rate | `sum by(instance, pod, direction) (rate(greptime_flow_processed_rows[$__rate_interval]))` | `timeseries` | Flow Ingest / Output Rate. | `prometheus` | -- | `[{{pod}}]-[{{instance}}]-[{{direction}}]` | -| Flow Ingest Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod))`
`sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval]))` | `timeseries` | Flow Ingest Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-p95` | | Flow Operation Latency | `histogram_quantile(0.95, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))`
`histogram_quantile(0.99, sum(rate(greptime_flow_processing_time_bucket[$__rate_interval])) by (le,instance,pod,type))` | `timeseries` | Flow Operation Latency. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{type}}]-p95` | -| Flow Buffer Size per Instance | `greptime_flow_input_buf_size` | `timeseries` | Flow Buffer Size per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]` | -| Flow Processing Error per Instance | `sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval]))` | `timeseries` | Flow Processing Error per Instance. | `prometheus` | -- | `[{{instance}}]-[{{pod}}]-[{{code}}]` | | Flow Batching Attempt / Error Rate | `sum by(instance,pod,flow_id) (rate(greptime_flow_batching_start_query_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_error_count[$__rate_interval]))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_slow_query_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode task attempt, error, and slow-query rates by flow. | `prometheus` | `ops` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-attempt` | | Flow Batching Query Latency | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`histogram_quantile(0.99, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_time_secs_count[$__rate_interval]))` | `timeseries` | Flow batching mode query latency by flow. | `prometheus` | `s` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | | Flow Batching Query Window Count | `histogram_quantile(0.95, sum by(le,instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_bucket[$__rate_interval])))`
`sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_sum[$__rate_interval])) / sum by(instance,pod,flow_id) (rate(greptime_flow_batching_engine_query_window_cnt_count[$__rate_interval]))` | `timeseries` | Flow batching mode query window count by flow. | `prometheus` | `short` | `[{{instance}}]-[{{pod}}]-[flow={{flow_id}}]-p95` | diff --git a/grafana/dashboards/metrics/standalone/dashboard.yaml b/grafana/dashboards/metrics/standalone/dashboard.yaml index 86076d3d33a..beadb9c3215 100644 --- a/grafana/dashboards/metrics/standalone/dashboard.yaml +++ b/grafana/dashboards/metrics/standalone/dashboard.yaml @@ -515,11 +515,6 @@ groups: description: Derived-data and alerting pipeline failures. unit: eps queries: - - expr: sum by (code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: flow-{{code}} - expr: sum(rate(greptime_trigger_evaluate_failure_count[$__rate_interval])) datasource: type: prometheus @@ -2199,25 +2194,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{pod}}]-[{{instance}}]-[{{direction}}]' - - title: Flow Ingest Latency - type: timeseries - description: Flow Ingest Latency. - queries: - - expr: histogram_quantile(0.95, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p95' - - expr: histogram_quantile(0.99, sum(rate(greptime_flow_insert_elapsed_bucket[$__rate_interval])) by (le, instance, pod)) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-p99' - - expr: sum by(instance, pod) (rate(greptime_flow_insert_elapsed_sum[$__rate_interval])) / sum by(instance, pod) (rate(greptime_flow_insert_elapsed_count[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-avg' - title: Flow Operation Latency type: timeseries description: Flow Operation Latency. @@ -2232,24 +2208,6 @@ groups: type: prometheus uid: ${metrics} legendFormat: '[{{instance}}]-[{{pod}}]-[{{type}}]-p99' - - title: Flow Buffer Size per Instance - type: timeseries - description: Flow Buffer Size per Instance. - queries: - - expr: greptime_flow_input_buf_size - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]' - - title: Flow Processing Error per Instance - type: timeseries - description: Flow Processing Error per Instance. - queries: - - expr: sum by(instance,pod,code) (rate(greptime_flow_errors[$__rate_interval])) - datasource: - type: prometheus - uid: ${metrics} - legendFormat: '[{{instance}}]-[{{pod}}]-[{{code}}]' - title: Flow Batching Attempt / Error Rate type: timeseries description: Flow batching mode task attempt, error, and slow-query rates by flow. diff --git a/src/flow/src/metrics.rs b/src/flow/src/metrics.rs index 00f93d47ab8..db524c669a3 100644 --- a/src/flow/src/metrics.rs +++ b/src/flow/src/metrics.rs @@ -20,14 +20,6 @@ use prometheus::*; lazy_static! { pub static ref METRIC_FLOW_TASK_COUNT: IntGauge = register_int_gauge!("greptime_flow_task_count", "flow task count").unwrap(); - pub static ref METRIC_FLOW_INPUT_BUF_SIZE: IntGauge = - register_int_gauge!("greptime_flow_input_buf_size", "flow input buf size").unwrap(); - pub static ref METRIC_FLOW_INSERT_ELAPSED: HistogramVec = register_histogram_vec!( - "greptime_flow_insert_elapsed", - "flow insert elapsed", - &["table_id"] - ) - .unwrap(); pub static ref METRIC_FLOW_BATCHING_ENGINE_QUERY_TIME: HistogramVec = register_histogram_vec!( "greptime_flow_batching_engine_query_time_secs", "flow batching engine query time(seconds)", @@ -101,8 +93,6 @@ lazy_static! { &["flow_id", "mode"], ) .unwrap(); - pub static ref METRIC_FLOW_RUN_INTERVAL_MS: IntGauge = - register_int_gauge!("greptime_flow_run_interval_ms", "flow run interval in ms").unwrap(); pub static ref METRIC_FLOW_ROWS: IntCounterVec = register_int_counter_vec!( "greptime_flow_processed_rows", "Count of rows flowing through the system.", @@ -115,10 +105,4 @@ lazy_static! { &["type"] ) .unwrap(); - pub static ref METRIC_FLOW_ERRORS: IntCounterVec = register_int_counter_vec!( - "greptime_flow_errors", - "Count of errors in flow processing", - &["code"] - ) - .unwrap(); } diff --git a/src/pipeline/src/etl/transform/index.rs b/src/pipeline/src/etl/transform/index.rs index 53e16dca222..b0cdd9c4fe0 100644 --- a/src/pipeline/src/etl/transform/index.rs +++ b/src/pipeline/src/etl/transform/index.rs @@ -25,7 +25,7 @@ const INDEX_INVERTED: &str = "inverted"; #[allow(clippy::enum_variant_names)] pub enum Index { Time, - // deprecated, use Inverted instead + // Deprecated, equivalent to `tag: true`. Tag, Fulltext, Skipping, From 180221240a756d6ea53179d773040fb02a3d484e Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 12:34:54 +0000 Subject: [PATCH 062/133] test: remove unused legacy compatibility test suites (#9378) Signed-off-by: Dennis Zhuang --- tests/compat/README.md | 31 ---- tests/compat/case/read/standalone/read.result | 13 -- tests/compat/case/read/standalone/read.sql | 1 - .../compat/case/write/standalone/write.result | 25 --- tests/compat/case/write/standalone/write.sql | 16 -- tests/compat/test-compat.sh | 37 ----- tests/compat/util.sh | 82 ---------- tests/upgrade-compat/distributed/common | 1 - .../common/table_engine_0_10_2.result | 137 ---------------- .../standalone/common/table_engine_0_10_2.sql | 60 ------- .../common/table_engine_v0_11_0.result | 137 ---------------- .../common/table_engine_v0_11_0.sql | 60 ------- .../common/table_engine_v0_9_5.result | 137 ---------------- .../standalone/common/table_engine_v0_9_5.sql | 60 ------- .../common/test_alter_drop_column.result | 62 ------- .../common/test_alter_drop_column.sql | 23 --- .../standalone/common/test_alter_table.result | 54 ------- .../standalone/common/test_alter_table.sql | 20 --- .../standalone/common/test_simple.result | 47 ------ .../standalone/common/test_simple.sql | 22 --- .../standalone/common/test_ttl.result | 153 ------------------ .../standalone/common/test_ttl.sql | 42 ----- .../standalone/common/test_update.result | 32 ---- .../standalone/common/test_update.sql | 16 -- 24 files changed, 1268 deletions(-) delete mode 100644 tests/compat/README.md delete mode 100644 tests/compat/case/read/standalone/read.result delete mode 100644 tests/compat/case/read/standalone/read.sql delete mode 100644 tests/compat/case/write/standalone/write.result delete mode 100644 tests/compat/case/write/standalone/write.sql delete mode 100755 tests/compat/test-compat.sh delete mode 100755 tests/compat/util.sh delete mode 120000 tests/upgrade-compat/distributed/common delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_0_10_2.result delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result delete mode 100644 tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql delete mode 100644 tests/upgrade-compat/standalone/common/test_alter_drop_column.result delete mode 100644 tests/upgrade-compat/standalone/common/test_alter_drop_column.sql delete mode 100644 tests/upgrade-compat/standalone/common/test_alter_table.result delete mode 100644 tests/upgrade-compat/standalone/common/test_alter_table.sql delete mode 100644 tests/upgrade-compat/standalone/common/test_simple.result delete mode 100644 tests/upgrade-compat/standalone/common/test_simple.sql delete mode 100644 tests/upgrade-compat/standalone/common/test_ttl.result delete mode 100644 tests/upgrade-compat/standalone/common/test_ttl.sql delete mode 100644 tests/upgrade-compat/standalone/common/test_update.result delete mode 100644 tests/upgrade-compat/standalone/common/test_update.sql diff --git a/tests/compat/README.md b/tests/compat/README.md deleted file mode 100644 index 07c4cb494f1..00000000000 --- a/tests/compat/README.md +++ /dev/null @@ -1,31 +0,0 @@ -# GreptimeDB compatibility test - -The compatibility test check whether a newer version of GreptimeDB can read from the data written by an old version of -GreptimeDB (the "backward" compatibility), and vice-versa (the "forward" compatibility). It's often ran in the Github -Actions to ensure there are no breaking changes by accident. - -The test work like this: reuse the sqlness-runner two times but each for a read or write side. For example, if we are -testing backward compatibility, we use sqlness-runner to run the SQLs with writes against the old GreptimeDB binary, and -use the same sqlness-runner to run the SQLs with reads against the new GreptimeDB binary. If the reads were executed -expectedly, we have achieved backward compatibility. - -This compatibility test is inspired by [Databend](https://github.com/datafuselabs/databend/). - -## Usage - -```shell -tests/compat/test-compat.sh -``` - -E.g. `tests/compat/test-compat.sh 0.6.0` tests if the data written by GreptimeDB **v0.6.0** can be read by **current** -version of GreptimeDB, and vice-versa. By "current", it's meant the fresh binary built by current codes. - -## Prerequisites - -Current version of GreptimeDB's binaries must reside in `./bins`: - -- `./bins/current/greptime` -- `./bins/current/sqlness-runner` - -The steps in Github Action already assure that. When running in local host, you have to `cp` them from target directory -manually. diff --git a/tests/compat/case/read/standalone/read.result b/tests/compat/case/read/standalone/read.result deleted file mode 100644 index ba59e7a4628..00000000000 --- a/tests/compat/case/read/standalone/read.result +++ /dev/null @@ -1,13 +0,0 @@ -select ts, i, s, f from foo order by ts; - -+---------------------+---+----------+-----+ -| ts | i | s | f | -+---------------------+---+----------+-----+ -| 2024-02-01T17:00:00 | 1 | my_tag_1 | | -| 2024-02-01T18:00:00 | 2 | my_tag_2 | | -| 2024-02-01T19:00:00 | 3 | my_tag_3 | | -| 2024-02-01T20:00:00 | 4 | my_tag_4 | 4.4 | -| 2024-02-01T21:00:00 | 5 | my_tag_5 | 5.5 | -| 2024-02-01T22:00:00 | 6 | my_tag_6 | 6.6 | -+---------------------+---+----------+-----+ - diff --git a/tests/compat/case/read/standalone/read.sql b/tests/compat/case/read/standalone/read.sql deleted file mode 100644 index e83d62d1ce9..00000000000 --- a/tests/compat/case/read/standalone/read.sql +++ /dev/null @@ -1 +0,0 @@ -select ts, i, s, f from foo order by ts; diff --git a/tests/compat/case/write/standalone/write.result b/tests/compat/case/write/standalone/write.result deleted file mode 100644 index 85ea219981c..00000000000 --- a/tests/compat/case/write/standalone/write.result +++ /dev/null @@ -1,25 +0,0 @@ -create table foo(ts timestamp time index, s string primary key, i int); - -Affected Rows: 0 - -insert into foo values -("2024-02-02 01:00:00+0800", "my_tag_1", 1), -("2024-02-02 02:00:00+0800", "my_tag_2", 2), -("2024-02-02 03:00:00+0800", "my_tag_3", 3); - -Affected Rows: 3 - --- Alter the table to trigger a flush (will be executed before process being terminated). --- Otherwise the SST might not be generated (the data could be remained in WAL). --- If we have the explicitly flush table interface in the future, it's still good to have the alter table in the test. -alter table foo add column f float; - -Affected Rows: 0 - -insert into foo values -("2024-02-02 04:00:00+0800", "my_tag_4", 4, 4.4), -("2024-02-02 05:00:00+0800", "my_tag_5", 5, 5.5), -("2024-02-02 06:00:00+0800", "my_tag_6", 6, 6.6); - -Affected Rows: 3 - diff --git a/tests/compat/case/write/standalone/write.sql b/tests/compat/case/write/standalone/write.sql deleted file mode 100644 index d7403a8e267..00000000000 --- a/tests/compat/case/write/standalone/write.sql +++ /dev/null @@ -1,16 +0,0 @@ -create table foo(ts timestamp time index, s string primary key, i int); - -insert into foo values -("2024-02-02 01:00:00+0800", "my_tag_1", 1), -("2024-02-02 02:00:00+0800", "my_tag_2", 2), -("2024-02-02 03:00:00+0800", "my_tag_3", 3); - --- Alter the table to trigger a flush (will be executed before process being terminated). --- Otherwise the SST might not be generated (the data could be remained in WAL). --- If we have the explicitly flush table interface in the future, it's still good to have the alter table in the test. -alter table foo add column f float; - -insert into foo values -("2024-02-02 04:00:00+0800", "my_tag_4", 4, 4.4), -("2024-02-02 05:00:00+0800", "my_tag_5", 5, 5.5), -("2024-02-02 06:00:00+0800", "my_tag_6", 6, 6.6); diff --git a/tests/compat/test-compat.sh b/tests/compat/test-compat.sh deleted file mode 100755 index 837a523b2ec..00000000000 --- a/tests/compat/test-compat.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash - -set -o errexit - -usage() { - echo " Tests the compatibility between different versions of GreptimeDB." - echo " Expects the directory './bins/current' contains the newly built binaries." - echo " Usage: $0 " -} - -# The previous version of GreptimeDB to test compatibility with. -# e.g. old_ver="0.6.0" -old_ver="$1" - -if [ -z $old_ver ] -then - usage - exit -1 -fi - -SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" -echo " === SCRIPT_PATH: $SCRIPT_PATH" -source "${SCRIPT_PATH}/util.sh" - -# go to work tree root -cd "$SCRIPT_PATH/../../" - -download_binary "$old_ver" - -run_test $old_ver "backward" - -echo " === Clear GreptimeDB data before running forward compatibility test" -rm -rf /tmp/greptimedb-standalone - -run_test $old_ver "forward" - -echo "Compatibility test run successfully!" diff --git a/tests/compat/util.sh b/tests/compat/util.sh deleted file mode 100755 index 94ec7ea3a2d..00000000000 --- a/tests/compat/util.sh +++ /dev/null @@ -1,82 +0,0 @@ -#!/bin/bash - -# Assemble the GreptimeDB binary download URL for a specific version. -binary_url() { - local ver="$1" - local bin_tar="greptime-$(uname -s | tr '[:upper:]' '[:lower:]')-amd64-v$ver.tar.gz" - echo "https://github.com/GreptimeTeam/greptimedb/releases/download/v$ver/$bin_tar" -} - -# Download a specific version of GreptimeDB binary tar file, untar it to folder `./bins/$ver`. -# `ver` is semver without prefix `v` -download_binary() { - local ver="$1" - local url="$(binary_url $ver)" - local bin_tar="greptime-$(uname -s | tr '[:upper:]' '[:lower:]')-amd64-v$ver.tar.gz" - - if [ -f ./bins/$ver/greptime ]; then - echo " === binaries exist: $(ls ./bins/$ver/* | tr '\n' ' ')" - chmod +x ./bins/$ver/* - return - fi - - if [ -f "$bin_tar" ]; then - echo " === tar file exists: $bin_tar" - else - echo " === Download binary ver: $ver" - echo " === Download binary url: $url" - curl --connect-timeout 5 --retry 5 --retry-delay 1 -L "$url" -o "$bin_tar" - fi - - mkdir -p ./bins/$ver - tar -xf "$bin_tar" --strip-components=1 -C ./bins/$ver - - echo " === unpacked: ./bins/$ver:" - ls -lh ./bins/$ver - - chmod +x ./bins/$ver/* -} - -# Test data compatibility that: -# - the data written by an old version of GreptimeDB can be read by the current one -# - the data written by the current version of GreptimeDB can be read by an old one ("forward" compatibility) -run_test() { - local old_ver="$1" - local forward="$2" - - local write_case_dir="./tests/compat/case/write" - local read_case_dir="./tests/compat/case/read" - local bin_old="./bins/$old_ver/greptime" - local bin_new="./bins/current/greptime" - local runner="./bins/current/sqlness-runner" - - echo " === Test with:" - echo " === old greptimedb version:" - "$bin_old" --version - echo " === new greptimedb version:" - "$bin_new" --version - - # "forward" means we are testing forward compatibility: - # the data generated by current version GreptimeDB can be used by old. - # So we run new GreptimeDB binary first to write, then run old to read. - # And the opposite for backward compatibility. - if [ "$forward" == "forward" ] - then - echo " === Running forward compat test ..." - echo " === Run test: write with current GreptimeDB" - $runner bare --bins-dir $(dirname $bin_new) --case-dir $write_case_dir - else - echo " === Running backward compat test ..." - echo " === Run test: write with old GreptimeDB" - $runner bare --bins-dir $(dirname $bin_old) --case-dir $write_case_dir - fi - - if [ "$forward" == 'forward' ] - then - echo " === Run test: read with old GreptimeDB" - $runner bare --bins-dir $(dirname $bin_old) --case-dir $read_case_dir - else - echo " === Run test: read with current GreptimeDB" - $runner bare --bins-dir $(dirname $bin_new) --case-dir $read_case_dir - fi -} diff --git a/tests/upgrade-compat/distributed/common b/tests/upgrade-compat/distributed/common deleted file mode 120000 index 2b0920287dc..00000000000 --- a/tests/upgrade-compat/distributed/common +++ /dev/null @@ -1 +0,0 @@ -../standalone/common \ No newline at end of file diff --git a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result b/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result deleted file mode 100644 index 046255a641f..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.10.2 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql b/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql deleted file mode 100644 index 1907533b159..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_0_10_2.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.10.2 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result b/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result deleted file mode 100644 index 7ce230a6881..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.11.0 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql b/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql deleted file mode 100644 index 963170fdf58..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_11_0.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.11.0 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result b/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result deleted file mode 100644 index 41b81f01c08..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.result +++ /dev/null @@ -1,137 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -Affected Rows: 0 - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -Affected Rows: 0 - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -Affected Rows: 0 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -+---------------------+-----------------------------------------------------------+ -| Table | Create Table | -+---------------------+-----------------------------------------------------------+ -| mito_system_metrics | CREATE TABLE IF NOT EXISTS "mito_system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------------+-----------------------------------------------------------+ - -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "cpu_util" DOUBLE NULL, | -| | "host" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host") | -| | ) | -| | | -| | ENGINE=metric | -| | WITH( | -| | on_physical_table = 'phy' | -| | ) | -+----------------+-----------------------------------------------------------+ - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -Affected Rows: 3 - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -Affected Rows: 3 - -SELECT * FROM mito_system_metrics; - -+-------+-------+----------+-------------+-----------+-------------------------+ -| host | idc | cpu_util | memory_util | disk_util | ts | -+-------+-------+----------+-------------+-----------+-------------------------+ -| host1 | idc_a | 11.8 | 10.3 | 10.3 | 2022-11-03T03:39:57.450 | -| host1 | idc_b | 50.0 | 66.7 | 40.6 | 2022-11-03T03:39:57.450 | -| host2 | idc_a | 80.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host3 | idc_a | 90.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host4 | idc_a | 70.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -| host5 | idc_a | 60.0 | 70.3 | 90.0 | 2022-11-03T03:39:57.450 | -+-------+-------+----------+-------------+-----------+-------------------------+ - -SELECT * FROM system_metrics; - -+----------+-------+-------------------------+ -| cpu_util | host | ts | -+----------+-------+-------------------------+ -| 80.0 | host2 | 2022-11-03T03:39:57.450 | -| 70.0 | host4 | 2022-11-03T03:39:57.450 | -| 60.0 | host5 | 2022-11-03T03:39:57.450 | -| 90.0 | host3 | 2022-11-03T03:39:57.450 | -| 50.0 | host1 | 2022-11-03T03:39:57.450 | -+----------+-------+-------------------------+ - -DROP TABLE mito_system_metrics; - -Affected Rows: 0 - -DROP TABLE system_metrics; - -Affected Rows: 0 - -DROP TABLE phy; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql b/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql deleted file mode 100644 index 9908085213e..00000000000 --- a/tests/upgrade-compat/standalone/common/table_engine_v0_9_5.sql +++ /dev/null @@ -1,60 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE mito_system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -)ENGINE=mito; - -INSERT INTO mito_system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -CREATE TABLE phy (ts timestamp time index, cpu_util double) engine=metric with ("physical_metric_table" = ""); - -CREATE TABLE system_metrics ( - host STRING, - cpu_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host), - TIME INDEX(ts) -)ENGINE=metric with ("on_physical_table" = "phy"); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host1', 11.8, 1667446797450), - ('host2', 80.0, 1667446797450), - ('host1', 50.0, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE mito_system_metrics; - -SHOW CREATE TABLE system_metrics; - -INSERT INTO mito_system_metrics -VALUES - ("host3", "idc_a", 90.0, 70.3, 90.0, 1667446797450), - ("host4", "idc_a", 70.0, 70.3, 90.0, 1667446797450), - ("host5", "idc_a", 60.0, 70.3, 90.0, 1667446797450); - -INSERT INTO system_metrics (host, cpu_util, ts) -VALUES - ('host3', 90.0, 1667446797450), - ('host4', 70.0, 1667446797450), - ('host5', 60.0, 1667446797450); - -SELECT * FROM mito_system_metrics; - -SELECT * FROM system_metrics; - -DROP TABLE mito_system_metrics; - -DROP TABLE system_metrics; - -DROP TABLE phy; diff --git a/tests/upgrade-compat/standalone/common/test_alter_drop_column.result b/tests/upgrade-compat/standalone/common/test_alter_drop_column.result deleted file mode 100644 index 605f8711642..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_drop_column.result +++ /dev/null @@ -1,62 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_drop_test ( - id INT, - val_to_drop STRING, - val_to_keep STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO alter_drop_test(id, val_to_drop, val_to_keep, ts) VALUES (1, 'a', 'x', 1672531200000); - -Affected Rows: 1 - --- SQLNESS ARG version=latest -ALTER TABLE alter_drop_test DROP COLUMN val_to_drop; - -Affected Rows: 0 - -SHOW CREATE TABLE alter_drop_test; - -+-----------------+------------------------------------------------+ -| Table | Create Table | -+-----------------+------------------------------------------------+ -| alter_drop_test | CREATE TABLE IF NOT EXISTS "alter_drop_test" ( | -| | "id" INT NULL, | -| | "val_to_keep" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("id") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+-----------------+------------------------------------------------+ - -SELECT * FROM alter_drop_test ORDER BY id; - -+----+-------------+---------------------+ -| id | val_to_keep | ts | -+----+-------------+---------------------+ -| 1 | x | 2023-01-01T00:00:00 | -+----+-------------+---------------------+ - -INSERT INTO alter_drop_test(id, val_to_keep, ts) VALUES (2, 'y', 1672531201000); - -Affected Rows: 1 - -SELECT * FROM alter_drop_test ORDER BY id; - -+----+-------------+---------------------+ -| id | val_to_keep | ts | -+----+-------------+---------------------+ -| 1 | x | 2023-01-01T00:00:00 | -| 2 | y | 2023-01-01T00:00:01 | -+----+-------------+---------------------+ - -DROP TABLE alter_drop_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql b/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql deleted file mode 100644 index df52affe522..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_drop_column.sql +++ /dev/null @@ -1,23 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_drop_test ( - id INT, - val_to_drop STRING, - val_to_keep STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO alter_drop_test(id, val_to_drop, val_to_keep, ts) VALUES (1, 'a', 'x', 1672531200000); - --- SQLNESS ARG version=latest -ALTER TABLE alter_drop_test DROP COLUMN val_to_drop; - -SHOW CREATE TABLE alter_drop_test; - -SELECT * FROM alter_drop_test ORDER BY id; - -INSERT INTO alter_drop_test(id, val_to_keep, ts) VALUES (2, 'y', 1672531201000); - -SELECT * FROM alter_drop_test ORDER BY id; - -DROP TABLE alter_drop_test; diff --git a/tests/upgrade-compat/standalone/common/test_alter_table.result b/tests/upgrade-compat/standalone/common/test_alter_table.result deleted file mode 100644 index df9bbe1a91c..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_table.result +++ /dev/null @@ -1,54 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_test ( - id INT, - val STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO alter_test(id, val, ts) VALUES (1, 'a', 1672531200000); - -Affected Rows: 1 - --- SQLNESS ARG version=latest -ALTER TABLE alter_test ADD COLUMN new_col INT; - -Affected Rows: 0 - -SHOW CREATE TABLE alter_test; - -+------------+-------------------------------------------+ -| Table | Create Table | -+------------+-------------------------------------------+ -| alter_test | CREATE TABLE IF NOT EXISTS "alter_test" ( | -| | "id" INT NULL, | -| | "val" STRING NULL, | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "new_col" INT NULL, | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("id") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+------------+-------------------------------------------+ - -INSERT INTO alter_test(id, val, ts, new_col) VALUES (2, 'b', 1672531201000, 100); - -Affected Rows: 1 - -SELECT * FROM alter_test ORDER BY id; - -+----+-----+---------------------+---------+ -| id | val | ts | new_col | -+----+-----+---------------------+---------+ -| 1 | a | 2023-01-01T00:00:00 | | -| 2 | b | 2023-01-01T00:00:01 | 100 | -+----+-----+---------------------+---------+ - -DROP TABLE alter_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_alter_table.sql b/tests/upgrade-compat/standalone/common/test_alter_table.sql deleted file mode 100644 index 41b1555d8a8..00000000000 --- a/tests/upgrade-compat/standalone/common/test_alter_table.sql +++ /dev/null @@ -1,20 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE alter_test ( - id INT, - val STRING, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO alter_test(id, val, ts) VALUES (1, 'a', 1672531200000); - --- SQLNESS ARG version=latest -ALTER TABLE alter_test ADD COLUMN new_col INT; - -SHOW CREATE TABLE alter_test; - -INSERT INTO alter_test(id, val, ts, new_col) VALUES (2, 'b', 1672531201000, 100); - -SELECT * FROM alter_test ORDER BY id; - -DROP TABLE alter_test; diff --git a/tests/upgrade-compat/standalone/common/test_simple.result b/tests/upgrade-compat/standalone/common/test_simple.result deleted file mode 100644 index ff2c340598a..00000000000 --- a/tests/upgrade-compat/standalone/common/test_simple.result +++ /dev/null @@ -1,47 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -); - -Affected Rows: 0 - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - -Affected Rows: 3 - --- SQLNESS ARG version=latest -SHOW CREATE TABLE system_metrics; - -+----------------+-----------------------------------------------------------+ -| Table | Create Table | -+----------------+-----------------------------------------------------------+ -| system_metrics | CREATE TABLE IF NOT EXISTS "system_metrics" ( | -| | "host" STRING NULL, | -| | "idc" STRING NULL, | -| | "cpu_util" DOUBLE NULL, | -| | "memory_util" DOUBLE NULL, | -| | "disk_util" DOUBLE NULL, | -| | "ts" TIMESTAMP(3) NOT NULL DEFAULT current_timestamp(), | -| | TIME INDEX ("ts"), | -| | PRIMARY KEY ("host", "idc") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+----------------+-----------------------------------------------------------+ - -DROP TABLE system_metrics; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_simple.sql b/tests/upgrade-compat/standalone/common/test_simple.sql deleted file mode 100644 index 0f8daa0985a..00000000000 --- a/tests/upgrade-compat/standalone/common/test_simple.sql +++ /dev/null @@ -1,22 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE system_metrics ( - host STRING, - idc STRING, - cpu_util DOUBLE, - memory_util DOUBLE, - disk_util DOUBLE, - ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), - PRIMARY KEY(host, idc), - TIME INDEX(ts) -); - -INSERT INTO system_metrics -VALUES - ("host1", "idc_a", 11.8, 10.3, 10.3, 1667446797450), - ("host2", "idc_a", 80.0, 70.3, 90.0, 1667446797450), - ("host1", "idc_b", 50.0, 66.7, 40.6, 1667446797450); - --- SQLNESS ARG version=latest -SHOW CREATE TABLE system_metrics; - -DROP TABLE system_metrics; diff --git a/tests/upgrade-compat/standalone/common/test_ttl.result b/tests/upgrade-compat/standalone/common/test_ttl.result deleted file mode 100644 index c368d13e4c1..00000000000 --- a/tests/upgrade-compat/standalone/common/test_ttl.result +++ /dev/null @@ -1,153 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE test_ttl_0s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '0 second'); - -Affected Rows: 0 - -CREATE TABLE test_ttl_1s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '1 second'); - -Affected Rows: 0 - -CREATE TABLE test_ttl_none(ts TIMESTAMP TIME INDEX, val INT); - -Affected Rows: 0 - -CREATE DATABASE ttl_db_1s WITH (ttl = '1 second'); - -Affected Rows: 1 - -CREATE DATABASE ttl_db_0s WITH (ttl = '0 second'); - -Affected Rows: 1 - -CREATE DATABASE ttl_db_none; - -Affected Rows: 1 - --- SQLNESS ARG version=latest -SHOW TABLES; - -+------------------+ -| Tables_in_public | -+------------------+ -| numbers | -| test_ttl_0s | -| test_ttl_1s | -| test_ttl_none | -+------------------+ - -SHOW CREATE TABLE test_ttl_1s; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| test_ttl_1s | CREATE TABLE IF NOT EXISTS "test_ttl_1s" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | ttl = '1s' | -| | ) | -+-------------+--------------------------------------------+ - -SHOW CREATE TABLE test_ttl_0s; - -+-------------+--------------------------------------------+ -| Table | Create Table | -+-------------+--------------------------------------------+ -| test_ttl_0s | CREATE TABLE IF NOT EXISTS "test_ttl_0s" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | WITH( | -| | ttl = '0s' | -| | ) | -+-------------+--------------------------------------------+ - -SHOW CREATE TABLE test_ttl_none; - -+---------------+----------------------------------------------+ -| Table | Create Table | -+---------------+----------------------------------------------+ -| test_ttl_none | CREATE TABLE IF NOT EXISTS "test_ttl_none" ( | -| | "ts" TIMESTAMP(3) NOT NULL, | -| | "val" INT NULL, | -| | TIME INDEX ("ts") | -| | ) | -| | | -| | ENGINE=mito | -| | | -+---------------+----------------------------------------------+ - -DROP TABLE test_ttl_1s; - -Affected Rows: 0 - -DROP TABLE test_ttl_0s; - -Affected Rows: 0 - -DROP TABLE test_ttl_none; - -Affected Rows: 0 - -SHOW DATABASES; - -+--------------------+ -| Database | -+--------------------+ -| greptime_private | -| information_schema | -| public | -| ttl_db_0s | -| ttl_db_1s | -| ttl_db_none | -+--------------------+ - -SHOW CREATE DATABASE ttl_db_1s; - -+-----------+-----------------------------------------+ -| Database | Create Database | -+-----------+-----------------------------------------+ -| ttl_db_1s | CREATE DATABASE IF NOT EXISTS ttl_db_1s | -| | WITH( | -| | ttl = '1s' | -| | ) | -+-----------+-----------------------------------------+ - -SHOW CREATE DATABASE ttl_db_0s; - -+-----------+-----------------------------------------+ -| Database | Create Database | -+-----------+-----------------------------------------+ -| ttl_db_0s | CREATE DATABASE IF NOT EXISTS ttl_db_0s | -| | WITH( | -| | ttl = '0s' | -| | ) | -+-----------+-----------------------------------------+ - -SHOW CREATE DATABASE ttl_db_none; - -+-------------+-------------------------------------------+ -| Database | Create Database | -+-------------+-------------------------------------------+ -| ttl_db_none | CREATE DATABASE IF NOT EXISTS ttl_db_none | -+-------------+-------------------------------------------+ - -DROP DATABASE ttl_db_1s; - -Affected Rows: 0 - -DROP DATABASE ttl_db_0s; - -Affected Rows: 0 - -DROP DATABASE ttl_db_none; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_ttl.sql b/tests/upgrade-compat/standalone/common/test_ttl.sql deleted file mode 100644 index 3462fd22444..00000000000 --- a/tests/upgrade-compat/standalone/common/test_ttl.sql +++ /dev/null @@ -1,42 +0,0 @@ - --- SQLNESS ARG version=v0.9.5 -CREATE TABLE test_ttl_0s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '0 second'); - -CREATE TABLE test_ttl_1s(ts TIMESTAMP TIME INDEX, val INT) WITH (ttl = '1 second'); - -CREATE TABLE test_ttl_none(ts TIMESTAMP TIME INDEX, val INT); - -CREATE DATABASE ttl_db_1s WITH (ttl = '1 second'); - -CREATE DATABASE ttl_db_0s WITH (ttl = '0 second'); - -CREATE DATABASE ttl_db_none; - --- SQLNESS ARG version=latest -SHOW TABLES; - -SHOW CREATE TABLE test_ttl_1s; - -SHOW CREATE TABLE test_ttl_0s; - -SHOW CREATE TABLE test_ttl_none; - -DROP TABLE test_ttl_1s; - -DROP TABLE test_ttl_0s; - -DROP TABLE test_ttl_none; - -SHOW DATABASES; - -SHOW CREATE DATABASE ttl_db_1s; - -SHOW CREATE DATABASE ttl_db_0s; - -SHOW CREATE DATABASE ttl_db_none; - -DROP DATABASE ttl_db_1s; - -DROP DATABASE ttl_db_0s; - -DROP DATABASE ttl_db_none; diff --git a/tests/upgrade-compat/standalone/common/test_update.result b/tests/upgrade-compat/standalone/common/test_update.result deleted file mode 100644 index 872878cb152..00000000000 --- a/tests/upgrade-compat/standalone/common/test_update.result +++ /dev/null @@ -1,32 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE update_test ( - id INT, - val INT, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -Affected Rows: 0 - -INSERT INTO update_test(id, val, ts) VALUES (1, 10, 1672531200000), (2, 20, 1672531201000); - -Affected Rows: 2 - --- SQLNESS ARG version=latest -UPDATE update_test SET val = 100 WHERE id = 1; - -Error: 1001(Unsupported), SQL statement is not supported, keyword: UPDATE - -SELECT * FROM update_test ORDER BY id; - -+----+-----+---------------------+ -| id | val | ts | -+----+-----+---------------------+ -| 1 | 10 | 2023-01-01T00:00:00 | -| 2 | 20 | 2023-01-01T00:00:01 | -+----+-----+---------------------+ - -DROP TABLE update_test; - -Affected Rows: 0 - diff --git a/tests/upgrade-compat/standalone/common/test_update.sql b/tests/upgrade-compat/standalone/common/test_update.sql deleted file mode 100644 index 33eedcd2061..00000000000 --- a/tests/upgrade-compat/standalone/common/test_update.sql +++ /dev/null @@ -1,16 +0,0 @@ --- SQLNESS ARG version=v0.9.5 -CREATE TABLE update_test ( - id INT, - val INT, - ts TIMESTAMP TIME INDEX, - PRIMARY KEY(id) -); - -INSERT INTO update_test(id, val, ts) VALUES (1, 10, 1672531200000), (2, 20, 1672531201000); - --- SQLNESS ARG version=latest -UPDATE update_test SET val = 100 WHERE id = 1; - -SELECT * FROM update_test ORDER BY id; - -DROP TABLE update_test; From 079ffec1167ab790a03a9833026d3b9311e64020 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 28 Sep 2026 12:35:47 +0000 Subject: [PATCH 063/133] chore: remove dead code left behind by removed features (#9377) * chore: remove dead code left behind by removed features Signed-off-by: Dennis Zhuang * chore: remove unused RouteInfoCorrupted error Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.toml | 1 - src/cmd/src/datanode/objbench.rs | 1 - src/common/base/Cargo.toml | 3 - src/common/meta/src/error.rs | 8 - src/common/meta/src/rpc/router.rs | 228 -------------------- src/mito-codec/src/key_values.rs | 55 ----- src/mito-codec/src/row_converter/sparse.rs | 34 --- src/mito2/src/access_layer.rs | 2 - src/mito2/src/cache.rs | 2 - src/mito2/src/cache/write_cache.rs | 3 - src/mito2/src/compaction/compactor.rs | 2 - src/mito2/src/flush.rs | 1 - src/mito2/src/manifest/manager.rs | 2 +- src/mito2/src/read.rs | 240 +-------------------- src/mito2/src/read/last_row.rs | 155 +------------ src/mito2/src/read/scan_util.rs | 3 - src/mito2/src/read/stream.rs | 1 - src/mito2/src/sst/index.rs | 1 - src/mito2/src/sst/parquet/file_range.rs | 2 +- src/mito2/src/test_util.rs | 43 +--- src/object-store/src/lib.rs | 1 - src/object-store/src/metrics.rs | 15 -- src/operator/src/statement/ddl.rs | 52 +---- src/pipeline/Cargo.toml | 1 - src/pipeline/src/etl.rs | 18 -- src/puffin/Cargo.toml | 3 - src/query/Cargo.toml | 1 - src/servers/Cargo.toml | 1 - src/servers/src/http/pprof/nix.rs | 123 ----------- src/servers/src/opentsdb/codec.rs | 231 -------------------- src/servers/src/postgres/types.rs | 39 ---- src/servers/src/query_handler.rs | 8 +- src/session/Cargo.toml | 3 - src/session/src/context.rs | 1 - src/sql/src/statements/create.rs | 19 +- src/store-api/src/manifest/action.rs | 114 ---------- src/store-api/src/manifest/storage.rs | 78 ------- src/store-api/src/mito_engine_options.rs | 2 - src/store-api/src/region_request.rs | 7 +- src/table/src/requests.rs | 20 -- tests-integration/Cargo.toml | 1 - 41 files changed, 10 insertions(+), 1515 deletions(-) delete mode 100644 src/object-store/src/metrics.rs delete mode 100644 src/servers/src/http/pprof/nix.rs delete mode 100644 src/store-api/src/manifest/action.rs delete mode 100644 src/store-api/src/manifest/storage.rs diff --git a/Cargo.toml b/Cargo.toml index 5f497a3faff..4f842313d8c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -272,7 +272,6 @@ url = "2.3" # completion" panic (greptimedb#8265 / grpc-rust#2522). Remove once tonic ships a fix. tonic = { version = "0.14, <0.14.5", features = ["tls-aws-lc", "gzip", "zstd"] } tower = "0.5" -tower-http = "0.6" tracing = "0.1" tracing-appender = "0.2" tracing-opentelemetry = "0.33.0" diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index d9f7a00c8c4..03696c96fc0 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -200,7 +200,6 @@ impl ObjbenchCommand { metadata: region_meta, source, cache_manager, - storage: None, max_sequence: None, sst_write_format: FormatType::PrimaryKey, index_options: Default::default(), diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 760241a3e6a..702171e9d67 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -4,9 +4,6 @@ version.workspace = true edition.workspace = true license.workspace = true -[features] -testing = [] - [lints] workspace = true diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3f4868ffba0..62559fd6f1b 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -355,13 +355,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Corrupted table route data, err: {}", err_msg))] - RouteInfoCorrupted { - err_msg: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Illegal state from server, code: {}, error: {}", code, err_msg))] IllegalServerState { code: i32, @@ -1213,7 +1206,6 @@ impl ErrorExt for Error { SerdeJson { .. } | ParseOption { .. } - | RouteInfoCorrupted { .. } | InvalidProtoMsg { .. } | InvalidMetadata { .. } | Unexpected { .. } diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 519b12e081f..929f57515ed 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -14,22 +14,15 @@ use std::collections::{BTreeMap, HashMap, HashSet}; -use api::v1::meta::{ - Partition as PbPartition, Peer as PbPeer, Region as PbRegion, Table as PbTable, - TableRoute as PbTableRoute, -}; use common_time::util::current_time_millis; use derive_builder::Builder; use serde::ser::SerializeSeq; use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use snafu::OptionExt; use store_api::region_engine::RegionRole; use store_api::storage::{RegionId, RegionNumber}; use strum::AsRefStr; -use table::table_name::TableName; use crate::DatanodeId; -use crate::error::{self, Result}; use crate::key::RegionDistribution; use crate::peer::Peer; @@ -62,13 +55,6 @@ pub fn region_distribution(region_routes: &[RegionRoute]) -> RegionDistribution regions_id_map } -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] -pub struct TableRoute { - pub table: Table, - pub region_routes: Vec, - region_leaders: HashMap>, -} - /// Returns the leader peers of the table. pub fn find_leaders(region_routes: &[RegionRoute]) -> HashSet { region_routes @@ -167,94 +153,6 @@ pub fn find_follower_regions(region_routes: &[RegionRoute], datanode: &Peer) -> .collect() } -impl TableRoute { - pub fn new(table: Table, region_routes: Vec) -> Self { - let region_leaders = region_routes - .iter() - .map(|x| (x.region.id.region_number(), x.leader_peer.clone())) - .collect::>(); - Self { - table, - region_routes, - region_leaders, - } - } - - pub fn try_from_raw(peers: &[PbPeer], table_route: PbTableRoute) -> Result { - let table = table_route - .table - .context(error::RouteInfoCorruptedSnafu { - err_msg: "'table' is empty in table route", - })? - .try_into()?; - - let mut region_routes = Vec::with_capacity(table_route.region_routes.len()); - for region_route in table_route.region_routes.into_iter() { - let region = region_route - .region - .context(error::RouteInfoCorruptedSnafu { - err_msg: "'region' is empty in region route", - })? - .into(); - - let leader_peer = peers.get(region_route.leader_peer_index as usize).cloned(); - - let follower_peers = region_route - .follower_peer_indexes - .into_iter() - .filter_map(|x| peers.get(x as usize).cloned()) - .collect::>(); - - region_routes.push(RegionRoute { - region, - leader_peer, - follower_peers, - leader_state: None, - leader_down_since: None, - write_route_policy: None, - }); - } - - Ok(Self::new(table, region_routes)) - } -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] -pub struct Table { - pub id: u64, - pub table_name: TableName, - #[serde(serialize_with = "as_utf8", deserialize_with = "from_utf8")] - pub table_schema: Vec, -} - -impl TryFrom for Table { - type Error = error::Error; - - fn try_from(t: PbTable) -> Result { - let table_name = t - .table_name - .context(error::RouteInfoCorruptedSnafu { - err_msg: "table name required", - })? - .into(); - Ok(Self { - id: t.id, - table_name, - table_schema: t.table_schema, - }) - } -} - -impl From
for PbTable { - fn from(table: Table) -> Self { - PbTable { - id: table.id, - table_name: Some(table.table_name.into()), - table_schema: table.table_schema, - } - } -} - #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Builder)] pub struct RegionRoute { pub region: Region, @@ -430,18 +328,6 @@ impl RegionRoute { } } -pub struct RegionRoutes(pub Vec); - -impl RegionRoutes { - pub fn region_leader_map(&self) -> HashMap { - convert_to_region_leader_map(&self.0) - } - - pub fn find_region_leader(&self, region_number: RegionNumber) -> Option<&Peer> { - self.region_leader_map().get(®ion_number).copied() - } -} - #[derive(Debug, Clone, Default, PartialEq, Serialize)] pub struct Region { pub id: RegionId, @@ -508,49 +394,6 @@ impl Region { } } -/// Gets the partition expression of the `PbRegion` in compatible mode. -#[allow(deprecated)] -pub fn pb_region_partition_expr(r: &PbRegion) -> String { - if let Some(partition) = &r.partition { - if !partition.expression.is_empty() { - partition.expression.clone() - } else if !partition.value_list.is_empty() { - String::from_utf8_lossy(&partition.value_list[0]).to_string() - } else { - "".to_string() - } - } else { - "".to_string() - } -} - -impl From for Region { - fn from(r: PbRegion) -> Self { - let partition_expr = pb_region_partition_expr(&r); - Self { - id: r.id.into(), - name: r.name, - partition_expr, - attrs: r.attrs.into_iter().collect::>(), - } - } -} - -impl From for PbRegion { - fn from(region: Region) -> Self { - let partition_expr = region.partition_expr(); - Self { - id: region.id.into(), - name: region.name, - partition: Some(PbPartition { - expression: partition_expr, - ..Default::default() - }), - attrs: region.attrs.into_iter().collect::>(), - } - } -} - #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] pub struct LegacyPartition { #[serde(serialize_with = "as_utf8_vec", deserialize_with = "from_utf8_vec")] @@ -559,23 +402,6 @@ pub struct LegacyPartition { pub value_list: Vec>, } -fn as_utf8(val: &[u8], serializer: S) -> std::result::Result { - serializer.serialize_str( - String::from_utf8(val.to_vec()) - .unwrap_or_else(|_| "".to_string()) - .as_str(), - ) -} - -pub fn from_utf8<'de, D>(deserializer: D) -> std::result::Result, D::Error> -where - D: Deserializer<'de>, -{ - let s = String::deserialize(deserializer)?; - - Ok(s.into_bytes()) -} - fn as_utf8_vec( val: &[Vec], serializer: S, @@ -888,58 +714,4 @@ mod tests { assert_eq!(got, p); } - - #[test] - #[allow(deprecated)] - fn test_region_partition_expr() { - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: None, - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), ""); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), ""); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, ""); - - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: Some(PbPartition { - column_list: vec![b"a".to_vec()], - value_list: vec![b"{}".to_vec()], - expression: Default::default(), - }), - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), "{}"); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), "{}"); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, "{}"); - - let r = PbRegion { - id: 1, - name: "r1".to_string(), - partition: Some(PbPartition { - column_list: vec![b"a".to_vec()], - value_list: vec![b"{}".to_vec()], - expression: "a>b".to_string(), - }), - attrs: Default::default(), - }; - assert_eq!(pb_region_partition_expr(&r), "a>b"); - - let r2: Region = r.into(); - assert_eq!(r2.partition_expr(), "a>b"); - - let r3: PbRegion = r2.into(); - assert_eq!(r3.partition.as_ref().unwrap().expression, "a>b"); - } } diff --git a/src/mito-codec/src/key_values.rs b/src/mito-codec/src/key_values.rs index d66110bacfa..aba3cff988c 100644 --- a/src/mito-codec/src/key_values.rs +++ b/src/mito-codec/src/key_values.rs @@ -98,61 +98,6 @@ impl KeyValues { } } -/// Key value view of a mutation. -#[derive(Debug)] -pub struct KeyValuesRef<'a> { - /// Mutation to read. - /// - /// This mutation must be a valid mutation and rows in the mutation - /// must not be `None`. - mutation: &'a Mutation, - /// Key value read helper. - helper: SparseReadRowHelper, - /// Primary key encoding hint. - primary_key_encoding: PrimaryKeyEncoding, -} - -impl<'a> KeyValuesRef<'a> { - /// Creates [crate::memtable::KeyValues] from specific `mutation`. - /// - /// Returns `None` if `rows` of the `mutation` is `None`. - pub fn new(metadata: &RegionMetadata, mutation: &'a Mutation) -> Option> { - let rows = mutation.rows.as_ref()?; - let primary_key_encoding = - infer_primary_key_encoding_from_hint(mutation.write_hint.as_ref()); - let helper = SparseReadRowHelper::new(metadata, rows, primary_key_encoding); - - Some(KeyValuesRef { - mutation, - helper, - primary_key_encoding, - }) - } - - /// Returns a key value iterator. - pub fn iter(&self) -> impl Iterator> { - let rows = self.mutation.rows.as_ref().unwrap(); - let schema = &rows.schema; - rows.rows.iter().enumerate().map(|(idx, row)| { - KeyValue { - row, - schema, - helper: &self.helper, - sequence: self.mutation.sequence + idx as u64, // Calculate sequence for each row. - // Safety: This is a valid mutation. - op_type: OpType::try_from(self.mutation.op_type).unwrap(), - primary_key_encoding: self.primary_key_encoding, - } - }) - } - - /// Returns number of rows. - pub fn num_rows(&self) -> usize { - // Safety: rows is not None. - self.mutation.rows.as_ref().unwrap().rows.len() - } -} - /// Key value view of a row. /// /// A key value view divides primary key columns and field (value) columns. diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index 1be2fbacb2b..5497a7dc034 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -581,40 +581,6 @@ impl PrimaryKeyCodec for SparsePrimaryKeyCodec { } } -/// Field with column id. -pub struct FieldWithId { - pub field: SortField, - pub column_id: ColumnId, -} - -/// A special encoder for memtable. -pub struct SparseEncoder { - fields: Vec, -} - -impl SparseEncoder { - pub fn new(fields: Vec) -> Self { - Self { fields } - } - - pub fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec) -> Result<()> - where - I: Iterator>, - { - let mut serializer = Serializer::new(buffer); - for (value, field) in row.zip(self.fields.iter()) { - if !value.is_null() { - field - .column_id - .serialize(&mut serializer) - .context(SerializeFieldSnafu)?; - field.field.serialize(&mut serializer, &value)?; - } - } - Ok(()) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 06fdfd79240..3959e0c26a4 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -578,8 +578,6 @@ pub struct SstWriteRequest { pub metadata: RegionMetadataRef, pub source: FlatSource, pub cache_manager: CacheManagerRef, - #[allow(dead_code)] - pub storage: Option, /// Optional uniform row sequence for writes that do not preserve sequences. /// Compaction passes `None` to retain the reader's effective input sequences. pub max_sequence: Option, diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 70dee2f7dd8..cf0e8de3458 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -898,7 +898,6 @@ impl CacheStrategy { /// Calls [CacheManager::get_range_result()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. - #[allow(dead_code)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -1391,7 +1390,6 @@ impl CacheManager { } /// Gets cached result for range scan. - #[allow(dead_code)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index e625556e087..afe943eef50 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -701,7 +701,6 @@ mod tests { op_type, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: Default::default(), @@ -819,7 +818,6 @@ mod tests { op_type: OperationType::Flush, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: cache_manager.clone(), @@ -912,7 +910,6 @@ mod tests { op_type: OperationType::Flush, metadata, source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: cache_manager.clone(), diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index c6a9999ee41..1a025bcaeec 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -461,7 +461,6 @@ impl SstMerger for DefaultSstMerger { write_opts: WriteOptions, ) -> Result<(Vec, Vec)> { let region_id = compaction_region.region_id; - let storage = compaction_region.region_options.storage.clone(); let index_options = compaction_region .current_version .options @@ -508,7 +507,6 @@ impl SstMerger for DefaultSstMerger { metadata: region_metadata.clone(), source, cache_manager: compaction_region.cache_manager.clone(), - storage, // Readers resolve file overrides before merge/dedup. Replacing // their effective sequences here could promote old rows above // versions in SSTs that were not part of this merge. diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 66182693f78..a5d9a7d0a73 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -836,7 +836,6 @@ impl RegionFlushTask { metadata: version.metadata.clone(), source, cache_manager: self.cache_manager.clone(), - storage: version.options.storage.clone(), max_sequence: Some(max_sequence), sst_write_format: if flat_format { FormatType::Flat diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index c6a0d3f3228..436488f3c3c 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -49,7 +49,7 @@ pub struct RegionManifestOptions { pub manifest_dir: String, pub object_store: ObjectStore, pub compress_type: CompressionType, - /// Interval of version ([ManifestVersion](store_api::manifest::ManifestVersion)) between two checkpoints. + /// Interval of version ([ManifestVersion](store_api::ManifestVersion)) between two checkpoints. /// Set to 0 to disable checkpoint. pub checkpoint_distance: u64, pub remove_file_options: RemoveFileOptions, diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs index d5abe29f021..65d356d79f2 100644 --- a/src/mito2/src/read.rs +++ b/src/mito2/src/read.rs @@ -45,7 +45,6 @@ use std::time::Duration; use api::v1::OpType; use arrow_schema::SchemaRef; -use async_trait::async_trait; use common_time::Timestamp; use datafusion_common::arrow::array::UInt8Array; use datatypes::arrow; @@ -73,7 +72,7 @@ use crate::error::{ ComputeArrowSnafu, ComputeVectorSnafu, ConvertVectorSnafu, DecodeSnafu, InvalidBatchSnafu, Result, }; -use crate::memtable::{BoxedBatchIterator, BoxedRecordBatchIterator}; +use crate::memtable::BoxedRecordBatchIterator; pub(crate) fn timestamp_array_to_i64_slice(arr: &ArrayRef) -> &[i64] { use datatypes::arrow::array::{ @@ -706,86 +705,6 @@ impl Batch { self.sequences.get_data(index).unwrap() } - /// Checks the batch is monotonic by timestamps. - #[cfg(debug_assertions)] - #[allow(dead_code)] - pub(crate) fn check_monotonic(&self) -> Result<(), String> { - use std::cmp::Ordering; - if self.timestamps_native().is_none() { - return Ok(()); - } - - let timestamps = self.timestamps_native().unwrap(); - let sequences = self.sequences.as_arrow().values(); - for (i, window) in timestamps.windows(2).enumerate() { - let current = window[0]; - let next = window[1]; - let current_sequence = sequences[i]; - let next_sequence = sequences[i + 1]; - match current.cmp(&next) { - Ordering::Less => { - // The current timestamp is less than the next timestamp. - continue; - } - Ordering::Equal => { - // The current timestamp is equal to the next timestamp. - if current_sequence < next_sequence { - return Err(format!( - "sequence are not monotonic: ts {} == {} but current sequence {} < {}, index: {}", - current, next, current_sequence, next_sequence, i - )); - } - } - Ordering::Greater => { - // The current timestamp is greater than the next timestamp. - return Err(format!( - "timestamps are not monotonic: {} > {}, index: {}", - current, next, i - )); - } - } - } - - Ok(()) - } - - /// Returns Ok if the given batch is behind the current batch. - #[cfg(debug_assertions)] - #[allow(dead_code)] - pub(crate) fn check_next_batch(&self, other: &Batch) -> Result<(), String> { - // Checks the primary key - if self.primary_key() < other.primary_key() { - return Ok(()); - } - if self.primary_key() > other.primary_key() { - return Err(format!( - "primary key is not monotonic: {:?} > {:?}", - self.primary_key(), - other.primary_key() - )); - } - // Checks the timestamp. - if self.last_timestamp() < other.first_timestamp() { - return Ok(()); - } - if self.last_timestamp() > other.first_timestamp() { - return Err(format!( - "timestamps are not monotonic: {:?} > {:?}", - self.last_timestamp(), - other.first_timestamp() - )); - } - // Checks the sequence. - if self.last_sequence() >= other.first_sequence() { - return Ok(()); - } - Err(format!( - "sequences are not monotonic: {:?} < {:?}", - self.last_sequence(), - other.first_sequence() - )) - } - /// 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<()> { @@ -863,112 +782,6 @@ impl Batch { } } -/// A struct to check the batch is monotonic. -#[cfg(debug_assertions)] -#[derive(Default)] -#[allow(dead_code)] -pub(crate) struct BatchChecker { - last_batch: Option, - start: Option, - end: Option, -} - -#[cfg(debug_assertions)] -#[allow(dead_code)] -impl BatchChecker { - /// Attaches the given start timestamp to the checker. - pub(crate) fn with_start(mut self, start: Option) -> Self { - self.start = start; - self - } - - /// Attaches the given end timestamp to the checker. - pub(crate) fn with_end(mut self, end: Option) -> Self { - self.end = end; - self - } - - /// Returns true if the given batch is monotonic and behind - /// the last batch. - pub(crate) fn check_monotonic(&mut self, batch: &Batch) -> Result<(), String> { - batch.check_monotonic()?; - - if let (Some(start), Some(first)) = (self.start, batch.first_timestamp()) - && start > first - { - return Err(format!( - "batch's first timestamp is before the start timestamp: {:?} > {:?}", - start, first - )); - } - if let (Some(end), Some(last)) = (self.end, batch.last_timestamp()) - && end <= last - { - return Err(format!( - "batch's last timestamp is after the end timestamp: {:?} <= {:?}", - end, last - )); - } - - // Checks the batch is behind the last batch. - // Then Updates the last batch. - let res = self - .last_batch - .as_ref() - .map(|last| last.check_next_batch(batch)) - .unwrap_or(Ok(())); - self.last_batch = Some(batch.clone()); - res - } - - /// Formats current batch and last batch for debug. - pub(crate) fn format_batch(&self, batch: &Batch) -> String { - use std::fmt::Write; - - let mut message = String::new(); - if let Some(last) = &self.last_batch { - write!( - message, - "last_pk: {:?}, last_ts: {:?}, last_seq: {:?}, ", - last.primary_key(), - last.last_timestamp(), - last.last_sequence() - ) - .unwrap(); - } - write!( - message, - "batch_pk: {:?}, batch_ts: {:?}, batch_seq: {:?}", - batch.primary_key(), - batch.timestamps(), - batch.sequences() - ) - .unwrap(); - - message - } - - /// Checks batches from the part range are monotonic. Otherwise, panics. - pub(crate) fn ensure_part_range_batch( - &mut self, - scanner: &str, - region_id: store_api::storage::RegionId, - partition: usize, - part_range: store_api::region_engine::PartitionRange, - batch: &Batch, - ) { - if let Err(e) = self.check_monotonic(batch) { - let err_msg = format!( - "{}: batch is not sorted, {}, region_id: {}, partition: {}, part_range: {:?}", - scanner, e, region_id, partition, part_range, - ); - common_telemetry::error!("{err_msg}, {}", self.format_batch(batch)); - // Only print the number of row in the panic message. - panic!("{err_msg}, batch rows: {}", batch.num_rows()); - } - } -} - /// Len of timestamp in arrow row format. const TIMESTAMP_KEY_LEN: usize = 9; @@ -1169,29 +982,6 @@ impl From for BatchBuilder { } } -/// Async [Batch] reader and iterator wrapper. -/// -/// This is the data source for SST writers or internal readers. -pub enum Source { - /// Source from a [BoxedBatchReader]. - Reader(BoxedBatchReader), - /// Source from a [BoxedBatchIterator]. - Iter(BoxedBatchIterator), - /// Source from a [BoxedBatchStream]. - Stream(BoxedBatchStream), -} - -impl Source { - /// Returns next [Batch] from this data source. - pub async fn next_batch(&mut self) -> Result> { - match self { - Source::Reader(reader) => reader.next_batch().await, - Source::Iter(iter) => iter.next().transpose(), - Source::Stream(stream) => stream.try_next().await, - } - } -} - /// Async [RecordBatch] reader and iterator wrapper for flat format. pub struct FlatSource { schema: SchemaRef, @@ -1249,37 +1039,9 @@ impl FlatSourceInner { } } -/// Async batch reader. -/// -/// The reader must guarantee [Batch]es returned by it have the same schema. -#[async_trait] -pub trait BatchReader: Send { - /// Fetch next [Batch]. - /// - /// Returns `Ok(None)` when the reader has reached its end and calling `next_batch()` - /// again won't return batch again. - /// - /// If `Err` is returned, caller should not call this method again, the implementor - /// may or may not panic in such case. - async fn next_batch(&mut self) -> Result>; -} - -/// Pointer to [BatchReader]. -pub type BoxedBatchReader = Box; - -/// Pointer to a stream that yields [Batch]. -pub type BoxedBatchStream = BoxStream<'static, Result>; - /// Pointer to a stream that yields [RecordBatch]. pub type BoxedRecordBatchStream = BoxStream<'static, Result>; -#[async_trait::async_trait] -impl BatchReader for Box { - async fn next_batch(&mut self) -> Result> { - (**self).next_batch().await - } -} - /// Local metrics for scanners. #[derive(Debug, Default)] pub(crate) struct ScannerMetrics { diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index c526bfcaa86..2bafe02dbbe 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -16,7 +16,6 @@ use std::sync::Arc; -use async_trait::async_trait; use datatypes::arrow::array::{Array, BinaryArray}; use datatypes::arrow::compute::concat_batches; use datatypes::arrow::record_batch::RecordBatch; @@ -30,98 +29,13 @@ use crate::cache::{ }; use crate::error::{ComputeArrowSnafu, Result}; use crate::read::read_columns::JsonTargetTypes; -use crate::read::{ - Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice, -}; +use crate::read::{BoxedRecordBatchStream, timestamp_array_to_i64_slice}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::parquet::format::{PrimaryKeyArray, primary_key_offsets}; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::parquet::reader::FlatRowGroupReader; -/// Reader to keep the last row for each time series. -/// It assumes that batches from the input reader are -/// - sorted -/// - all deleted rows has been filtered. -/// - not empty -/// -/// This reader is different from the [MergeMode](crate::region::options::MergeMode) as -/// it focus on time series (the same key). -#[allow(dead_code)] -pub(crate) struct LastRowReader { - /// Inner reader. - reader: BoxedBatchReader, - /// The last batch pending to return. - selector: LastRowSelector, -} - -#[allow(dead_code)] -impl LastRowReader { - /// Creates a new `LastRowReader`. - pub(crate) fn new(reader: BoxedBatchReader) -> Self { - Self { - reader, - selector: LastRowSelector::default(), - } - } - - /// Returns the last row of the next key. - pub(crate) async fn next_last_row(&mut self) -> Result> { - while let Some(batch) = self.reader.next_batch().await? { - if let Some(yielded) = self.selector.on_next(batch) { - return Ok(Some(yielded)); - } - } - Ok(self.selector.finish()) - } -} - -#[async_trait] -impl BatchReader for LastRowReader { - async fn next_batch(&mut self) -> Result> { - self.next_last_row().await - } -} - -/// Common struct that selects only the last row of each time series. -#[derive(Default)] -pub struct LastRowSelector { - last_batch: Option, -} - -impl LastRowSelector { - /// Handles next batch. Return the yielding batch if present. - pub fn on_next(&mut self, batch: Batch) -> Option { - if let Some(last) = &self.last_batch { - if last.primary_key() == batch.primary_key() { - // Same key, update last batch. - self.last_batch = Some(batch); - None - } else { - // Different key, return the last row in `last` and update `last_batch` by - // current batch. - debug_assert!(!last.is_empty()); - let last_row = last.slice(last.num_rows() - 1, 1); - self.last_batch = Some(batch); - Some(last_row) - } - } else { - self.last_batch = Some(batch); - None - } - } - - /// Finishes the selector and returns the pending batch if any. - pub fn finish(&mut self) -> Option { - if let Some(last) = self.last_batch.take() { - // This is the last key. - let last_row = last.slice(last.num_rows() - 1, 1); - return Some(last_row); - } - None - } -} - /// Cached last row reader for flat format row group. /// If the last rows are already cached (as flat `RecordBatch`), returns cached values. /// Otherwise, reads from the row group, selects last rows, and updates the cache. @@ -514,7 +428,6 @@ fn last_timestamp_start(ts_values: &[i64], range_start: usize, range_end: usize) mod tests { use std::sync::Arc; - use api::v1::OpType; use datatypes::arrow::array::{ ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array, UInt64Array, @@ -523,72 +436,6 @@ mod tests { use datatypes::arrow::record_batch::RecordBatch; use super::*; - use crate::test_util::{VecBatchReader, check_reader_result, new_batch}; - - #[tokio::test] - async fn test_last_row_one_batch() { - let input = [new_batch( - b"k1", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[21, 22], - )]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[new_batch(b"k1", &[2], &[11], &[OpType::Put], &[22])], - ) - .await; - - // Only one row. - let input = [new_batch(b"k1", &[1], &[11], &[OpType::Put], &[21])]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[new_batch(b"k1", &[1], &[11], &[OpType::Put], &[21])], - ) - .await; - } - - #[tokio::test] - async fn test_last_row_multi_batch() { - let input = [ - new_batch( - b"k1", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[21, 22], - ), - new_batch( - b"k1", - &[3, 4], - &[11, 11], - &[OpType::Put, OpType::Put], - &[23, 24], - ), - new_batch( - b"k2", - &[1, 2], - &[11, 11], - &[OpType::Put, OpType::Put], - &[31, 32], - ), - ]; - let reader = VecBatchReader::new(&input); - let mut reader = LastRowReader::new(Box::new(reader)); - check_reader_result( - &mut reader, - &[ - new_batch(b"k1", &[4], &[11], &[OpType::Put], &[24]), - new_batch(b"k2", &[2], &[11], &[OpType::Put], &[32]), - ], - ) - .await; - } /// Helper to build a flat format RecordBatch for testing. fn new_flat_batch(primary_keys: &[&[u8]], timestamps: &[i64], fields: &[i64]) -> RecordBatch { diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index b58757930eb..c5eed8853ae 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1112,21 +1112,18 @@ impl PartitionMetrics { } /// Increments the total bytes added to the range cache. - #[allow(dead_code)] pub(crate) fn inc_range_cache_size(&self, size: usize) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_size += size; } /// Increments the range cache hit counter. - #[allow(dead_code)] pub(crate) fn inc_range_cache_hit(&self) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_hit += 1; } /// Increments the range cache miss counter. - #[allow(dead_code)] pub(crate) fn inc_range_cache_miss(&self) { let mut metrics = self.0.metrics.lock().unwrap(); metrics.range_cache_miss += 1; diff --git a/src/mito2/src/read/stream.rs b/src/mito2/src/read/stream.rs index 27ae0e8c3a1..0c882d374d2 100644 --- a/src/mito2/src/read/stream.rs +++ b/src/mito2/src/read/stream.rs @@ -43,7 +43,6 @@ pub type ScanBatchStream = BoxStream<'static, Result>; pub(crate) struct ConvertBatchStream { inner: ScanBatchStream, projection_mapper: Arc, - #[allow(dead_code)] cache_strategy: CacheStrategy, partition_metrics: PartitionMetrics, pending: VecDeque, diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 32d0419a33c..502f6bcdc05 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -1588,7 +1588,6 @@ mod tests { op_type: OperationType::Flush, metadata: metadata.clone(), source, - storage: None, max_sequence: None, sst_write_format: Default::default(), cache_manager: Default::default(), diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 60b6a955e93..1c15afd386d 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -215,7 +215,7 @@ impl FileRange { selector, Some(TimeSeriesRowSelector::LastRow { after_merge: false }) ) { - // Only use LastRowReader if row group does not contain DELETE, all + // Only use the last row reader if row group does not contain DELETE, all // rows are selected, and filters that still run after this reader // cannot change which row is last. Tag filters are safe because a // tag is constant within a series. Timestamp and field filters are diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index dbfbc53db25..165f2930f41 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -79,7 +79,7 @@ use crate::error::Result; use crate::flush::{WriteBufferManager, WriteBufferManagerRef}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::manifest::storage::{is_checkpoint_file, is_delta_file}; -use crate::read::{Batch, BatchBuilder, BatchReader}; +use crate::read::{Batch, BatchBuilder}; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::sst::FormatType; use crate::sst::file_purger::{FilePurgerRef, NoopFilePurger}; @@ -1106,34 +1106,6 @@ pub(crate) fn ts_ms_value(data: i64) -> v1::Value { } } -/// A reader for test that pop [Batch] from a vector. -pub struct VecBatchReader { - batches: Vec, -} - -impl VecBatchReader { - pub fn new(batches: &[Batch]) -> VecBatchReader { - let batches = batches.iter().rev().cloned().collect(); - - VecBatchReader { batches } - } -} - -#[async_trait::async_trait] -impl BatchReader for VecBatchReader { - async fn next_batch(&mut self) -> Result> { - Ok(self.batches.pop()) - } -} - -impl Iterator for VecBatchReader { - type Item = Result; - - fn next(&mut self) -> Option> { - self.batches.pop().map(Ok) - } -} - pub fn new_batch_builder( primary_key: &[u8], timestamps: &[i64], @@ -1177,19 +1149,6 @@ pub fn new_batch( .unwrap() } -/// Ensure the reader returns batch as `expect`. -pub async fn check_reader_result(reader: &mut R, expect: &[Batch]) { - let mut result = Vec::new(); - while let Some(mut batch) = reader.next_batch().await.unwrap() { - batch.remove_pk_values(); - result.push(batch); - } - - assert_eq!(expect, result); - // Next call to `next_batch()` still returns None. - assert!(reader.next_batch().await.unwrap().is_none()); -} - /// A mock [WriteBufferManager] that supports controlling whether to flush/stall. #[derive(Debug, Default)] pub struct MockWriteBufferManager { diff --git a/src/object-store/src/lib.rs b/src/object-store/src/lib.rs index 14a2c3d1c68..6c8956445d8 100644 --- a/src/object-store/src/lib.rs +++ b/src/object-store/src/lib.rs @@ -24,7 +24,6 @@ pub mod error; pub mod factory; pub mod layers; pub mod manager; -mod metrics; pub mod secure_fs; #[cfg(feature = "testing")] pub mod test_util; diff --git a/src/object-store/src/metrics.rs b/src/object-store/src/metrics.rs deleted file mode 100644 index 28aee460f6c..00000000000 --- a/src/object-store/src/metrics.rs +++ /dev/null @@ -1,15 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! object-store metrics diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 772bff4aad9..2e3c88af463 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -2970,7 +2970,7 @@ mod test { use datafusion::functions_aggregate::expr_fn::count; use datafusion::logical_expr::builder::LogicalTableSource; use datafusion::logical_expr::{LogicalPlanBuilder, col}; - use session::context::{QueryContext, QueryContextBuilder}; + use session::context::QueryContext; use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; @@ -3967,54 +3967,4 @@ WITH ('repartition.column.hint' = 'host')", .contains("cannot set repartition.column.hint on a table with partition metadata") ); } - - #[tokio::test] - #[ignore = "TODO(ruihang): WIP new partition rule"] - async fn test_parse_partitions() { - common_telemetry::init_default_ut_logging(); - let cases = [ - ( - r" -CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) ) -PARTITION ON COLUMNS (b) ( - b < 'hz', - b >= 'hz' AND b < 'sh', - b >= 'sh' -) -ENGINE=mito", - r#"[{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"hz\"}}"]},{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"sh\"}}"]},{"column_list":["b"],"value_list":["\"MaxValue\""]}]"#, - ), - ( - r" -CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) ) -PARTITION BY RANGE COLUMNS (b, a) ( - PARTITION r0 VALUES LESS THAN ('hz', 10), - b < 'hz' AND a < 10, - b >= 'hz' AND b < 'sh' AND a >= 10 AND a < 20, - b >= 'sh' AND a >= 20 -) -ENGINE=mito", - r#"[{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"hz\"}}","{\"Value\":{\"Int32\":10}}"]},{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"sh\"}}","{\"Value\":{\"Int32\":20}}"]},{"column_list":["b","a"],"value_list":["\"MaxValue\"","\"MaxValue\""]}]"#, - ), - ]; - let ctx = QueryContextBuilder::default().build().into(); - for (sql, expected) in cases { - let result = ParserContext::create_with_dialect( - sql, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .unwrap(); - match &result[0] { - Statement::CreateTable(c) => { - let expr = expr_helper::create_to_expr(c, &QueryContext::arc()).unwrap(); - let (partitions, _) = - parse_partitions(&expr, c.partitions.clone(), &ctx).unwrap(); - let json = serde_json::to_string(&partitions).unwrap(); - assert_eq!(json, expected); - } - _ => unreachable!(), - } - } - } } diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 9a355f4c8dc..5a9cb3e9aa6 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -60,7 +60,6 @@ yaml-rust = "0.4" [dev-dependencies] catalog = { workspace = true, features = ["testing"] } criterion = { workspace = true, features = ["html_reports"] } -session = { workspace = true, features = ["testing"] } [[bench]] name = "processor" diff --git a/src/pipeline/src/etl.rs b/src/pipeline/src/etl.rs index f83b5c2626c..f78a515fb59 100644 --- a/src/pipeline/src/etl.rs +++ b/src/pipeline/src/etl.rs @@ -278,24 +278,6 @@ impl PipelineExecOutput { } } - // Backward compatibility helper that returns first ContextOpt with all its rows - // or merges all rows with default ContextOpt for multi-context scenarios - pub fn into_legacy_format(self) -> Option<(ContextOpt, Vec)> { - if let Self::Transformed(TransformedOutput { rows_by_context }) = self { - if rows_by_context.len() == 1 { - let (opt, rows) = rows_by_context.into_iter().next().unwrap(); - Some((opt, rows)) - } else { - // Multiple contexts: merge all rows with default ContextOpt for test compatibility - let all_rows: Vec = - rows_by_context.into_values().flatten().collect(); - Some((ContextOpt::default(), all_rows)) - } - } else { - None - } - } - // Note: This is a test only function, do not use it in production. pub fn into_dispatched(self) -> Option { if let Self::DispatchedTo(d, _) = self { diff --git a/src/puffin/Cargo.toml b/src/puffin/Cargo.toml index 078e3aebeb6..a51ca99d444 100644 --- a/src/puffin/Cargo.toml +++ b/src/puffin/Cargo.toml @@ -34,6 +34,3 @@ snafu.workspace = true tokio.workspace = true tokio-util.workspace = true uuid.workspace = true - -[dev-dependencies] -common-base = { workspace = true, features = ["testing"] } diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index c91bb477c78..1d95f7856d5 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -91,7 +91,6 @@ nalgebra.workspace = true paste.workspace = true pretty_assertions.workspace = true rand.workspace = true -session = { workspace = true, features = ["testing"] } store-api.workspace = true table = { workspace = true, features = ["testing"] } toml.workspace = true diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 06648254a06..c1af56a6281 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -165,7 +165,6 @@ permutation = "0.4" rand.workspace = true serde_json.workspace = true servers = { workspace = true, features = ["testing"] } -session = { workspace = true, features = ["testing"] } simdutf8 = "0.1" table.workspace = true tempfile = "3.0.0" diff --git a/src/servers/src/http/pprof/nix.rs b/src/servers/src/http/pprof/nix.rs deleted file mode 100644 index 12a518e1de0..00000000000 --- a/src/servers/src/http/pprof/nix.rs +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::any::Any; -use std::time::Duration; - -use common_error::ext::ErrorExt; -use common_error::status_code::StatusCode; -use common_macro::stack_trace_debug; -use prost::Message; -use snafu::{Location, ResultExt, Snafu}; - -#[derive(Snafu)] -#[stack_trace_debug] -pub enum Error { - #[snafu(display("Failed to create profiler guard"))] - CreateGuard { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create report"))] - CreateReport { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create flamegraph"))] - CreateFlamegraph { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, - - #[snafu(display("Failed to create pprof report"))] - ReportPprof { - #[snafu(source)] - error: pprof::Error, - #[snafu(implicit)] - location: Location, - }, -} - -pub type Result = std::result::Result; - -impl ErrorExt for Error { - fn status_code(&self) -> StatusCode { - StatusCode::Unexpected - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -/// CPU profiler utility. -// Inspired by https://github.com/datafuselabs/databend/blob/67f445e83cd4eceda98f6c1c114858929d564029/src/common/base/src/base/profiling.rs -#[derive(Debug)] -pub struct Profiling { - /// Sample duration. - duration: Duration, - /// Sample frequency. - frequency: i32, -} - -impl Profiling { - /// Creates a new profiler. - pub fn new(duration: Duration, frequency: i32) -> Profiling { - Profiling { - duration, - frequency, - } - } - - /// Profiles and returns a generated pprof report. - pub async fn report(&self) -> Result { - let guard = pprof::ProfilerGuardBuilder::default() - .frequency(self.frequency) - .blocklist(&["libc", "libgcc", "pthread", "vdso"]) - .build() - .context(CreateGuardSnafu)?; - tokio::time::sleep(self.duration).await; - guard.report().build().context(CreateReportSnafu) - } - - /// Profiles and returns a generated flamegraph. - pub async fn dump_flamegraph(&self) -> Result> { - let mut body: Vec = Vec::new(); - - let report = self.report().await?; - report - .flamegraph(&mut body) - .context(CreateFlamegraphSnafu)?; - - Ok(body) - } - - /// Profiles and returns a generated proto. - pub async fn dump_proto(&self) -> Result> { - let report = self.report().await?; - // Generate google’s pprof format report. - let profile = report.pprof().context(ReportPprofSnafu)?; - let body = profile.encode_to_vec(); - - Ok(body) - } -} diff --git a/src/servers/src/opentsdb/codec.rs b/src/servers/src/opentsdb/codec.rs index c4760aa74de..55674eebf7a 100644 --- a/src/servers/src/opentsdb/codec.rs +++ b/src/servers/src/opentsdb/codec.rs @@ -12,11 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::{Column, ColumnDataType, InsertRequest as GrpcInsertRequest, SemanticType, column}; -use common_query::prelude::{greptime_timestamp, greptime_value}; - -use crate::error::{self, Result}; - #[derive(Debug, Clone)] pub struct DataPoint { metric: String, @@ -35,76 +30,6 @@ impl DataPoint { } } - pub fn try_create(line: &str) -> Result { - let tokens = line.split_whitespace().collect::>(); - let cmd = if tokens.is_empty() { "" } else { tokens[0] }; - // OpenTSDB command is case sensitive, verified in real OpenTSDB. - if cmd != "put" { - return error::InvalidQuerySnafu { - reason: format!("unknown command {cmd}."), - } - .fail(); - } - if tokens.len() < 4 { - return error::InvalidQuerySnafu { - reason: format!( - "put: illegal argument: not enough arguments (need least 4, got {})", - tokens.len() - ), - } - .fail(); - } - - let metric = tokens[1]; - - let ts_millis = match tokens[2].parse::() { - Ok(t) => Self::timestamp_to_millis(t), - Err(_) => { - return error::InvalidQuerySnafu { - reason: format!("put: invalid timestamp: {}", tokens[2]), - } - .fail(); - } - }; - - let value = match tokens[3].parse::() { - Ok(v) => v, - Err(_) => { - return error::InvalidQuerySnafu { - reason: format!("put: invalid value: {}", tokens[3]), - } - .fail(); - } - }; - - let mut tags = Vec::with_capacity(tokens.len() - 4); - for token in tokens.iter().skip(4) { - let tag = token.split('=').collect::>(); - if tag.len() != 2 || tag[0].is_empty() || tag[1].is_empty() { - return error::InvalidQuerySnafu { - reason: format!("put: invalid tag: {token}"), - } - .fail(); - } - let tagk = tag[0].to_string(); - let tagv = tag[1].to_string(); - if tags.iter().any(|(t, _)| t == &tagk) { - return error::InvalidQuerySnafu { - reason: format!("put: illegal argument: duplicate tag: {tagk}"), - } - .fail(); - } - tags.push((tagk, tagv)); - } - - Ok(DataPoint { - metric: metric.to_string(), - ts_millis, - value, - tags, - }) - } - pub fn metric(&self) -> &str { &self.metric } @@ -125,53 +50,6 @@ impl DataPoint { self.value } - pub fn as_grpc_insert(&self) -> GrpcInsertRequest { - let mut columns = Vec::with_capacity(2 + self.tags.len()); - - let ts_column = Column { - column_name: greptime_timestamp().to_string(), - values: Some(column::Values { - timestamp_millisecond_values: vec![self.ts_millis], - ..Default::default() - }), - semantic_type: SemanticType::Timestamp as i32, - datatype: ColumnDataType::TimestampMillisecond as i32, - ..Default::default() - }; - columns.push(ts_column); - - let field_column = Column { - column_name: greptime_value().to_string(), - values: Some(column::Values { - f64_values: vec![self.value], - ..Default::default() - }), - semantic_type: SemanticType::Field as i32, - datatype: ColumnDataType::Float64 as i32, - ..Default::default() - }; - columns.push(field_column); - - for (tagk, tagv) in self.tags.iter() { - columns.push(Column { - column_name: tagk.clone(), - values: Some(column::Values { - string_values: vec![tagv.clone()], - ..Default::default() - }), - semantic_type: SemanticType::Tag as i32, - datatype: ColumnDataType::String as i32, - ..Default::default() - }); - } - - GrpcInsertRequest { - table_name: self.metric.clone(), - columns, - row_count: 1, - } - } - pub fn timestamp_to_millis(t: i64) -> i64 { // 9999999999999 (13 digits) is of date "Sat Nov 20 2286 17:46:39 UTC", // 999999999999 (12 digits) is "Sun Sep 09 2001 01:46:39 UTC", @@ -184,112 +62,3 @@ impl DataPoint { } } } - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_try_create() { - fn test_illegal_line(line: &str, expected_err: &str) { - let result = DataPoint::try_create(line); - match result.unwrap_err() { - error::Error::InvalidQuery { reason, .. } => { - assert_eq!(reason, expected_err) - } - _ => unreachable!(), - } - } - - test_illegal_line("no_put", "unknown command no_put."); - test_illegal_line( - "put", - "put: illegal argument: not enough arguments (need least 4, got 1)", - ); - test_illegal_line( - "put metric.foo notatime 42 host=web01", - "put: invalid timestamp: notatime", - ); - test_illegal_line( - "put metric.foo 1000 notavalue host=web01", - "put: invalid value: notavalue", - ); - test_illegal_line("put metric.foo 1000 42 host=", "put: invalid tag: host="); - test_illegal_line( - "put metric.foo 1000 42 host=web01 host=web02", - "put: illegal argument: duplicate tag: host", - ); - - let data_point = DataPoint::try_create( - "put sys.if.bytes.out 1479496100 1.3E3 host=web01 interface=eth0", - ) - .unwrap(); - assert_eq!(data_point.metric, "sys.if.bytes.out"); - assert_eq!(data_point.ts_millis, 1479496100000); - assert_eq!(data_point.value, 1.3e3); - assert_eq!( - data_point.tags, - vec![ - ("host".to_string(), "web01".to_string()), - ("interface".to_string(), "eth0".to_string()) - ] - ); - - let data_point = - DataPoint::try_create("put sys.procs.running 1479496100 42 host=web01").unwrap(); - assert_eq!(data_point.metric, "sys.procs.running"); - assert_eq!(data_point.ts_millis, 1479496100000); - assert_eq!(data_point.value, 42f64); - assert_eq!( - data_point.tags, - vec![("host".to_string(), "web01".to_string())] - ); - } - - #[test] - fn test_as_grpc_insert() { - let data_point = DataPoint { - metric: "my_metric_1".to_string(), - ts_millis: 1000, - value: 1.0, - tags: vec![ - ("tagk1".to_string(), "tagv1".to_string()), - ("tagk2".to_string(), "tagv2".to_string()), - ], - }; - - let grpc_insert = data_point.as_grpc_insert(); - assert_eq!(grpc_insert.table_name, "my_metric_1"); - - let columns = &grpc_insert.columns; - let row_count = grpc_insert.row_count; - - assert_eq!(row_count, 1); - assert_eq!(columns.len(), 4); - - assert_eq!(columns[0].column_name, greptime_timestamp()); - assert_eq!( - columns[0] - .values - .as_ref() - .unwrap() - .timestamp_millisecond_values, - vec![1000] - ); - - assert_eq!(columns[1].column_name, greptime_value()); - assert_eq!(columns[1].values.as_ref().unwrap().f64_values, vec![1.0]); - - assert_eq!(columns[2].column_name, "tagk1"); - assert_eq!( - columns[2].values.as_ref().unwrap().string_values, - vec!["tagv1"] - ); - - assert_eq!(columns[3].column_name, "tagk2"); - assert_eq!( - columns[3].values.as_ref().unwrap().string_values, - vec!["tagv2"] - ); - } -} diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index c83debd39a1..25cb94662c0 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -489,45 +489,6 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { } } -#[allow(dead_code)] -pub(super) fn type_pg_to_gt(origin: &Type) -> Result { - // Note that we only support a small amount of pg data types - match origin { - &Type::BOOL => Ok(ConcreteDataType::boolean_datatype()), - &Type::INT2 => Ok(ConcreteDataType::int16_datatype()), - &Type::INT4 => Ok(ConcreteDataType::int32_datatype()), - &Type::INT8 => Ok(ConcreteDataType::int64_datatype()), - &Type::NUMERIC => Ok(ConcreteDataType::uint64_datatype()), - &Type::VARCHAR | &Type::CHAR | &Type::TEXT => Ok(ConcreteDataType::string_datatype()), - &Type::TIMESTAMP | &Type::TIMESTAMPTZ => Ok(ConcreteDataType::timestamp_datatype( - common_time::timestamp::TimeUnit::Millisecond, - )), - &Type::DATE => Ok(ConcreteDataType::date_datatype()), - &Type::TIME => Ok(ConcreteDataType::timestamp_datatype( - common_time::timestamp::TimeUnit::Microsecond, - )), - &Type::INT2_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int16_datatype(), - ))), - &Type::INT4_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int32_datatype(), - ))), - &Type::INT8_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::int64_datatype(), - ))), - &Type::NUMERIC_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new( - ConcreteDataType::uint64_datatype(), - ))), - &Type::VARCHAR_ARRAY | &Type::CHAR_ARRAY | &Type::TEXT_ARRAY => Ok( - ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())), - ), - _ => server_error::InternalSnafu { - err_msg: format!("unimplemented datatype {origin:?}"), - } - .fail(), - } -} - pub(super) fn invalid_parameter_error(msg: &str, detail: Option) -> PgWireError { let mut error_info = PgErrorCode::Ec22023.to_err_info(msg.to_string()); error_info.detail = detail; diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index e74088a2a5d..12a56076d6f 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -81,8 +81,6 @@ pub struct MetricsIngestOutcome { #[async_trait] pub trait InfluxdbLineProtocolHandler { - /// A successful request will not return a response. - /// Only on error will the socket return a line of data. async fn exec(&self, request: InfluxdbRequest, ctx: QueryContextRef) -> Result; } @@ -91,12 +89,10 @@ pub trait OpentsdbProtocolHandler { /// Checks all points in one external request before per-point debug execution. async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> Result<()>; - /// A successful request will not return a response. - /// Only on error will the socket return a line of data. async fn exec(&self, data_points: Vec, ctx: QueryContextRef) -> Result; - /// Executes an ordinary HTTP put with optional batching. Debug and socket - /// callers retain [`Self::exec`]; the frontend clears HTTP batching selection + /// Executes an ordinary HTTP put with optional batching. Debug callers + /// retain [`Self::exec`]; the frontend clears HTTP batching selection /// there so diagnostic requests preserve direct, per-point error attribution. async fn exec_batch(&self, data_points: Vec, ctx: QueryContextRef) -> Result { self.exec(data_points, ctx).await diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 01ef368f1f0..5485198c266 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -4,9 +4,6 @@ version.workspace = true edition.workspace = true license.workspace = true -[features] -testing = [] - [lints] workspace = true diff --git a/src/session/src/context.rs b/src/session/src/context.rs index ec10dc6cf3a..c1448053f35 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -64,7 +64,6 @@ pub struct QueryContext { current_catalog: String, /// mapping of RegionId to SequenceNumber, for snapshot read, meaning that the read should only /// container data that was committed before(and include) the given sequence number - /// this field will only be filled if extensions contains a pair of "snapshot_read" and "true" snapshot_seqs: Arc>>, /// Mappings of the RegionId to the minimal sequence of SST file to scan. sst_min_sequences: Arc>>, diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index cb048a94b9e..2cb78431969 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -25,7 +25,7 @@ use snafu::ResultExt; use sqlparser::ast::{ColumnOptionDef, DataType, Expr}; use sqlparser_derive::{Visit, VisitMut}; -use crate::ast::{ColumnDef, Ident, ObjectName, Value as SqlValue}; +use crate::ast::{ColumnDef, Ident, ObjectName}; use crate::error::{ InvalidFlowQuerySnafu, InvalidSqlSnafu, Result, SetFulltextOptionSnafu, SetSkippingIndexOptionSnafu, @@ -356,23 +356,6 @@ impl Partitions { } } -#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut)] -pub struct PartitionEntry { - pub name: Ident, - pub value_list: Vec, -} - -impl Display for PartitionEntry { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!( - f, - "PARTITION {} VALUES LESS THAN ({})", - self.name, - format_list_comma!(self.value_list), - ) - } -} - impl Display for Partitions { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if !self.column_list.is_empty() { diff --git a/src/store-api/src/manifest/action.rs b/src/store-api/src/manifest/action.rs deleted file mode 100644 index c9b0d3e7d0c..00000000000 --- a/src/store-api/src/manifest/action.rs +++ /dev/null @@ -1,114 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! Common actions for manifest -use serde::{Deserialize, Serialize}; - -use crate::manifest::ManifestVersion; - -pub type ProtocolVersion = u16; - -/// Current reader and writer versions -/// TODO(dennis): configurable -const READER_VERSION: ProtocolVersion = 0; -const WRITER_VERSION: ProtocolVersion = 0; - -/// The maximum protocol version we are currently allowed to use, -/// TODO(dennis): reading from configuration. -pub fn supported_protocol_version() -> (ProtocolVersion, ProtocolVersion) { - (READER_VERSION, WRITER_VERSION) -} - -/// Protocol action that used to block older clients from reading or writing the log when backwards -/// incompatible changes are made to the protocol. -/// -/// clients should be tolerant of messages and fields that they do not understand. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct ProtocolAction { - pub min_reader_version: ProtocolVersion, - pub min_writer_version: ProtocolVersion, -} - -impl std::fmt::Display for ProtocolAction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "Protocol({}, {})", - &self.min_reader_version, &self.min_writer_version, - ) - } -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] -pub struct VersionHeader { - pub prev_version: ManifestVersion, -} - -impl Default for ProtocolAction { - fn default() -> Self { - let (min_reader_version, min_writer_version) = supported_protocol_version(); - Self { - min_reader_version, - min_writer_version, - } - } -} - -impl ProtocolAction { - pub fn new() -> Self { - Self::default() - } - - pub fn is_readable(&self, reader_version: ProtocolVersion) -> bool { - reader_version >= self.min_reader_version - } - - pub fn is_writable(&self, writer_version: ProtocolVersion) -> bool { - writer_version >= self.min_writer_version - } -} - -#[cfg(test)] -mod tests { - use serde_json as json; - - use super::*; - - #[test] - fn test_protocol_action() { - let mut action = ProtocolAction::new(); - - assert!(action.is_readable(0)); - assert!(action.is_writable(0)); - action.min_reader_version = 2; - action.min_writer_version = 3; - assert!(!action.is_readable(0)); - assert!(!action.is_writable(0)); - assert!(action.is_readable(2)); - assert!(action.is_writable(3)); - assert!(action.is_readable(3)); - assert!(action.is_writable(4)); - - let s = json::to_string(&action).unwrap(); - assert_eq!("{\"min_reader_version\":2,\"min_writer_version\":3}", s); - - let action_decoded: ProtocolAction = json::from_str(&s).unwrap(); - assert!(!action_decoded.is_readable(0)); - assert!(!action_decoded.is_writable(0)); - assert!(action_decoded.is_readable(2)); - assert!(action_decoded.is_writable(3)); - assert!(action_decoded.is_readable(3)); - assert!(action_decoded.is_writable(4)); - } -} diff --git a/src/store-api/src/manifest/storage.rs b/src/store-api/src/manifest/storage.rs deleted file mode 100644 index 58a3d27fe83..00000000000 --- a/src/store-api/src/manifest/storage.rs +++ /dev/null @@ -1,78 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use async_trait::async_trait; -use common_error::ext::ErrorExt; - -use crate::manifest::ManifestVersion; - -#[async_trait] -pub trait LogIterator: Send + Sync { - type Error: ErrorExt + Send + Sync; - - async fn next_log(&mut self) -> Result)>, Self::Error>; -} - -#[async_trait] -pub trait ManifestLogStorage { - type Error: ErrorExt + Send + Sync; - type Iter: LogIterator; - - /// Scan the logs in [start, end) - async fn scan( - &self, - start: ManifestVersion, - end: ManifestVersion, - ) -> Result; - - /// Delete logs and checkpoints which version is less than specified `end`. - /// Keep the last checkpoint and remaining logs if `keep_last_checkpoint` is true. - /// - /// Returns the delete logs number. - async fn delete_until( - &self, - end: ManifestVersion, - keep_last_checkpoint: bool, - ) -> Result; - - /// Delete all logs and checkpoints, and remove the manifest directory. - /// The delta file corresponding to the `remove_action_version` will be deleted along with the manifest directory at the end. - async fn delete_all(&self, remove_action_version: ManifestVersion) -> Result<(), Self::Error>; - - /// Save a log - async fn save(&self, version: ManifestVersion, bytes: &[u8]) -> Result<(), Self::Error>; - - /// Delete logs in [start, end) and ignore checkpoints. - async fn delete(&self, start: ManifestVersion, end: ManifestVersion) - -> Result<(), Self::Error>; - - /// Save a checkpoint. - async fn save_checkpoint( - &self, - version: ManifestVersion, - bytes: &[u8], - ) -> Result<(), Self::Error>; - - /// Load the latest checkpoint - async fn load_last_checkpoint(&self) - -> Result)>, Self::Error>; - /// Delete the checkpoint by version - async fn delete_checkpoint(&self, version: ManifestVersion) -> Result<(), Self::Error>; - - /// Load the checkpoint by version - async fn load_checkpoint( - &self, - version: ManifestVersion, - ) -> Result)>, Self::Error>; -} diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index bc9b31ada95..07cfe6a6516 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -29,8 +29,6 @@ pub const MERGE_MODE_KEY: &str = "merge_mode"; pub const TTL_KEY: &str = "ttl"; /// Option key for the per-table auto flush interval. pub const AUTO_FLUSH_INTERVAL_KEY: &str = "auto_flush_interval"; -/// Option key for snapshot read. -pub const SNAPSHOT_READ: &str = "snapshot_read"; /// Option key for compaction type. pub const COMPACTION_TYPE: &str = "compaction.type"; /// Option key for forcing compaction options override. diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 762bf80c4b3..e1f0eb22497 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -67,7 +67,7 @@ use crate::mito_engine_options::{ TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY, }; use crate::path_utils::table_dir; -use crate::storage::{ColumnId, RegionId, ScanRequest}; +use crate::storage::{ColumnId, RegionId}; /// The type of path to generate. #[derive(Debug, Clone, Copy, PartialEq, TryFromPrimitive)] @@ -540,11 +540,6 @@ pub struct RegionPutRequest { pub partition_expr_version: Option, } -#[derive(Debug)] -pub struct RegionReadRequest { - pub request: ScanRequest, -} - /// Request to delete data from a region. #[derive(Debug)] pub struct RegionDeleteRequest { diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 899438edf52..81f396ec703 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -49,7 +49,6 @@ use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use crate::error::{ConflictingTableOptionsSnafu, ParseTableOptionSnafu, Result}; use crate::metadata::{TableId, TableVersion}; -use crate::table_reference::TableReference; mod semantic; pub use semantic::*; @@ -835,25 +834,6 @@ impl Default for CompactTableRequest { } } -/// Truncate table request -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TruncateTableRequest { - pub catalog_name: String, - pub schema_name: String, - pub table_name: String, - pub table_id: TableId, -} - -impl TruncateTableRequest { - pub fn table_ref(&self) -> TableReference<'_> { - TableReference { - catalog: &self.catalog_name, - schema: &self.schema_name, - table: &self.table_name, - } - } -} - #[derive(Debug, Clone, Default, Deserialize, Serialize)] pub struct CopyDatabaseRequest { pub catalog_name: String, diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index d281f3cecc8..2c90f4ee2ae 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -116,7 +116,6 @@ rstest.workspace = true rstest_reuse.workspace = true sea-query.workspace = true servers = { workspace = true, features = ["testing"] } -session = { workspace = true, features = ["testing"] } similar-asserts.workspace = true sqlx = { workspace = true, features = [ "mysql", From e2b7e5e0a8c4e95911cf8d62560461d88f27ebcd Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 28 Sep 2026 12:40:35 +0000 Subject: [PATCH 064/133] chore: disable unstable rustfmt features (#9379) * chore: disable unstable rustfmt features * typo fix --- rustfmt.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/rustfmt.toml b/rustfmt.toml index 64d94def266..fccaefbc69a 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,2 +1,3 @@ -group_imports = "StdExternalCrate" -imports_granularity = "Module" +# disable these features until they become stable +#group_imports = "StdExternalCrate" +#imports_granularity = "Module" From 3c0e2a8d553c4e7d1dbf5fe16227f8f991dca6ed Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 28 Sep 2026 15:22:02 +0000 Subject: [PATCH 065/133] feat: export Metric snapshots with packed Parquet objects (#9382) * feat: export Metric snapshots with packed Parquet objects Signed-off-by: jeremyhi * fix: reject empty packed export time ranges Signed-off-by: jeremyhi * test: cover cancellation during packed export I/O Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 2 + src/cli/Cargo.toml | 1 + src/cli/src/data/export_v2.rs | 4 +- src/cli/src/data/export_v2/command.rs | 69 +- src/cli/src/data/export_v2/coordinator.rs | 190 ++++- src/cli/src/data/export_v2/data.rs | 4 + src/cli/src/data/export_v2/manifest.rs | 2 +- src/cli/src/data/snapshot_storage.rs | 102 ++- src/cli/src/database.rs | 102 +-- src/common/datasource/Cargo.toml | 1 + src/common/datasource/src/lib.rs | 1 + src/common/datasource/src/packed_writer.rs | 669 ++++++++++++++++++ src/common/datasource/src/parquet_writer.rs | 116 ++- src/operator/src/statement/copy_table_to.rs | 54 +- src/operator/src/statement/database_copy.rs | 2 +- src/operator/src/statement/export_database.rs | 70 +- .../src/statement/export_logical_tables.rs | 30 +- .../statement/export_logical_tables/tests.rs | 39 +- .../export_logical_tables/writers.rs | 9 +- src/servers/src/http.rs | 6 +- .../tests/export_logical_tables.rs | 196 ++++- .../cases/snapshot_parquet_copy/setup.sql | 6 +- .../cases/snapshot_parquet_copy/verify.result | 33 +- .../cases/snapshot_parquet_copy/verify.sql | 3 + 24 files changed, 1584 insertions(+), 127 deletions(-) create mode 100644 src/common/datasource/src/packed_writer.rs diff --git a/Cargo.lock b/Cargo.lock index 4ced42fc860..96903faaa25 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2088,6 +2088,7 @@ dependencies = [ "table", "tempfile", "tokio", + "tokio-util", "url", "uuid", ] @@ -2397,6 +2398,7 @@ dependencies = [ "paste", "regex", "serde", + "serde_json", "snafu 0.8.6", "strum 0.27.1", "tokio", diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 8178fa1b0b5..6abe504a3ab 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -71,3 +71,4 @@ common-test-util.workspace = true object-store = { workspace = true, features = ["testing"] } serde.workspace = true tempfile.workspace = true +tokio-util.workspace = true diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index f9a3c2be095..9887e525030 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -41,11 +41,13 @@ //! `--experimental-metric-export` enables shared Metric physical scans for Parquet. //! Enable `experimental_metric_export = true` on the frontend or standalone server //! and use an endpoint whose frontends all support and enable this option. +//! Add `--metric-data-layout packed` to write version-2 packed snapshots. +//! This requires the server packed-export capability; schema-only exports remain version 1. //! The command checks the server capability before modifying the snapshot. //! Each snapshot path belongs to one export task. Resume requires stable source //! schemas/data and confirmation that the previous export and storage writes ended; //! an HTTP timeout does not establish that. Only unfinished chunks are cleaned and -//! rerun. Completed chunks and the V2 snapshot/import format remain unchanged. +//! rerun. Resume retains the snapshot version/layout and completed chunks. mod chunker; mod command; diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index e0f67eb796c..20e69ef5a99 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -296,6 +296,10 @@ pub struct ExportCreateCommand { #[clap(long)] experimental_metric_export: bool, + /// Store independent Metric Parquet streams in shared packed objects. + #[clap(long, value_parser = ["packed"], requires = "experimental_metric_export")] + metric_data_layout: Option, + /// Delete existing snapshot and recreate. #[clap(long)] force: bool, @@ -345,6 +349,16 @@ impl ExportCreateCommand { let time_range = TimeRange::parse(self.start_time.as_deref(), self.end_time.as_deref()) .map_err(BoxedError::new)?; + if self.metric_data_layout.is_some() + && time_range.is_bounded() + && time_range.start == time_range.end + { + return crate::error::InvalidArgumentsSnafu { + msg: "Packed export requires --start-time to be earlier than --end-time", + } + .fail() + .map_err(BoxedError::new); + } if self.chunk_time_window.is_some() && !time_range.is_bounded() { return ChunkTimeWindowRequiresBoundsSnafu .fail() @@ -425,6 +439,13 @@ impl ExportCreateCommand { .map_err(BoxedError::new); } } + if self.metric_data_layout.is_some() && !self.schema_only { + database_client + .require_packed_export() + .await + .context(DatabaseSnafu) + .map_err(BoxedError::new)?; + } let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; Ok(Box::new(ExportCreate { @@ -434,6 +455,7 @@ impl ExportCreateCommand { schema_only: self.schema_only, format: self.format, experimental_metric_export: self.experimental_metric_export, + packed: self.metric_data_layout.is_some(), force: self.force, time_range, chunk_time_window: self.chunk_time_window, @@ -462,6 +484,7 @@ struct ExportConfig { schema_only: bool, format: DataFormat, experimental_metric_export: bool, + packed: bool, force: bool, time_range: TimeRange, chunk_time_window: Option, @@ -504,7 +527,7 @@ impl ExportCreate { let mut manifest = self.storage.read_manifest().await?; // Check version compatibility - if manifest.version != MANIFEST_VERSION || manifest.data_layout.is_some() { + if manifest.validate_layout().is_err() { return ManifestVersionMismatchSnafu { expected: MANIFEST_VERSION, found: manifest.version, @@ -513,6 +536,15 @@ impl ExportCreate { } validate_resume_config(&manifest, &self.config)?; + if manifest.is_packed() { + self.database_client + .require_packed_export() + .await + .context(DatabaseSnafu)?; + if !self.config.experimental_metric_export { + return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail(); + } + } info!( "Resuming existing snapshot: {} (completed: {}/{} chunks)", @@ -571,10 +603,14 @@ impl ExportCreate { self.config.chunk_time_window, )?; + if self.config.packed && !self.config.schema_only { + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + } if self.config.experimental_metric_export { for chunk in &manifest.chunks { self.storage - .prepare_export_chunk(&schema_names, chunk.id, false) + .prepare_export_chunk(&schema_names, chunk.id, false, manifest.is_packed()) .await?; } } @@ -796,6 +832,14 @@ fn build_schema_ddl( } fn validate_resume_config(manifest: &Manifest, config: &ExportConfig) -> Result<()> { + if config.packed && !manifest.schema_only && !manifest.is_packed() { + return ResumeConfigMismatchSnafu { + field: "metric_data_layout", + existing: "standalone".to_string(), + requested: "packed".to_string(), + } + .fail(); + } if manifest.schema_only != config.schema_only { return SchemaOnlyModeMismatchSnafu { existing_schema_only: manifest.schema_only, @@ -1801,6 +1845,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1820,7 +1865,7 @@ mod tests { #[test] fn test_validate_resume_config_accepts_schema_selection_with_different_case_and_order() { - let manifest = Manifest::new_for_export( + let mut manifest = Manifest::new_for_export( "greptime".to_string(), vec!["public".to_string(), "analytics".to_string()], false, @@ -1829,7 +1874,7 @@ mod tests { None, ) .unwrap(); - let config = ExportConfig { + let mut config = ExportConfig { catalog: "greptime".to_string(), schemas: Some(vec![ "ANALYTICS".to_string(), @@ -1839,6 +1884,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1850,6 +1896,18 @@ mod tests { }; assert!(validate_resume_config(&manifest, &config).is_ok()); + config.packed = true; + assert!( + validate_resume_config(&manifest, &config) + .unwrap_err() + .to_string() + .contains("metric_data_layout") + ); + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + assert!(validate_resume_config(&manifest, &config).is_ok()); + config.packed = false; + assert!(validate_resume_config(&manifest, &config).is_ok()); } #[test] @@ -1872,6 +1930,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range, chunk_time_window: Some(Duration::from_secs(3600)), @@ -1906,6 +1965,7 @@ mod tests { schema_only: false, format: DataFormat::Csv, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1942,6 +2002,7 @@ mod tests { schema_only: false, format: DataFormat::Parquet, experimental_metric_export: false, + packed: false, force: false, time_range: TimeRange::new(Some(start), Some(start)), chunk_time_window: None, diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index 5f95c43603e..c6181e51211 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -40,6 +40,7 @@ struct ExportContext<'a> { format: DataFormat, parallelism: usize, experimental_metric_export: bool, + packed: bool, resume: bool, } @@ -71,6 +72,7 @@ pub async fn export_data( catalog: manifest.catalog.clone(), schemas: manifest.schemas.clone(), format: manifest.format, + packed: manifest.is_packed(), parallelism: options.parallelism, experimental_metric_export: options.experimental_metric_export, resume: options.resume, @@ -279,7 +281,7 @@ async fn export_chunk( if context.experimental_metric_export { context .storage - .prepare_export_chunk(&context.schemas, chunk_id, context.resume) + .prepare_export_chunk(&context.schemas, chunk_id, context.resume, context.packed) .await?; } let scheme = StorageScheme::from_uri(context.snapshot_uri)?; @@ -289,6 +291,7 @@ async fn export_chunk( time_range, parallelism: context.parallelism, experimental_metric_export: context.experimental_metric_export, + packed: context.packed, }; for schema in &context.schemas { @@ -314,6 +317,39 @@ async fn export_chunk( } let files = list_chunk_files(context.storage, &context.schemas, chunk_id).await?; + if context.packed { + use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; + let mut expected = Vec::new(); + for schema in &context.schemas { + let prefix = data_dir_for_schema_chunk(schema, chunk_id); + let index_path = format!("{prefix}{PACK_INDEX_FILE}"); + let text = context.storage.read_text(&index_path).await?; + let index: PackIndex = serde_json::from_str(&text).map_err(|e| { + crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: e.to_string(), + } + .build() + })?; + index.validate().map_err(|e| { + crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: e.to_string(), + } + .build() + })?; + expected.push(index_path); + expected.extend(index.objects.iter().map(|o| format!("{prefix}{}", o.path))); + } + expected.sort(); + if files != expected { + return crate::data::export_v2::error::InvalidUriSnafu { + uri: context.snapshot_uri, + reason: "chunk inventory differs from packed indexes", + } + .fail(); + } + } info!("Collected {} files for chunk {}", files.len(), chunk_id); Ok(files) } @@ -462,6 +498,158 @@ mod tests { } } + struct PausedClose { + inner: Option, + started: std::sync::Arc, + release: std::sync::Arc, + } + + impl object_store::layers::mock::oio::Write for PausedClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.inner.as_mut().unwrap().write(bytes).await + } + + async fn close(&mut self) -> object_store::Result { + let mut inner = self.inner.take().unwrap(); + let (started, release) = (self.started.clone(), self.release.clone()); + // Model blocking storage I/O that survives dropping its caller. + let (inner, result) = tokio::spawn(async move { + started.notify_one(); + release.notified().await; + let result = inner.close().await; + (inner, result) + }) + .await + .unwrap(); + self.inner = Some(inner); + result + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.as_mut().unwrap().abort().await + } + } + + #[tokio::test] + async fn packed_cancellation_drains_close_before_failing_chunk() { + use std::sync::Arc; + + use common_datasource::packed_snapshot::PACK_INDEX_FILE; + use common_datasource::packed_writer::{PackedTableWriter, PackedWriter}; + use common_datasource::parquet_writer::ParquetFileWriter; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; + use tokio_util::sync::CancellationToken; + + use crate::data::snapshot_storage::OpenDalStorage; + + for paused_path in ["pack-000000.bin", PACK_INDEX_FILE] { + let directory = tempfile::tempdir().unwrap(); + let uri = url::Url::from_directory_path(directory.path()).unwrap(); + let storage = + OpenDalStorage::from_uri(uri.as_str(), &ObjectStoreConfig::default()).unwrap(); + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let factory: MockWriterFactory = Arc::new({ + let (started, release) = (started.clone(), release.clone()); + move |path, _, inner| { + if path.trim_start_matches('/') == paused_path { + Box::new(PausedClose { + inner: Some(inner), + started: started.clone(), + release: release.clone(), + }) + } else { + inner + } + } + }); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let packed = PackedWriter::new(store.clone()).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + true, + )])); + let mut table = ParquetFileWriter::open_packed( + schema, + store.clone(), + "unused", + None, + PackedTableWriter::new(packed.clone(), "empty".into(), 1, false), + ) + .unwrap(); + table.finish(None).await.unwrap(); + let token = CancellationToken::new(); + let mut manifest = pending_manifest(1); + manifest.version = 2; + manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into()); + let export = export_data_concurrent( + &storage, + &mut manifest, + 2, + &crate::data::progress::NoopProgress, + |_, _| async { + let mut packed = packed.lock().await; + let result = packed.finish(&token).await; + if result.is_err() { + packed.abort().await.unwrap(); + } + result.map_err(|error| { + crate::data::export_v2::error::IoSnafu { + operation: "exporting packed chunk", + error: std::io::Error::other(error), + } + .build() + }) + }, + ); + tokio::pin!(export); + tokio::select! { + _ = started.notified() => {}, + result = &mut export => panic!("completed while close paused: {result:?}"), + } + token.cancel(); + assert!(futures::poll!(&mut export).is_pending()); + assert_eq!( + storage.read_manifest().await.unwrap().chunks[0].status, + ChunkStatus::InProgress + ); + release.notify_one(); + assert!(export.await.unwrap_err().to_string().contains("cancelled")); + let persisted = storage.read_manifest().await.unwrap(); + assert_eq!(persisted.chunks[0].status, ChunkStatus::Failed); + assert!(!persisted.is_complete()); + assert!(persisted.chunks[0].files.is_empty()); + assert!( + store + .stat("pack-000000.bin") + .await + .unwrap() + .content_length() + > 0 + ); + assert_eq!( + store.exists(PACK_INDEX_FILE).await.unwrap(), + paused_path == PACK_INDEX_FILE + ); + if paused_path == PACK_INDEX_FILE { + let index: common_datasource::packed_snapshot::PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()) + .unwrap(); + index.validate_membership(["empty"]).unwrap(); + } + } + } + #[test] fn test_next_eligible_chunk_scans_in_order() { let manifest = pending_manifest(3); diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index 8723154e4b1..53c3899f59e 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -31,6 +31,7 @@ pub(super) struct CopyOptions { pub(super) time_range: TimeRange, pub(super) parallelism: usize, pub(crate) experimental_metric_export: bool, + pub(super) packed: bool, } pub(super) struct CopyTarget { @@ -220,6 +221,9 @@ fn build_with_options(options: &CopyOptions) -> String { if options.experimental_metric_export { parts.push("experimental_metric_export='true'".to_string()); } + if options.packed { + parts.push("metric_data_layout='packed'".to_string()); + } if let Some(start) = options.time_range.start { parts.push(format!( "START_TIME='{}'", diff --git a/src/cli/src/data/export_v2/manifest.rs b/src/cli/src/data/export_v2/manifest.rs index 5dd9d6f1684..d56f2732490 100644 --- a/src/cli/src/data/export_v2/manifest.rs +++ b/src/cli/src/data/export_v2/manifest.rs @@ -28,7 +28,7 @@ use crate::data::export_v2::error::{ TimeParseInvalidFormatSnafu, }; -/// Manifest format version produced by the current exporter. +/// Default standalone manifest version. Packed data snapshots use version 2. pub const MANIFEST_VERSION: u32 = 1; /// Manifest file name within snapshot directory. diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index c5e9561d285..17008706066 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -295,8 +295,9 @@ pub trait SnapshotStorage: Send + Sync { schemas: &[String], chunk_id: u32, resume: bool, + packed: bool, ) -> Result<()> { - let _ = (schemas, chunk_id, resume); + let _ = (schemas, chunk_id, resume, packed); InvalidUriSnafu { uri: "snapshot", reason: "storage does not support preparing export chunks", @@ -730,6 +731,7 @@ impl SnapshotStorage for OpenDalStorage { schemas: &[String], chunk_id: u32, resume: bool, + packed: bool, ) -> Result<()> { let mut files = Vec::new(); for schema in schemas { @@ -751,10 +753,17 @@ impl SnapshotStorage for OpenDalStorage { continue; } let name = path.strip_prefix(&prefix).unwrap_or(""); - if !resume || entry.metadata().is_dir() || !valid_chunk_filename(name) { + if !resume + || entry.metadata().is_dir() + || !(if packed { + valid_packed_chunk_filename(name) + } else { + valid_chunk_filename(name) + }) + { return InvalidUriSnafu { uri: path, - reason: "expected an empty new chunk or direct Parquet files in an owned unfinished chunk", + reason: "expected an empty new chunk or recognized files in an owned unfinished chunk", }.fail(); } files.push(path.to_string()); @@ -786,6 +795,17 @@ impl SnapshotStorage for OpenDalStorage { } } +fn valid_packed_chunk_filename(name: &str) -> bool { + name == common_datasource::packed_snapshot::PACK_INDEX_FILE + || [("pack-", ".bin"), ("table-", ".parquet")] + .iter() + .any(|(prefix, suffix)| { + name.strip_prefix(prefix) + .and_then(|n| n.strip_suffix(suffix)) + .is_some_and(|id| !id.is_empty() && id.bytes().all(|b| b.is_ascii_digit())) + }) +} + fn valid_chunk_filename(name: &str) -> bool { name.strip_suffix(".parquet") .is_some_and(|stem| !stem.is_empty()) @@ -1074,7 +1094,7 @@ mod tests { let schemas = vec!["public".to_string(), "other".to_string()]; assert!( storage - .prepare_export_chunk(&schemas, 2, false) + .prepare_export_chunk(&schemas, 2, false, false) .await .is_err() ); @@ -1085,11 +1105,11 @@ mod tests { .unwrap() ); storage - .prepare_export_chunk(&schemas, 2, true) + .prepare_export_chunk(&schemas, 2, true, false) .await .unwrap(); storage - .prepare_export_chunk(&schemas, 2, true) + .prepare_export_chunk(&schemas, 2, true, false) .await .unwrap(); assert!( @@ -1129,7 +1149,7 @@ mod tests { } assert!( storage - .prepare_export_chunk(&["public".into(), "other".into()], 2, true) + .prepare_export_chunk(&["public".into(), "other".into()], 2, true, false) .await .is_err() ); @@ -1143,6 +1163,70 @@ mod tests { } } + #[tokio::test] + async fn packed_retry_preserves_foreign_files_and_completed_chunks() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + let owned = ["pack-000000.bin", "table-42.parquet", "pack-index.json"]; + for name in owned { + storage + .write_text(&format!("data/public/2/{name}"), "partial") + .await + .unwrap(); + } + for path in [ + "data/public/1/pack-000000.bin", + "data/public/2/foreign.parquet", + ] { + storage.write_text(path, "keep").await.unwrap(); + } + let schemas = ["public".into()]; + assert!( + storage + .prepare_export_chunk(&schemas, 2, true, true) + .await + .is_err() + ); + for name in owned { + assert!( + storage + .file_exists(&format!("data/public/2/{name}")) + .await + .unwrap() + ); + } + assert_eq!( + storage + .read_text("data/public/2/foreign.parquet") + .await + .unwrap(), + "keep" + ); + storage + .object_store + .delete("data/public/2/foreign.parquet") + .await + .unwrap(); + storage + .prepare_export_chunk(&schemas, 2, true, true) + .await + .unwrap(); + assert!( + storage + .list_files_recursive("data/public/2/") + .await + .unwrap() + .is_empty() + ); + assert_eq!( + storage + .read_text("data/public/1/pack-000000.bin") + .await + .unwrap(), + "keep" + ); + } + #[tokio::test] async fn test_prepare_export_chunk_reports_delete_failure() { let dir = tempdir().unwrap(); @@ -1159,7 +1243,7 @@ mod tests { .await .unwrap(); let error = storage - .prepare_export_chunk(&["public".into()], 2, true) + .prepare_export_chunk(&["public".into()], 2, true, false) .await .unwrap_err(); let Error::StorageOperation { @@ -1191,7 +1275,7 @@ mod tests { .await .unwrap(); storage - .prepare_export_chunk(&[schema], 2, true) + .prepare_export_chunk(&[schema], 2, true, false) .await .unwrap(); assert!(!storage.file_exists(&owned).await.unwrap()); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index b771ab48d05..ac314bc5899 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -99,6 +99,15 @@ impl DatabaseClient { /// Requires the explicit packed-import protocol before any restore mutation. pub async fn require_packed_import(&self) -> Result<()> { + self.require_packed_capability("metric_packed_import").await + } + + /// Requires packed export support before creating snapshot artifacts. + pub async fn require_packed_export(&self) -> Result<()> { + self.require_packed_capability("metric_packed_export").await + } + + async fn require_packed_capability(&self, capability: &str) -> Result<()> { let url = format!("http://{}/v1/capabilities", self.addr); let mut builder = reqwest::Client::builder().timeout(self.timeout); if let Some(proxy) = self.proxy.clone() { @@ -113,18 +122,18 @@ impl DatabaseClient { request = request.header("Authorization", auth); } let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { - reason: "packed import capability request failed", + reason: "packed snapshot capability request failed", })?; if response.status() == reqwest::StatusCode::NOT_FOUND { return crate::error::InvalidArgumentsSnafu { - msg: "target does not support packed import", + msg: format!("server does not support {capability}"), } .fail(); } let response = response .error_for_status() .with_context(|_| HttpQuerySqlSnafu { - reason: "packed import capability request rejected", + reason: "packed snapshot capability request rejected", })?; let body = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot read capability response", @@ -136,9 +145,9 @@ impl DatabaseClient { } .fail(); } - if value.get("metric_packed_import").and_then(Value::as_u64) != Some(1) { + if value.get(capability).and_then(Value::as_u64) != Some(1) { return crate::error::InvalidArgumentsSnafu { - msg: "target does not support packed import version 1", + msg: format!("server does not support {capability} version 1"), } .fail(); } @@ -221,41 +230,54 @@ mod tests { #[tokio::test] async fn packed_capability_probe_checks_auth_and_protocol_version() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; - for (status, body, supported) in [ - (200, r#"{"metric_packed_import":1,"future":2}"#, true), - (200, "{}", false), - (200, r#"{"metric_packed_import":2}"#, false), - (404, "{}", false), - (401, "{}", false), - (403, "{}", false), - (200, "[]", false), - (200, "invalid-json", false), - ] { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let n = socket.read(&mut request).await.unwrap(); - let request = String::from_utf8_lossy(&request[..n]).to_lowercase(); - assert!(request.starts_with("get /v1/capabilities ")); - assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza==")); - socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); - }); - let client = super::DatabaseClient::new( - address.to_string(), - "greptime".into(), - Some("user:password".into()), - std::time::Duration::from_secs(5), - None, - true, - ); - assert_eq!( - client.require_packed_import().await.is_ok(), - supported, - "{status}: {body}" - ); - server.await.unwrap(); + for export in [false, true] { + for (status, body, supported) in [ + (200, r#"{"metric_packed_import":1,"future":2}"#, true), + (200, "{}", false), + (200, r#"{"metric_packed_import":2}"#, false), + (404, "{}", false), + (401, "{}", false), + (403, "{}", false), + (200, "[]", false), + (200, "invalid-json", false), + ] { + let body = if export { + body.replace("metric_packed_import", "metric_packed_export") + } else { + body.to_string() + }; + let response_body = body.clone(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let body = response_body; + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let n = socket.read(&mut request).await.unwrap(); + let request = String::from_utf8_lossy(&request[..n]).to_lowercase(); + assert!(request.starts_with("get /v1/capabilities ")); + assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza==")); + socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + }); + let client = super::DatabaseClient::new( + address.to_string(), + "greptime".into(), + Some("user:password".into()), + std::time::Duration::from_secs(5), + None, + true, + ); + assert_eq!( + if export { + client.require_packed_export().await.is_ok() + } else { + client.require_packed_import().await.is_ok() + }, + supported, + "{status}: {body}" + ); + server.await.unwrap(); + } } } diff --git a/src/common/datasource/Cargo.toml b/src/common/datasource/Cargo.toml index d1ca609c606..3953d144bad 100644 --- a/src/common/datasource/Cargo.toml +++ b/src/common/datasource/Cargo.toml @@ -39,6 +39,7 @@ parquet.workspace = true paste.workspace = true regex.workspace = true serde.workspace = true +serde_json.workspace = true snafu.workspace = true strum.workspace = true tokio.workspace = true diff --git a/src/common/datasource/src/lib.rs b/src/common/datasource/src/lib.rs index c832831da80..ebec28d7129 100644 --- a/src/common/datasource/src/lib.rs +++ b/src/common/datasource/src/lib.rs @@ -20,6 +20,7 @@ pub mod file_format; pub mod lister; pub mod object_store; pub mod packed_snapshot; +pub mod packed_writer; pub mod parquet_writer; pub mod share_buffer; #[cfg(test)] diff --git a/src/common/datasource/src/packed_writer.rs b/src/common/datasource/src/packed_writer.rs new file mode 100644 index 00000000000..635a1a7a771 --- /dev/null +++ b/src/common/datasource/src/packed_writer.rs @@ -0,0 +1,669 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Streaming destinations for one schema chunk of independent Parquet streams. + +use std::sync::Arc; + +use bytes::Bytes; +use object_store::{ObjectStore, Writer}; +use snafu::{ResultExt, ensure}; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, Result}; +use crate::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex, PackObject, PackTable}; + +pub const WRITE_BYTES: usize = 8 * 1024 * 1024; + +/// One request owns this sink. Finishers retain their encoder admission while +/// waiting on the mutex, bounding both completed buffers and pending appends. +pub type PackedWriterRef = Arc>; + +#[derive(Clone, Copy)] +struct UploadLimits { + part: usize, + object: u64, +} + +impl UploadLimits { + fn for_store(store: &ObjectStore) -> Result { + let info = store.info(); + let caps = info.capability(); + let part = caps + .write_multi_max_size + .unwrap_or(WRITE_BYTES) + .min(WRITE_BYTES); + ensure!( + part > 0 && caps.write_multi_min_size.unwrap_or(0) <= part, + error::ParquetWriterResourceSnafu { + reason: "backend cannot accept bounded 8 MiB parts" + } + ); + // OpenDAL exposes part size but not part count. These are the limits of + // its multipart implementations, including GCS's XML multipart writer. + let parts = match info.scheme() { + "s3" | "oss" | "gcs" => 10_000, + "azblob" => 50_000, + _ => u64::MAX, + }; + let object = (caps.write_total_max_size.unwrap_or(usize::MAX) as u64) + .min(parts.saturating_mul(part as u64)); + Ok(Self { part, object }) + } +} + +struct Upload { + writer: Writer, + store: ObjectStore, + path: String, + limits: UploadLimits, + length: u64, + conditional: bool, + close_started: bool, +} + +impl Upload { + async fn open(store: &ObjectStore, path: String, limits: UploadLimits) -> Result { + let conditional = store.info().capability().write_with_if_not_exists; + if !conditional { + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + error::InvalidPackedSnapshotSnafu { + reason: format!("output already exists: {path}") + } + ); + } + let writer = store + .writer_with(&path) + .chunk(limits.part) + .concurrent(1) + .if_not_exists(conditional) + .await + .context(error::WriteObjectSnafu { path: &path })?; + Ok(Self { + writer, + store: store.clone(), + path, + limits, + length: 0, + conditional, + close_started: false, + }) + } + + async fn write(&mut self, bytes: Bytes) -> Result<()> { + ensure!( + (bytes.len() as u64) <= self.limits.object.saturating_sub(self.length), + error::ParquetWriterResourceSnafu { + reason: "packed export object or multipart part-count limit exceeded" + } + ); + for offset in (0..bytes.len()).step_by(self.limits.part) { + self.writer + .write(bytes.slice(offset..(offset + self.limits.part).min(bytes.len()))) + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + } + self.length += bytes.len() as u64; + Ok(()) + } + + async fn write_json_array( + &mut self, + values: &[T], + token: &CancellationToken, + ) -> Result<()> { + self.write(Bytes::from_static(b"[")).await?; + for (i, value) in values.iter().enumerate() { + check_cancelled(Some(token))?; + if i > 0 { + self.write(Bytes::from_static(b",")).await?; + } + let bytes = serde_json::to_vec(value).map_err(|e| { + error::InvalidPackedSnapshotSnafu { + reason: e.to_string(), + } + .build() + })?; + self.write(bytes.into()).await?; + } + self.write(Bytes::from_static(b"]")).await + } + + async fn close(&mut self) -> Result<()> { + self.close_started = true; + self.writer + .close() + .await + .context(error::WriteObjectSnafu { path: &self.path })?; + Ok(()) + } + + async fn abort(mut self) -> Result<()> { + let result = self.writer.abort().await; + if !self.conditional + && result.as_ref().is_err_and(|e| { + e.kind() == object_store::ErrorKind::Unsupported + && (!self.close_started + || object_store::secure_fs::is_unsynced_overwrite_abort(e)) + }) + { + let store = self.store.clone(); + let path = self.path.clone(); + drop(self); + store + .delete(&path) + .await + .context(error::WriteObjectSnafu { path })?; + } else { + result.context(error::WriteObjectSnafu { path: &self.path })?; + } + Ok(()) + } +} + +/// Pack/index metadata is accounted separately from retained Arrow payloads. +pub struct PackedWriter { + store: ObjectStore, + limits: UploadLimits, + pack: Option, + index: PackIndex, + next_pack: usize, +} + +impl PackedWriter { + pub fn new(store: ObjectStore) -> Result { + let limits = UploadLimits::for_store(&store)?; + Ok(Arc::new(Mutex::new(Self { + store, + limits, + pack: None, + index: PackIndex { + version: 1, + objects: vec![], + tables: vec![], + }, + next_pack: 0, + }))) + } + + async fn close_pack(&mut self) -> Result<()> { + if let Some(pack) = &mut self.pack { + pack.close().await?; + self.index.objects.push(PackObject { + path: pack.path.clone(), + kind: ObjectKind::Pack, + length: pack.length, + }); + self.pack = None; + } + Ok(()) + } + + async fn append( + &mut self, + name: String, + bytes: Bytes, + rows: u64, + token: Option<&CancellationToken>, + ) -> Result<()> { + check_cancelled(token)?; + if self + .pack + .as_ref() + .is_some_and(|p| bytes.len() as u64 > self.limits.object.saturating_sub(p.length)) + { + self.close_pack().await?; + } + check_cancelled(token)?; + if self.pack.is_none() { + let path = format!("pack-{:06}.bin", self.next_pack); + self.next_pack += 1; + self.pack = Some(Upload::open(&self.store, path, self.limits).await?); + } + let pack = self.pack.as_mut().ok_or_else(|| { + error::InvalidPackedSnapshotSnafu { + reason: "missing pack writer", + } + .build() + })?; + let table = PackTable { + table_name: name, + object: pack.path.clone(), + offset: pack.length, + length: bytes.len() as u64, + row_count: rows, + }; + pack.write(bytes).await?; + check_cancelled(token)?; + self.index.tables.push(table); + Ok(()) + } + + /// Publish the index only after every data object has closed successfully. + pub async fn finish(&mut self, token: &CancellationToken) -> Result> { + check_cancelled(Some(token))?; + self.close_pack().await?; + check_cancelled(Some(token))?; + self.index + .tables + .sort_unstable_by(|a, b| a.table_name.cmp(&b.table_name)); + self.index.validate()?; + let mut upload = Upload::open(&self.store, PACK_INDEX_FILE.into(), self.limits).await?; + let result = async { + upload + .write(Bytes::from_static(b"{\"version\":1,\"objects\":")) + .await?; + upload.write_json_array(&self.index.objects, token).await?; + upload.write(Bytes::from_static(b",\"tables\":")).await?; + upload.write_json_array(&self.index.tables, token).await?; + upload.write(Bytes::from_static(b"}")).await?; + check_cancelled(Some(token))?; + upload.close().await?; + check_cancelled(Some(token)) + } + .await; + if let Err(error) = result { + if let Err(secondary) = upload.abort().await { + common_telemetry::warn!(secondary; "Failed to abort pack index"); + } + return Err(error); + } + Ok(self + .index + .objects + .iter() + .map(|o| o.path.clone()) + .chain([PACK_INDEX_FILE.into()]) + .collect()) + } + + /// Call only after all admitted encoders have drained their I/O. + pub async fn abort(&mut self) -> Result<()> { + if let Some(pack) = self.pack.take() { + pack.abort().await?; + } + Ok(()) + } +} + +/// Buffers a small stream, or spills its prefix and continues the same stream. +pub struct PackedTableWriter { + shared: PackedWriterRef, + name: String, + path: String, + buffer: Vec, + ordinary: bool, + standalone: Option, +} + +impl PackedTableWriter { + /// The caller retains encoder admission until finish/abort returns. + pub fn new(shared: PackedWriterRef, name: String, id: u32, ordinary: bool) -> Self { + Self { + shared, + name, + path: format!("table-{id}.parquet"), + buffer: if ordinary { + Vec::new() + } else { + Vec::with_capacity(WRITE_BYTES) + }, + ordinary, + standalone: None, + } + } + + pub(crate) async fn write(&mut self, bytes: Bytes) -> Result<()> { + if self.standalone.is_none() + && (self.ordinary || self.buffer.len().saturating_add(bytes.len()) > WRITE_BYTES) + { + let sink = self.shared.lock().await; + self.standalone = + Some(Upload::open(&sink.store, self.path.clone(), sink.limits).await?); + } + if let Some(upload) = &mut self.standalone { + if !self.buffer.is_empty() { + upload + .write(std::mem::take(&mut self.buffer).into()) + .await?; + } + upload.write(bytes).await + } else { + self.buffer.extend_from_slice(&bytes); + Ok(()) + } + } + + pub(crate) async fn finish( + &mut self, + rows: u64, + token: Option<&CancellationToken>, + ) -> Result<()> { + check_cancelled(token)?; + if let Some(upload) = &mut self.standalone { + upload.close().await?; + check_cancelled(token)?; + let mut shared = self.shared.lock().await; + shared.index.objects.push(PackObject { + path: self.path.clone(), + kind: ObjectKind::Parquet, + length: upload.length, + }); + shared.index.tables.push(PackTable { + table_name: self.name.clone(), + object: self.path.clone(), + offset: 0, + length: upload.length, + row_count: rows, + }); + } else { + self.shared + .lock() + .await + .append( + self.name.clone(), + // OpenDAL may retain each slice until the pack part fills. + // Release the encoder's 8 MiB capacity before queuing it. + Bytes::from(std::mem::take(&mut self.buffer).into_boxed_slice()), + rows, + token, + ) + .await?; + } + Ok(()) + } + + pub(crate) async fn abort(&mut self) -> Result<()> { + if let Some(upload) = self.standalone.take() { + upload.abort().await?; + } + Ok(()) + } +} + +fn check_cancelled(token: Option<&CancellationToken>) -> Result<()> { + ensure!( + token.is_none_or(|t| !t.is_cancelled()), + error::ParquetWriteCancelledSnafu + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use arrow::array::{ArrayRef, Int64Array, StringArray}; + use arrow::record_batch::RecordBatch; + use futures::TryStreamExt; + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use parquet::arrow::ParquetRecordBatchStreamBuilder; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + use super::*; + use crate::file_format::parquet::packed_reader::{PackReadWindows, PackedParquetReader}; + use crate::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; + + #[tokio::test] + async fn independent_streams_empty_and_spilled_roundtrip() { + let directory = common_test_util::temp_dir::create_temp_dir("packed-writer"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + let small = RecordBatch::try_from_iter([( + "tag", + Arc::new(StringArray::from(vec![Some(""), None, Some("value")])) as ArrayRef, + )]) + .unwrap(); + let mut state = 17u64; + let values = (0..600_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as i64 + }) + .collect::>(); + let random = Arc::new(Int64Array::from(values)) as ArrayRef; + let large = RecordBatch::try_from_iter([("a", random.clone()), ("b", random)]).unwrap(); + let batches = [small.clone(), RecordBatch::new_empty(large.schema()), large]; + for (id, batch) in batches.iter().enumerate() { + let sink = + PackedTableWriter::new(shared.clone(), format!("table{id}"), id as u32, false); + let mut writer = ParquetFileWriter::open_packed( + batch.schema(), + store.clone(), + "unused", + Some(ParquetWriterLimits { + row_group_rows: 8192, + flush_threshold_bytes: WRITE_BYTES, + max_row_groups: 4096, + }), + sink, + ) + .unwrap(); + writer.write(batch.clone(), None).await.unwrap(); + writer.finish(None).await.unwrap(); + } + let inventory = shared + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap(); + index + .validate_membership(["table0", "table1", "table2"]) + .unwrap(); + assert_eq!(inventory.len(), 3); + assert!(sizes.lock().unwrap().iter().all(|n| *n <= WRITE_BYTES)); + assert_eq!( + index + .objects + .iter() + .filter(|o| o.kind == ObjectKind::Pack) + .count(), + 1 + ); + assert!( + index + .objects + .iter() + .any(|o| o.kind == ObjectKind::Parquet && o.length > WRITE_BYTES as u64) + ); + let windows = PackReadWindows::new(store.clone()); + for (entry, expected) in index.tables.iter().zip(batches) { + let object = index + .objects + .iter() + .find(|o| o.path == entry.object) + .unwrap(); + let actual = if object.kind == ObjectKind::Pack { + let reader = PackedParquetReader::new( + windows.clone(), + object.path.clone(), + object.length, + entry.offset, + entry.length, + ) + .unwrap(); + ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect::>() + .await + .unwrap() + } else { + ParquetRecordBatchReaderBuilder::try_new( + store.read(&object.path).await.unwrap().to_bytes(), + ) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap() + }; + assert_eq!(entry.row_count, expected.num_rows() as u64); + assert_eq!( + arrow::compute::concat_batches(&expected.schema(), &actual).unwrap(), + expected + ); + } + } + + struct Observe(oio::Writer, Arc>>); + impl oio::Write for Observe { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.1.lock().unwrap().push(bytes.len()); + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + self.0.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + #[tokio::test] + async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() { + let directory = common_test_util::temp_dir::create_temp_dir("packed-limit"); + let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let factory: MockWriterFactory = Arc::new({ + let sizes = sizes.clone(); + move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + }); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator() + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + // Two data parts per object; no oversized fixture is needed. + shared.lock().await.limits = UploadLimits { + part: 512, + object: 1024, + }; + for id in 0..3 { + let mut table = PackedTableWriter::new(shared.clone(), format!("t{id}"), id, false); + table.write(vec![id as u8; 600].into()).await.unwrap(); + table.finish(0, None).await.unwrap(); + } + let mut table = PackedTableWriter::new(shared.clone(), "large".into(), 4, true); + table.write(vec![0; 1024].into()).await.unwrap(); + assert!(matches!( + table.write(Bytes::from_static(b"x")).await, + Err(error::Error::ParquetWriterResource { .. }) + )); + table.abort().await.unwrap(); + assert!(!store.exists("table-4.parquet").await.unwrap()); + shared + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = + serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap(); + assert_eq!(index.objects.len(), 3); + assert!( + index + .tables + .iter() + .all(|t| t.offset == 0 && t.length == 600) + ); + assert!(sizes.lock().unwrap().iter().all(|n| *n <= 512)); + } + + struct FailClose(oio::Writer); + impl oio::Write for FailClose { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.0.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected close failure", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + self.0.abort().await + } + } + + #[tokio::test] + async fn close_failure_cancellation_and_collision_never_publish_index() { + for failure in ["pack-000000.bin", PACK_INDEX_FILE, "cancel", "collision"] { + let directory = common_test_util::temp_dir::create_temp_dir("packed-failure"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + if failure == "collision" { + store.write("pack-000000.bin", "winner").await.unwrap(); + } + let factory: MockWriterFactory = Arc::new(move |path, _, writer| { + if path.trim_start_matches('/') == failure { + Box::new(FailClose(writer)) + } else { + writer + } + }); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let shared = PackedWriter::new(store.clone()).unwrap(); + let mut table = PackedTableWriter::new(shared.clone(), "t".into(), 0, false); + table.write(vec![0; 12].into()).await.unwrap(); + table.finish(0, None).await.unwrap(); + let token = CancellationToken::new(); + if failure == "cancel" { + token.cancel(); + } + assert!(shared.lock().await.finish(&token).await.is_err()); + let _ = shared.lock().await.abort().await; + assert!(!store.exists(PACK_INDEX_FILE).await.unwrap()); + if failure == "collision" { + assert_eq!( + store.read("pack-000000.bin").await.unwrap().to_bytes(), + Bytes::from_static(b"winner") + ); + } + } + } +} diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index ed35c314339..243a46b4e1a 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -28,6 +28,7 @@ use tokio_util::sync::CancellationToken; use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::error::{self, Result}; +use crate::packed_writer::PackedTableWriter; /// Destination creation policy; conditional failures never authorize path deletion. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -52,7 +53,8 @@ pub struct ParquetWriterLimits { /// operation can leave a filesystem worker running after cleanup. pub struct ParquetFileWriter { encoder: Option>>, - sink: Writer, + sink: ParquetSink, + rows: u64, store: ObjectStore, path: String, limits: Option, @@ -60,6 +62,11 @@ pub struct ParquetFileWriter { close_started: bool, } +enum ParquetSink { + Object(Writer), + Packed(Box), +} + impl ParquetFileWriter { /// Open a file using COPY's encoding settings. Destination ownership and /// overwrite policy belong to the caller. None preserves Parquet's defaults. @@ -90,31 +97,7 @@ impl ParquetFileWriter { limits: Option, creation: ParquetCreationPolicy, ) -> Result { - let mut props = WriterProperties::builder() - .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_statistics_truncate_length(None) - .set_column_index_truncate_length(None); - if let Some(limits) = limits { - ensure!( - limits.row_group_rows > 0 - && limits.flush_threshold_bytes > 0 - && limits.max_row_groups > 0, - error::InvalidParquetWriterLimitsSnafu - ); - props = props - .set_max_row_group_row_count(Some(limits.row_group_rows)) - .set_max_row_group_bytes(None); - } - for field in schema.fields() { - if matches!(field.data_type(), DataType::Timestamp(_, _)) { - let column = ColumnPath::new(vec![field.name().clone()]); - props = props - .set_column_dictionary_enabled(column.clone(), false) - .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED); - } - } - let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build())) - .context(error::WriteParquetSnafu { path })?; + let encoder = build_encoder(schema, limits, path)?; let sink = store .writer_with(path) .concurrent(concurrency) @@ -124,7 +107,8 @@ impl ParquetFileWriter { .context(error::WriteObjectSnafu { path })?; Ok(Self { encoder: Some(encoder), - sink, + sink: ParquetSink::Object(sink), + rows: 0, store, path: path.to_owned(), limits, @@ -133,6 +117,26 @@ impl ParquetFileWriter { }) } + /// Encode into a request-owned pack or standalone destination. + pub fn open_packed( + schema: SchemaRef, + store: ObjectStore, + path: &str, + limits: Option, + packed: PackedTableWriter, + ) -> Result { + Ok(Self { + encoder: Some(build_encoder(schema, limits, path)?), + sink: ParquetSink::Packed(Box::new(packed)), + rows: 0, + store, + path: path.into(), + limits, + creation: ParquetCreationPolicy::IfNotExists, + close_started: false, + }) + } + /// Write a batch, enforcing file limits across batch and row-group boundaries. pub async fn write( &mut self, @@ -182,6 +186,7 @@ impl ParquetFileWriter { check_cancelled(cancellation)?; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; + self.rows += len as u64; offset += len; } Ok(()) @@ -189,11 +194,14 @@ impl ParquetFileWriter { async fn write_bytes(&mut self, bytes: Vec) -> Result<()> { let bytes = Bytes::from(bytes); + let sink = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.write(bytes).await, + ParquetSink::Object(sink) => sink, + }; let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; // Slices retain the complete encoded allocation until its last submission. for offset in (0..bytes.len()).step_by(chunk) { - self.sink - .write(bytes.slice(offset..(offset + chunk).min(bytes.len()))) + sink.write(bytes.slice(offset..(offset + chunk).min(bytes.len()))) .await .context(error::WriteObjectSnafu { path: &self.path })?; } @@ -218,9 +226,12 @@ impl ParquetFileWriter { .context(error::JoinHandleSnafu)??; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; + let sink = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.finish(self.rows, cancellation).await, + ParquetSink::Object(sink) => sink, + }; self.close_started = true; - self.sink - .close() + sink.close() .await .context(error::WriteObjectSnafu { path: &self.path })?; check_cancelled(cancellation)?; @@ -230,7 +241,10 @@ impl ParquetFileWriter { /// Abort after all in-flight operations complete. Preserve ambiguous commits; /// conditional callers delegate cleanup exclusively to the backend. pub async fn abort(mut self) -> Result<()> { - let result = self.sink.abort().await; + let result = match &mut self.sink { + ParquetSink::Packed(packed) => return packed.abort().await, + ParquetSink::Object(sink) => sink.abort().await, + }; if self.creation == ParquetCreationPolicy::Overwrite && result.as_ref().is_err_and(|error| { error.kind() == object_store::ErrorKind::Unsupported @@ -253,6 +267,39 @@ impl ParquetFileWriter { } } +fn build_encoder( + schema: SchemaRef, + limits: Option, + path: &str, +) -> Result>> { + let mut props = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_statistics_truncate_length(None) + .set_column_index_truncate_length(None); + if let Some(limits) = limits { + ensure!( + limits.row_group_rows > 0 + && limits.flush_threshold_bytes > 0 + && limits.max_row_groups > 0, + error::InvalidParquetWriterLimitsSnafu + ); + props = props + .set_max_row_group_row_count(Some(limits.row_group_rows)) + .set_max_row_group_bytes(None); + } + for field in schema.fields() { + if matches!(field.data_type(), DataType::Timestamp(_, _)) { + let column = ColumnPath::new(vec![field.name().clone()]); + props = props + .set_column_dictionary_enabled(column.clone(), false) + .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED); + } + } + let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build())) + .context(error::WriteParquetSnafu { path })?; + Ok(encoder) +} + fn check_cancelled(cancellation: Option<&CancellationToken>) -> Result<()> { ensure!( cancellation.is_none_or(|token| !token.is_cancelled()), @@ -498,7 +545,7 @@ mod tests { .await .unwrap(); // No storage-layer chunking: observe the application's actual submissions. - writer.sink = store.writer("large.parquet").await.unwrap(); + writer.sink = ParquetSink::Object(store.writer("large.parquet").await.unwrap()); writer.write(batch.clone(), None).await.unwrap(); writer.finish(None).await.unwrap(); let sizes = sizes.lock().unwrap().clone(); @@ -701,7 +748,8 @@ mod tests { .await .unwrap(); // Force footer bytes through the sink before its close operation. - writer.sink = store.writer_with("cancel.parquet").chunk(1).await.unwrap(); + writer.sink = + ParquetSink::Object(store.writer_with("cancel.parquet").chunk(1).await.unwrap()); let cancellation = CancellationToken::new(); let result = { let finish = writer.finish(Some(&cancellation)); diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index d3d44f7655f..f9727be6c12 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -24,6 +24,7 @@ use common_datasource::file_format::csv::stream_to_csv; use common_datasource::file_format::json::stream_to_json; use common_datasource::file_format::parquet::stream_to_parquet; use common_datasource::object_store::build_backend_for_write_with_path; +use common_datasource::packed_writer::PackedTableWriter; use common_datasource::parquet_writer::ParquetFileWriter; use common_query::Output; use common_recordbatch::adapter::DfRecordBatchStreamAdapter; @@ -127,7 +128,7 @@ impl StatementExecutor { req: CopyTableRequest, query_ctx: QueryContextRef, ) -> Result { - self.copy_captured_table_to_managed(table, req, query_ctx, None) + self.copy_captured_table_to_managed(table, req, query_ctx, None, None) .await } @@ -137,6 +138,7 @@ impl StatementExecutor { req: CopyTableRequest, query_ctx: QueryContextRef, managed: Option<(&ExportWriteBudget, &CancellationToken)>, + packed: Option, ) -> Result { let info = table.table_info(); let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name); @@ -185,7 +187,7 @@ impl StatementExecutor { } = &req; debug!("Copy table: {table_id} to location: {location}"); - self.copy_to_file_managed(&format, output, location, connection, managed) + self.copy_to_file_managed(&format, output, location, connection, managed, packed) .await } @@ -196,7 +198,7 @@ impl StatementExecutor { location: &str, connection: &HashMap, ) -> Result { - self.copy_to_file_managed(format, output, location, connection, None) + self.copy_to_file_managed(format, output, location, connection, None, None) .await } @@ -207,6 +209,7 @@ impl StatementExecutor { location: &str, connection: &HashMap, managed: Option<(&ExportWriteBudget, &CancellationToken)>, + packed: Option, ) -> Result { let output = if managed.is_none() { output @@ -229,7 +232,15 @@ impl StatementExecutor { violated: format!("Expected filename, path: {location}"), })?; if let Some((budget, token)) = managed { - stream_to_managed_parquet(stream, backend.object_store, &filename, budget, token).await + stream_to_managed_parquet_with_packed( + stream, + backend.object_store, + &filename, + budget, + token, + packed, + ) + .await } else { self.stream_to_file(stream, format, backend.object_store, &filename) .await @@ -237,12 +248,24 @@ impl StatementExecutor { } } +#[cfg(test)] pub(crate) async fn stream_to_managed_parquet( + stream: SendableRecordBatchStream, + store: ObjectStore, + path: &str, + budget: &ExportWriteBudget, + token: &CancellationToken, +) -> Result { + stream_to_managed_parquet_with_packed(stream, store, path, budget, token, None).await +} + +async fn stream_to_managed_parquet_with_packed( mut stream: SendableRecordBatchStream, store: ObjectStore, path: &str, budget: &ExportWriteBudget, token: &CancellationToken, + packed: Option, ) -> Result { use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema}; let original = stream.schema(); @@ -261,10 +284,21 @@ pub(crate) async fn stream_to_managed_parquet( } else { expanded_schema.clone() }; - let mut writer = - ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None) - .await - .context(error::WriteStreamToFileSnafu { path })?; + let packed_destination = packed.is_some(); + let mut writer = if let Some(packed) = packed { + ParquetFileWriter::open_packed( + output_schema.arrow_schema().clone(), + store, + path, + Some( + crate::statement::export_logical_tables::LogicalTableExportLimits::default().writer, + ), + packed, + ) + } else { + ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None).await + } + .context(error::WriteStreamToFileSnafu { path })?; let mut started = false; let result = async { let mut rows = 0; @@ -358,7 +392,9 @@ pub(crate) async fn stream_to_managed_parquet( .await; if result.is_err() { token.cancel(); - if started && let Err(error) = writer.abort().await { + if (started || packed_destination) + && let Err(error) = writer.abort().await + { common_telemetry::warn!(error; "Failed to abort ordinary export file"); } } diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs index 9878560afaf..08162c5d5e6 100644 --- a/src/operator/src/statement/database_copy.rs +++ b/src/operator/src/statement/database_copy.rs @@ -60,7 +60,7 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap) -> Result<()> { if let Some(layout) = options.get("metric_data_layout") { return error::InvalidCopyParameterSnafu { diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index dbe18ba763a..f363c73d7a2 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -20,6 +20,7 @@ use std::future::Future; use common_datasource::file_format::Format; use common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url}; +use common_datasource::packed_writer::{PackedTableWriter, PackedWriter}; use common_meta::key::table_route::TableRouteValue; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -70,7 +71,13 @@ impl StatementExecutor { req: CopyDatabaseRequest, tables: Vec, ) -> Result { - validate_database_export_layout(&req.with)?; + if req + .with + .get("metric_data_layout") + .is_none_or(|v| v != "packed") + { + validate_database_export_layout(&req.with)?; + } validate_database_directory(&req.location)?; let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; ensure!( @@ -205,9 +212,27 @@ impl StatementExecutor { let req = &plan.request; let parallelism = parse_parallelism_from_option_map(&req.with); let budget = ExportWriteBudget::new(parallelism); + let packed = if req + .with + .get("metric_data_layout") + .is_some_and(|v| v == "packed") + { + let store = + build_backend_for_write(&req.location, &req.connection, &self.local_file_access) + .await + .context(error::BuildBackendSnafu)?; + Some( + PackedWriter::new(store).context(error::WriteStreamToFileSnafu { + path: &req.location, + })?, + ) + } else { + None + }; let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| { let ctx = ctx.clone(); let budget = budget.clone(); + let packed = packed.clone(); async move { match job { DatabaseExportJob::Metric(unit) => self @@ -220,6 +245,7 @@ impl StatementExecutor { &token, ctx, budget, + packed, ) .await .map(|summary| summary.rows), @@ -238,19 +264,57 @@ impl StatementExecutor { limit: None, }; let _permit = budget.writer(&token).await?; + let destination = packed.map(|shared| { + PackedTableWriter::new(shared, info.name.clone(), info.table_id(), true) + }); self.copy_captured_table_to_managed( table, copy, ctx, Some((&budget, &token)), + destination, ) .await } } } }) - .await?; - Ok(DatabaseExportSummary { rows, output_files }) + .await; + if let Some(packed) = packed { + let mut packed = packed.lock().await; + let result = match rows { + Ok(rows) => packed + .finish(cancellation) + .await + .context(error::WriteStreamToFileSnafu { + path: &req.location, + }) + .and_then(|files| { + Ok(DatabaseExportSummary { + rows, + output_files: files + .into_iter() + .map(|file| { + DatabaseExportFile::new(&req.location, &file, "") + .map(|f| f.location) + }) + .collect::>()?, + }) + }), + Err(error) => Err(error), + }; + if result.is_err() + && let Err(error) = packed.abort().await + { + common_telemetry::warn!(error; "Failed to abort Metric pack"); + } + result + } else { + Ok(DatabaseExportSummary { + rows: rows?, + output_files, + }) + } } } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 4a2e186d477..368335f8450 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -31,6 +31,7 @@ use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; use common_datasource::object_store::build_backend_for_write; +use common_datasource::packed_writer::{PackedTableWriter, PackedWriterRef}; use common_datasource::parquet_writer::{ ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits, }; @@ -108,6 +109,7 @@ pub struct LogicalTableExport { } pub(crate) struct LogicalTableProjection { + name: String, output: DatabaseExportFile, schema: SchemaRef, projection: Vec, @@ -196,6 +198,7 @@ impl LogicalTableExport { .insert( info.table_id(), LogicalTableProjection { + name: name.clone(), output: DatabaseExportFile::new(directory, name, ".parquet")?, schema, projection: indices, @@ -322,6 +325,7 @@ impl StatementExecutor { &cancellation.child_token(), query_ctx, ExportWriteBudget::new(1), + None, ) .await } @@ -337,6 +341,7 @@ impl StatementExecutor { cancellation: &CancellationToken, query_ctx: QueryContextRef, budget: Arc, + packed: Option, ) -> Result { limits.validate()?; let (store, stream) = tokio::select! { @@ -356,7 +361,7 @@ impl StatementExecutor { Ok((store, stream)) } => result?, }; - export_stream_managed(unit, stream, &store, limits, cancellation, budget).await + export_stream_managed(unit, stream, &store, limits, cancellation, budget, packed).await } } @@ -375,6 +380,7 @@ async fn export_stream( limits, &cancellation.child_token(), ExportWriteBudget::new(1), + None, ) .await } @@ -386,8 +392,10 @@ async fn export_stream_managed( limits: LogicalTableExportLimits, cancellation: &CancellationToken, budget: Arc, + packed: Option, ) -> Result { let mut writers = TableWriters::new(budget.clone()); + writers.packed = packed; let result = write_tables( unit, stream, @@ -535,6 +543,26 @@ struct ActiveWriter { } impl ActiveWriter { + fn open_packed( + table: &LogicalTableProjection, + id: u32, + store: &ObjectStore, + limits: LogicalTableExportLimits, + shared: PackedWriterRef, + ) -> Result { + let path = format!("table-{id}.parquet"); + let packed = PackedTableWriter::new(shared, table.name.clone(), id, false); + let writer = ParquetFileWriter::open_packed( + table.schema.clone(), + store.clone(), + &path, + Some(limits.writer), + packed, + ) + .map_err(|e| map_writer_error(e, &path))?; + Ok(Self { path, writer }) + } + async fn open( table: &LogicalTableProjection, store: &ObjectStore, diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index a03fdf1a80a..7aa6daab127 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -146,6 +146,7 @@ async fn routes_across_batches_and_writes_empty_files() { export_limits(), &CancellationToken::new(), budget.clone(), + None, ) .await .unwrap(); @@ -400,12 +401,17 @@ async fn native_histogram_parquet_roundtrip() { struct PausedFileWriter { inner: Option, + paused: bool, started: Arc, release: Arc, } impl object_store::layers::mock::oio::Write for PausedFileWriter { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + if self.paused { + return self.inner.as_mut().unwrap().write(bytes).await; + } + self.paused = true; let mut inner = self.inner.take().unwrap(); let started = self.started.clone(); let release = self.release.clone(); @@ -439,7 +445,7 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter { #[tokio::test] async fn cancellation_drains_storage_and_preserves_committed_files() { - for large in [false, true] { + for (large, packed) in [(false, false), (true, false), (false, true)] { use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory}; let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open"); let access = @@ -459,6 +465,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { move |_, _, inner| { Box::new(PausedFileWriter { inner: Some(inner), + paused: false, started: started.clone(), release: release.clone(), }) @@ -470,6 +477,16 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { .build() .unwrap(), ); + let store = store.layer(object_store::layers::CapabilityOverrideLayer::new( + move |mut capability| { + if packed { + capability.write_multi_max_size = Some(256); + } + capability + }, + )); + let packed = packed + .then(|| common_datasource::packed_writer::PackedWriter::new(store.clone()).unwrap()); let cancellation = CancellationToken::new(); let (unit, input) = if large { let field = Field::new("value", DataType::Utf8, true); @@ -518,6 +535,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { limits, &cancellation, budget.clone(), + packed.clone(), ); tokio::pin!(export); tokio::select! { @@ -528,6 +546,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { assert!(budget.available().1 < 64 * 1024 * 1024); } let held = budget.available(); + assert_eq!(held.0, 0); cancellation.cancel(); assert!(futures::poll!(&mut export).is_pending()); assert_eq!(budget.available().0, held.0); @@ -542,8 +561,16 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { result, Err(error::Error::LogicalTableExportCancelled { .. }) )); - assert_eq!(store.exists("cpu.v1.parquet").await.unwrap(), !large); - if !large { + if let Some(packed) = &packed { + packed.lock().await.abort().await.unwrap(); + assert!(!store.exists("pack-000000.bin").await.unwrap()); + assert!(!store.exists("pack-index.json").await.unwrap()); + } + assert_eq!( + store.exists("cpu.v1.parquet").await.unwrap(), + !large && packed.is_none() + ); + if !large && packed.is_none() { let (_, batches) = read(&store, "cpu.v1.parquet").await; assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); } @@ -829,7 +856,8 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() { &store, export_limits(), &token, - budget.clone() + budget.clone(), + None, ), export_stream_managed( &b, @@ -837,7 +865,8 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() { &store, export_limits(), &token, - budget.clone() + budget.clone(), + None, ), ordinary, ) diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs index 38a63936a5b..70b3bcd2e32 100644 --- a/src/operator/src/statement/export_logical_tables/writers.rs +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use arrow::record_batch::RecordBatch; +use common_datasource::packed_writer::PackedWriterRef; use futures::stream::FuturesUnordered; use futures::{FutureExt, StreamExt}; use object_store::ObjectStore; @@ -111,6 +112,7 @@ pub(crate) struct Payload { } pub(crate) struct TableWriters { + pub(crate) packed: Option, current: Option<(u32, mpsc::Sender)>, tasks: FuturesUnordered>>, budget: Arc, @@ -120,6 +122,7 @@ impl TableWriters { pub(crate) fn new(budget: Arc) -> Self { Self { current: None, + packed: None, tasks: FuturesUnordered::new(), budget, } @@ -150,7 +153,11 @@ impl TableWriters { self.close_input(); let permit = self.budget.writer(token).await?; self.reap_for_admission(token).await?; - let writer = ActiveWriter::open(table, store, limits).await?; + let writer = if let Some(packed) = &self.packed { + ActiveWriter::open_packed(table, id, store, limits, packed.clone())? + } else { + ActiveWriter::open(table, store, limits).await? + }; let (sender, receiver) = mpsc::channel(2); self.current = Some((id, sender)); let token = token.clone(); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index a5266576cbe..b9fcd0a9420 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -1381,7 +1381,9 @@ impl HttpServer { .route( "/capabilities", routing::get(|| async { - axum::Json(serde_json::json!({"metric_packed_import": 1})) + axum::Json( + serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}), + ) }), ) .route( @@ -1897,7 +1899,7 @@ mod test { assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.json::().await, - serde_json::json!({"metric_packed_import": 1}) + serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}) ); } } diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 1f9161d970f..80169a391cd 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1451,7 +1451,7 @@ impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { #[tokio::test] async fn metric_export_v2_cli_resume_roundtrip() { - metric_export_v2_cli_roundtrip(false).await; + metric_export_v2_cli_roundtrip(false, false).await; } #[tokio::test] @@ -1461,11 +1461,54 @@ async fn metric_export_v2_cli_s3_resume_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) { - metric_export_v2_cli_roundtrip(true).await; + metric_export_v2_cli_roundtrip(true, false).await; } } -async fn metric_export_v2_cli_roundtrip(s3: bool) { +#[tokio::test(flavor = "multi_thread")] +async fn packed_export_v2_cli_local_roundtrip() { + metric_export_v2_cli_roundtrip(false, true).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn packed_export_v2_cli_s3_roundtrip() { + if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) + && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) + { + metric_export_v2_cli_roundtrip(true, true).await; + } +} + +async fn exported_table_path( + store: &object_store::ObjectStore, + packed: bool, + chunk: u32, + name: &str, +) -> String { + let prefix = format!("data/public/{chunk}/"); + if !packed { + return format!("{prefix}{name}.parquet"); + } + let index: common_datasource::packed_snapshot::PackIndex = serde_json::from_slice( + &store + .read(&format!("{prefix}pack-index.json")) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + format!( + "{prefix}{}", + index + .tables + .iter() + .find(|t| t.table_name == name) + .unwrap() + .object + ) +} + +async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { use servers::interceptor::SqlQueryInterceptorRef; let plugins = common_base::Plugins::new(); let faults = Arc::new(FailSecondChunk { @@ -1499,6 +1542,9 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { sql(instance, "INSERT INTO audit VALUES ('a',1,1),('z',NULL,3)").await; sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; sql(instance, "CREATE DATABASE z_later").await; + if packed { + sql(instance, "CREATE DATABASE empty_schema").await; + } sql( instance, "CREATE TABLE z_later.events (ts TIMESTAMP TIME INDEX, val BIGINT)", @@ -1530,6 +1576,9 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); for experimental in [false, true] { for layout in ["packed", "invalid"] { + if experimental && layout == "packed" { + continue; + } let path = destination .path() .join(format!("unsupported-{experimental}-{layout}")); @@ -1544,7 +1593,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { ) .await .remove(0); - let error = result.expect_err("export must reject import-only layouts"); + let error = result.expect_err("export must reject unsupported layouts"); assert!( format!("{error:?}").contains("metric_data_layout"), "{error:?}" @@ -1552,6 +1601,41 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { assert!(!path.exists()); } } + if packed && !s3 { + for window in [None, Some("3ms")] { + let path = destination.path().join("empty-range"); + let uri = url::Url::from_directory_path(&path).unwrap(); + let mut args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + uri.as_str(), + "--schemas", + "public", + "--experimental-metric-export", + "--metric-data-layout", + "packed", + "--no-proxy", + "--start-time", + "1970-01-01T00:00:00Z", + "--end-time", + "1970-01-01T00:00:00Z", + ]; + if let Some(window) = window { + args.extend(["--chunk-time-window", window]); + } + let error = run_data_cli(&args).await.unwrap_err(); + assert!( + error + .to_string() + .contains("Packed export requires --start-time to be earlier than --end-time"), + "{error}" + ); + assert!(!path.exists()); + } + } let (uri, store, storage_args) = if s3 { let endpoint = std::env::var("GT_S3_ENDPOINT_URL").unwrap(); let bucket = std::env::var("GT_S3_BUCKET").unwrap(); @@ -1604,7 +1688,11 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { "--to", &uri, "--schemas", - "public,z_later", + if packed { + "public,z_later,empty_schema" + } else { + "public,z_later" + }, "--experimental-metric-export", "--no-proxy", "--start-time", @@ -1617,6 +1705,9 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { "never", ]; args.extend(storage_args.iter().map(String::as_str)); + if packed { + args.extend(["--metric-data-layout", "packed"]); + } assert!(run_data_cli(&args).await.is_err()); let before: cli::export_v2::manifest::Manifest = serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); @@ -1631,7 +1722,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { if s3 { assert!( store - .stat("data/public/1/bulk.parquet") + .stat(&exported_table_path(&store, packed, 1, "bulk").await) .await .unwrap() .content_length() @@ -1642,7 +1733,12 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { for path in &before.chunks[0].files { preserved.push((path.clone(), store.read(path).await.unwrap().to_vec())); } - assert!(store.exists("data/public/2/audit.parquet").await.unwrap()); + assert!( + store + .exists(&exported_table_path(&store, packed, 2, "audit").await) + .await + .unwrap() + ); store .write("data/public/2/unknown.txt", "keep") .await @@ -1666,6 +1762,59 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { let after: cli::export_v2::manifest::Manifest = serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); assert!(after.is_complete()); + assert_eq!(after.version, if packed { 2 } else { 1 }); + assert_eq!(before.data_layout, after.data_layout); + if packed { + use common_datasource::packed_snapshot::{ObjectKind, PackIndex}; + for chunk in &after.chunks { + let empty: PackIndex = serde_json::from_slice( + &store + .read(&format!("data/empty_schema/{}/pack-index.json", chunk.id)) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + empty.validate_membership([]).unwrap(); + let prefix = format!("data/public/{}/", chunk.id); + let index: PackIndex = serde_json::from_slice( + &store + .read(&format!("{prefix}pack-index.json")) + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + index + .validate_membership(names.iter().map(String::as_str)) + .unwrap(); + assert_eq!( + index + .objects + .iter() + .filter(|o| o.kind == ObjectKind::Pack) + .count(), + 1 + ); + assert!(index.tables.iter().filter(|t| t.row_count == 0).count() >= 2); + let mut expected = index + .objects + .iter() + .map(|o| format!("{prefix}{}", o.path)) + .collect::>(); + expected.push(format!("{prefix}pack-index.json")); + expected.sort(); + assert_eq!( + chunk + .files + .iter() + .filter(|f| f.starts_with(&prefix)) + .cloned() + .collect::>(), + expected + ); + } + } assert_eq!( serde_json::to_value(&before.chunks[0]).unwrap(), serde_json::to_value(&after.chunks[0]).unwrap() @@ -1727,6 +1876,39 @@ async fn metric_export_v2_cli_roundtrip(s3: bool) { ) .await ); + if packed { + let schema_uri = format!("{uri}/schema-only"); + let mut schema_args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &schema_uri, + "--schemas", + "public", + "--schema-only", + "--experimental-metric-export", + "--metric-data-layout", + "packed", + "--no-proxy", + "--progress", + "never", + ]; + schema_args.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&schema_args).await.unwrap(); + let schema_manifest: cli::export_v2::manifest::Manifest = serde_json::from_slice( + &store + .read("schema-only/manifest.json") + .await + .unwrap() + .to_bytes(), + ) + .unwrap(); + assert!(schema_manifest.schema_only && schema_manifest.chunks.is_empty()); + assert_eq!(schema_manifest.version, 1); + assert!(schema_manifest.data_layout.is_none()); + } server.abort(); target_server.abort(); if s3 { diff --git a/tests/compatibility/cases/snapshot_parquet_copy/setup.sql b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql index 8342e1847b5..38bc7c74aec 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/setup.sql +++ b/tests/compatibility/cases/snapshot_parquet_copy/setup.sql @@ -1,4 +1,6 @@ -CREATE TABLE snapshot_values (ts TIMESTAMP TIME INDEX, val BIGINT); -INSERT INTO snapshot_values VALUES (1, 42), (2, 7); +CREATE TABLE snapshot_values (ts TIMESTAMP TIME INDEX, val BIGINT, host STRING); +INSERT INTO snapshot_values VALUES (1, 42, ''), (2, 7, NULL); COPY snapshot_values TO '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); TRUNCATE TABLE snapshot_values; +CREATE TABLE snapshot_empty (ts TIMESTAMP TIME INDEX, val DOUBLE); +COPY snapshot_empty TO '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.result b/tests/compatibility/cases/snapshot_parquet_copy/verify.result index d3e4d2483ab..6c590e1a1ad 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/verify.result +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.result @@ -4,9 +4,30 @@ Affected Rows: 2 SELECT * FROM snapshot_values ORDER BY ts; -+-------------------------+-----+ -| ts | val | -+-------------------------+-----+ -| 1970-01-01T00:00:00.001 | 42 | -| 1970-01-01T00:00:00.002 | 7 | -+-------------------------+-----+ ++-------------------------+-----+------+ +| ts | val | host | ++-------------------------+-----+------+ +| 1970-01-01T00:00:00.001 | 42 | | +| 1970-01-01T00:00:00.002 | 7 | | ++-------------------------+-----+------+ + +SELECT val, host IS NULL AS host_is_null, host = '' AS host_is_empty FROM snapshot_values ORDER BY ts; + ++-----+--------------+---------------+ +| val | host_is_null | host_is_empty | ++-----+--------------+---------------+ +| 42 | false | true | +| 7 | true | | ++-----+--------------+---------------+ + +COPY snapshot_empty FROM '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); + +Affected Rows: 0 + +SELECT COUNT(*) FROM snapshot_empty; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ diff --git a/tests/compatibility/cases/snapshot_parquet_copy/verify.sql b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql index c782dcacb1b..0e35e3c127f 100644 --- a/tests/compatibility/cases/snapshot_parquet_copy/verify.sql +++ b/tests/compatibility/cases/snapshot_parquet_copy/verify.sql @@ -1,2 +1,5 @@ COPY snapshot_values FROM '${SQLNESS_HOME}/snapshot_parquet_copy/values.parquet' WITH (FORMAT='parquet'); SELECT * FROM snapshot_values ORDER BY ts; +SELECT val, host IS NULL AS host_is_null, host = '' AS host_is_empty FROM snapshot_values ORDER BY ts; +COPY snapshot_empty FROM '${SQLNESS_HOME}/snapshot_parquet_copy/empty.parquet' WITH (FORMAT='parquet'); +SELECT COUNT(*) FROM snapshot_empty; From dd2c1d1acaa924449071839332a49619ca8ba28c Mon Sep 17 00:00:00 2001 From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Date: Tue, 29 Sep 2026 00:07:49 +0000 Subject: [PATCH 066/133] fix(meta-srv): use NoTls for disabled and Unix socket Postgres KV backends (#9059) * fix(meta-srv): use NoTls for disabled and Unix socket Postgres KV backends Signed-off-by: Tyagiquamar * fix(meta-srv): treat Postgres config as unix socket only when every host is a socket tokio-postgres dials hostaddr over TCP even when host is a socket path, and a mixed host list with Require/VerifyFull TLS would otherwise end up sending plaintext over TCP. is_unix_socket_url now parses via tokio_postgres::Config and requires no hostaddr and all-Unix hosts. Adds regression cases for the libpq keyword form, the user@ percent encoded socket URL, and the mixed host / hostaddr negatives. Signed-off-by: Tyagiquamar --------- Signed-off-by: Tyagiquamar --- src/meta-srv/src/utils/postgres.rs | 120 ++++++++++++++++++++++++++--- 1 file changed, 108 insertions(+), 12 deletions(-) diff --git a/src/meta-srv/src/utils/postgres.rs b/src/meta-srv/src/utils/postgres.rs index ec46c076e26..148a654e2cd 100644 --- a/src/meta-srv/src/utils/postgres.rs +++ b/src/meta-srv/src/utils/postgres.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(unix)] +use std::str::FromStr; + use common_error::ext::BoxedError; use common_meta::election::ElectionRef; use common_meta::election::rds::postgres::{ElectionPgClient, PgElection}; @@ -20,10 +23,13 @@ use common_meta::kv_backend::rds::PgStore; use common_meta::kv_backend::rds::postgres::{ TlsMode as PgTlsMode, TlsOption as PgTlsOption, create_postgres_tls_connector, }; +use common_telemetry::warn; use deadpool_postgres::{Config, Runtime}; use servers::tls::TlsOption; use snafu::{OptionExt, ResultExt}; use tokio_postgres::NoTls; +#[cfg(unix)] +use tokio_postgres::config::Host; use crate::error::{self, Result}; @@ -60,23 +66,57 @@ pub async fn create_postgres_pool( })?; cfg.url = Some(postgres_url.clone()); - let pool = if let Some(tls_config) = tls_config { - let pg_tls_config = convert_tls_option(&tls_config); - let tls_connector = - create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other { - source: BoxedError::new(e), - location: snafu::Location::new(file!(), line!(), 0), - })?; - cfg.create_pool(Some(Runtime::Tokio1), tls_connector) - .context(error::CreatePostgresPoolSnafu)? - } else { - cfg.create_pool(Some(Runtime::Tokio1), NoTls) - .context(error::CreatePostgresPoolSnafu)? + let is_unix_socket = is_unix_socket_url(postgres_url); + if is_unix_socket + && matches!(tls_config.as_ref(), Some(t) if t.mode != servers::tls::TlsMode::Disable) + { + warn!( + "TLS is not supported for Unix domain socket PostgreSQL connections, falling back to NoTls" + ); + } + + let pool = match tls_config { + Some(tls_config) + if tls_config.mode != servers::tls::TlsMode::Disable && !is_unix_socket => + { + let pg_tls_config = convert_tls_option(&tls_config); + let tls_connector = + create_postgres_tls_connector(&pg_tls_config).map_err(|e| error::Error::Other { + source: BoxedError::new(e), + location: snafu::Location::new(file!(), line!(), 0), + })?; + cfg.create_pool(Some(Runtime::Tokio1), tls_connector) + .context(error::CreatePostgresPoolSnafu)? + } + _ => cfg + .create_pool(Some(Runtime::Tokio1), NoTls) + .context(error::CreatePostgresPoolSnafu)?, }; Ok(pool) } +#[cfg(unix)] +fn is_unix_socket_url(url: &str) -> bool { + let Ok(cfg) = tokio_postgres::Config::from_str(url) else { + return false; + }; + // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path, + // so treat the config as a socket only when every host is a Unix socket and + // no `hostaddr` is set. + cfg.get_hostaddrs().is_empty() + && !cfg.get_hosts().is_empty() + && cfg + .get_hosts() + .iter() + .all(|host| matches!(host, Host::Unix(_))) +} + +#[cfg(not(unix))] +fn is_unix_socket_url(_: &str) -> bool { + false +} + /// Builds a Postgres-backed metadata [`KvBackendRef`]. /// /// * `store_addrs` - Postgres connection URLs; only the first address is used. @@ -150,3 +190,59 @@ pub async fn build_postgres_election( .await .context(error::KvBackendSnafu) } + +#[cfg(test)] +mod tests { + use super::is_unix_socket_url; + + #[test] + fn detects_postgres_unix_socket_url() { + #[cfg(unix)] + { + // libpq keyword-value form (issue #7734) + assert!(is_unix_socket_url( + "host=/var/run/postgresql dbname=greptime user=greptime password=secret" + )); + // standard postgres URL with percent-encoded unix socket directory + assert!(is_unix_socket_url( + "postgresql://user:pw@%2Fvar%2Frun%2Fpostgresql/mydb" + )); + assert!(is_unix_socket_url( + "postgresql://user@%2Fvar%2Frun%2Fpostgresql/db" + )); + // postgres URL with socket dir in query param + assert!(is_unix_socket_url( + "postgresql:///mydb?host=%2Fvar%2Frun%2Fpostgresql" + )); + assert!(is_unix_socket_url( + "postgresql://user:secret@/mydb?host=%2Fvar%2Frun%2Fpostgresql" + )); + + // TCP URLs should not be classified as unix socket + assert!(!is_unix_socket_url( + "postgresql://user:pw@localhost:5432/mydb" + )); + assert!(!is_unix_socket_url("postgresql://user@localhost/db")); + assert!(!is_unix_socket_url( + "host=127.0.0.1 port=5432 dbname=greptime user=greptime password=secret" + )); + // mixed socket and TCP hosts must not be treated as a socket-only config: + // a Require/VerifyFull TLS config would otherwise end up in plaintext over TCP + assert!(!is_unix_socket_url( + "host=/var/run/postgresql,db.example.com" + )); + // tokio-postgres dials `hostaddr` over TCP even when `host` is a socket path + assert!(!is_unix_socket_url("host=/tmp hostaddr=10.0.0.5")); + } + + #[cfg(not(unix))] + { + assert!(!is_unix_socket_url( + "host=/var/run/postgresql dbname=greptime user=greptime password=secret" + )); + assert!(!is_unix_socket_url( + "postgresql://user:pw@localhost:5432/mydb" + )); + } + } +} From 0686f72dc6ad32a3ad8a684c82b2eda6fee1e5fa Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 29 Sep 2026 00:37:34 +0000 Subject: [PATCH 067/133] chore: add cargo min-publish-age of 7 days (#9384) Signed-off-by: Dennis Zhuang --- .cargo/config.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.cargo/config.toml b/.cargo/config.toml index 9ffc1961b09..fd372c95b84 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -7,6 +7,11 @@ linker = "riscv64-linux-gnu-gcc" [alias] sqlness = "run --bin sqlness-runner --target-dir target/sqlness --" +[registry] +# Requires cargo 1.100+; older cargo ignores this key. Only registry crates +# carry a publish time, so git dependencies (our GreptimeTeam forks) are exempt. +global-min-publish-age = "7 days" + [unstable.git] shallow_index = true shallow_deps = true From 28f01d2ffeff9ba3018ed24ec2de6a202b7211c6 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 29 Sep 2026 02:21:27 +0000 Subject: [PATCH 068/133] revert(ci): pin the query-regression runner toolchain to nightly-2026-03-21 (#9389) * revert(ci): pin the query-regression runner toolchain to nightly-2026-03-21 The query-regression benchmark compiles both the candidate and the BASE checkout (the previous nightly build). Base refs can predate the stable-toolchain migration and still use #![feature] gates, so the runner toolchain must stay a nightly that can build historical revisions; deriving it from the workspace rust-toolchain.toml (now stable 1.96.1) breaks base builds. Revert the toolchain-toml coupling introduced in #9369 and keep the parts that were correct: - query-regression.yml: RUSTUP_TOOLCHAIN hard-pinned to nightly-2026-03-21 again (with a comment explaining why), Verify assertions back to the exact nightly versions, and the test-tooling pin-derivation machinery removed - runner Dockerfile: ARG RUST_TOOLCHAIN=nightly-2026-03-21 + baked ENV restored; the COPY rust-toolchain.toml parsing removed - build-ecs-image.py: the toml staging in the builder user-data removed; base-image auto-resolution kept but retargeted to Ubuntu 26.04 to match the Verify tool pins (python3 3.14 etc.) - the rebuild job keeps the epoch-bump lockstep, now as a PR from a timestamped ci/ branch mirroring update-dev-builder-version.sh (direct pushes to main fail GH006 under branch protection) Complements #9376 (actions-runner bump), which requires a rebuilt image with the non-deprecated runner. Part of #9289. Signed-off-by: Ning Sun * chore: skip qreg for rust-toolchain change --------- Signed-off-by: Ning Sun --- .../query-regression/Dockerfile | 20 +++---- .../query-regression/README.md | 32 +++++----- .../ecs-image/build-ecs-image.py | 39 +++++------- .github/workflows/query-regression.yml | 32 +++++----- .../workflows/release-dev-builder-images.yaml | 60 ++++++++++++++----- 5 files changed, 100 insertions(+), 83 deletions(-) diff --git a/.github/runner-scale-sets/query-regression/Dockerfile b/.github/runner-scale-sets/query-regression/Dockerfile index d02425d684e..617278adf95 100644 --- a/.github/runner-scale-sets/query-regression/Dockerfile +++ b/.github/runner-scale-sets/query-regression/Dockerfile @@ -22,14 +22,14 @@ ARG SCCACHE_SHA256=aec995a83ad3dff3d14b6314e08858b7b73d35ca85a5bcf3d3a9ec07dee35 ARG RUSTUP_INIT_VERSION=1.29.0 ARG RUSTUP_INIT_TARGET=x86_64-unknown-linux-gnu ARG RUSTUP_INIT_SHA256=4acc9acc76d5079515b46346a485974457b5a79893cfb01112423c89aeb5aa10 - -# Single source of truth for the Rust toolchain: parsed from -# rust-toolchain.toml at build time, so the image always bakes the -# workspace toolchain without a second hard-coded pin here. -COPY rust-toolchain.toml /opt/rust-toolchain.toml +# Pinned to a nightly deliberately: the benchmark compiles historical base +# checkouts (previous nightly builds) that still use #![feature] gates, so +# the runner cannot follow the workspace's rust-toolchain.toml (now stable). +ARG RUST_TOOLCHAIN=nightly-2026-03-21 ENV RUSTUP_HOME=/opt/rustup \ CARGO_HOME=/opt/cargo \ + RUSTUP_TOOLCHAIN=${RUST_TOOLCHAIN} \ RUSTUP_AUTO_INSTALL=0 \ PATH=/opt/cargo/bin:${PATH} @@ -73,14 +73,12 @@ RUN curl --fail --location --silent --show-error \ && install --mode=0755 /tmp/sccache/sccache /usr/local/bin/sccache \ && rm -rf /tmp/sccache.tar.gz /tmp/sccache -RUN rust_toolchain="$(grep -E '^channel' /opt/rust-toolchain.toml | cut -d'"' -f2)" \ - && test -n "${rust_toolchain}" \ - && curl --fail --location --silent --show-error \ +RUN curl --fail --location --silent --show-error \ --output /tmp/rustup-init \ "https://static.rust-lang.org/rustup/archive/${RUSTUP_INIT_VERSION}/${RUSTUP_INIT_TARGET}/rustup-init" \ && echo "${RUSTUP_INIT_SHA256} /tmp/rustup-init" | sha256sum --check --status - \ && chmod 0755 /tmp/rustup-init \ - && /tmp/rustup-init -y --profile minimal --default-toolchain "${rust_toolchain}" --no-modify-path \ + && /tmp/rustup-init -y --profile minimal --default-toolchain "${RUST_TOOLCHAIN}" --no-modify-path \ && rm -f /tmp/rustup-init \ && chown -R root:root /opt/rustup /opt/cargo \ && chmod -R go-w /opt/rustup /opt/cargo \ @@ -91,8 +89,6 @@ USER runner RUN set -eu \ && test "$(id -u)" = "1001" \ - && rust_toolchain="$(grep -E '^channel' /opt/rust-toolchain.toml | cut -d'"' -f2)" \ - && test -n "${rust_toolchain}" \ && temporary_cargo_home="$(mktemp --directory)" \ && temporary_proto_dir="" \ && cleanup() { rm -rf "${temporary_cargo_home}" "${temporary_proto_dir}"; } \ @@ -104,7 +100,7 @@ RUN set -eu \ && rustc --version \ && active_toolchain="$(rustup show active-toolchain)" \ && printf 'Active toolchain: %s\n' "${active_toolchain}" \ - && case "${active_toolchain}" in "${rust_toolchain}-x86_64-unknown-linux-gnu"|"${rust_toolchain}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ + && case "${active_toolchain}" in "${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu"|"${RUST_TOOLCHAIN}-x86_64-unknown-linux-gnu "*) ;; *) exit 1;; esac \ && test ! -w /opt/rustup \ && test ! -w /opt/cargo/bin \ && test -r /usr/include/google/protobuf/any.proto \ diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index 70ed4a88345..44500a8e50f 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -88,12 +88,17 @@ sweep. The runner `Dockerfile` in the parent directory stays the single source of the tool contract. The image is rebuilt **automatically** by the `rebuild-query-regression-runner-image` job in -`.github/workflows/release-dev-builder-images.yaml` whenever -`rust-toolchain.toml` or anything under this directory changes on main (or via -manual dispatch): it runs the ops tool below, then bumps -`RUNNER_IMAGE_EPOCH` in `query-regression.yml` and points the -`QUERY_REGRESSION_ECS_IMAGE_ID` repo variable at the new image, so the next -regression run picks up image and epoch together. +`.github/workflows/release-dev-builder-images.yaml` whenever anything under +this directory changes on main (or via manual dispatch): it runs the ops tool +below, opens an epoch-bump PR for `RUNNER_IMAGE_EPOCH` in +`query-regression.yml`, and points the `QUERY_REGRESSION_ECS_IMAGE_ID` repo +variable at the new image. + +The runner toolchain is **pinned inside the Dockerfile** to +`nightly-2026-03-21` and deliberately does NOT follow the workspace +`rust-toolchain.toml`: the benchmark also compiles the BASE checkout (the +previous nightly build), which may predate the stable-toolchain migration and +still require a nightly compiler. The manual fallback (also what the workflow runs): @@ -105,10 +110,10 @@ uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ ``` `--base-image-id` is optional: the script defaults to the latest public -Ubuntu 24.04 image in the region (the Dockerfile pins every tool version -itself, so base drift is low-risk); pass it — or set the -`ALIYUN_ECS_BASE_IMAGE_ID` repo variable consumed by the automated job — -to pin a specific base image. +Ubuntu 26.04 image in the region (26.04 matches the tool-version pins the +Verify step asserts, e.g. python3 3.14; keep the two in sync); pass it — or +set the `ALIYUN_ECS_BASE_IMAGE_ID` repo variable consumed by the automated +job — to pin a specific base image. The script boots a temporary builder instance, `docker build`s the runner image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and @@ -173,11 +178,10 @@ overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, `Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and -the image-baked Rust toolchain matching `rust-toolchain.toml` (the image -parses the pin from the toml at build time, and the workflow asserts it -dynamically at run time — there is no separately pinned toolchain version). +the image-baked Rust toolchain `nightly-2026-03-21` pinned in the runner +Dockerfile (deliberately independent of `rust-toolchain.toml` — see above). `mold` and `python3` -come from apt at image-build time (not Ubuntu 24.04's default 3.12); if the +come from apt at image-build time; if the Ubuntu archive ships a newer package revision between rebuilds, the Verify step fails with the observed version — bump those pins in `query-regression.yml` when that happens. Everything else (toolchain, uv, sccache, otelgen, rustup, diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py index 78b01e721c3..f9ecb6b2333 100644 --- a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -25,7 +25,7 @@ """Build the query-regression ECS custom image (manual ops tool). -Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 26.04 image, builds the existing runner container image (the Dockerfile in the parent directory remains the single source of the tool contract), materializes the tool directories onto the host filesystem so the workflow's "Verify runner @@ -52,17 +52,15 @@ import time from pathlib import Path ASSETS_DIR = Path(__file__).resolve().parent -# Repo root: ecs-image -> query-regression -> runner-scale-sets -> .github -> root. -REPO_ROOT = ASSETS_DIR.parents[3] DONE_MARKER = "QREG_IMAGE_BUILD_DONE" FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" POLL_INTERVAL_SECONDS = 15 CONSOLE_POLL_INTERVAL_SECONDS = 30 BUILD_TIMEOUT_SECONDS = 60 * 60 -# Same apt package contract as the runner Dockerfile; the base -# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the -# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Same apt package contract as the runner Dockerfile; the base is an +# Ubuntu 26.04 image, so an Ubuntu 26.04 host resolves the same tool +# versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). # Docker itself comes from Docker's official repository (docker-ce), not the # distribution-packaged docker.io. APT_PACKAGES = [ @@ -94,9 +92,8 @@ APT_PACKAGES = [ DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" -def render_user_data(dockerfile: str, rust_toolchain_toml: str, start_runner: str, unit: str) -> str: +def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() - rust_toolchain_b64 = base64.b64encode(rust_toolchain_toml.encode()).decode() start_runner_b64 = base64.b64encode(start_runner.encode()).decode() unit_b64 = base64.b64encode(unit.encode()).decode() packages = " ".join(APT_PACKAGES) @@ -125,11 +122,6 @@ base64 -d > /tmp/Dockerfile <<'EOF' {dockerfile_b64} EOF mkdir -p /tmp/image-context -# The Dockerfile COPYies rust-toolchain.toml (single source of truth for -# the baked toolchain); stage it into the build context. -base64 -d > /tmp/image-context/rust-toolchain.toml <<'EOF' -{rust_toolchain_b64} -EOF docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context # Materialize the tool contract onto the host filesystem. @@ -238,20 +230,22 @@ def call_api_with_retry(fn, description: str, attempts: int = 5): def resolve_base_image_id(client, region_id: str) -> str: - """Resolve the latest public Ubuntu 24.04 x86_64 system image. + """Resolve the latest public Ubuntu 26.04 x86_64 system image. Used as the default for --base-image-id: the runner Dockerfile pins - every tool version itself, so a current stock Ubuntu 24.04 base is all + every tool version itself, so a current stock Ubuntu LTS base is all the builder needs. Pass --base-image-id (or ALIYUN_ECS_BASE_IMAGE_ID) to pin a specific base image deterministically. """ from alibabacloud_ecs20140526 import models as ecs_models - def _is_ubuntu_2404(image) -> bool: - # osname is localized (e.g. "Ubuntu 24.04 64ä½"), osname_en the - # English form; accept either. + def _is_target_ubuntu(image) -> bool: + # osname is localized (e.g. "Ubuntu 26.04 64ä½"), osname_en the + # English form; accept either. Keep the Ubuntu version in sync with + # the tool-version pins asserted by the Verify step in + # query-regression.yml (e.g. python3 3.14 comes from 26.04). for os_name in (image.osname_en, image.osname): - if os_name and "Ubuntu" in os_name and "24.04" in os_name: + if os_name and "Ubuntu" in os_name and "26.04" in os_name: return True return False @@ -277,10 +271,10 @@ def resolve_base_image_id(client, region_id: str) -> str: break page_number += 1 - candidates = [image for image in images if _is_ubuntu_2404(image)] + candidates = [image for image in images if _is_target_ubuntu(image)] if not candidates: raise SystemExit( - "No public Ubuntu 24.04 x86_64 system image found in region " + "No public Ubuntu 26.04 x86_64 system image found in region " f"{region_id}; pass --base-image-id explicitly" ) candidates.sort(key=lambda image: image.creation_time or "", reverse=True) @@ -324,7 +318,7 @@ def main() -> int: from alibabacloud_ecs20140526 import models as ecs_models client = make_ecs_client(args.region_id) - # --base-image-id is optional: default to the latest public Ubuntu 24.04 + # --base-image-id is optional: default to the latest public Ubuntu 26.04 # image in the region (the Dockerfile pins every tool version itself, so # base drift is low-risk; pass --base-image-id or set # ALIYUN_ECS_BASE_IMAGE_ID to pin deterministically). @@ -337,7 +331,6 @@ def main() -> int: user_data = base64.b64encode( render_user_data( (ASSETS_DIR.parent / "Dockerfile").read_text(), - (REPO_ROOT / "rust-toolchain.toml").read_text(), (ASSETS_DIR / "start-runner.sh").read_text(), (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), ).encode() diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index cc8afaf8091..47daa3f0435 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -120,23 +120,12 @@ jobs: # Ordinary PRs also run the same tests from checks.yml. runs-on: ubuntu-latest timeout-minutes: 10 - outputs: - # The Rust toolchain pin, resolved from rust-toolchain.toml so the - # benchmark always runs the workspace toolchain without a second - # hard-coded copy in this workflow. - rust_toolchain: ${{ steps.rust-toolchain.outputs.pin }} steps: - name: Checkout uses: actions/checkout@v4 with: persist-credentials: false - - name: Resolve Rust toolchain pin - id: rust-toolchain - run: | - pin="$(grep -E '^channel' rust-toolchain.toml | cut -d'"' -f2)" - echo "pin=${pin}" >> "$GITHUB_OUTPUT" - - name: Test query regression tooling run: | python3 tests/perf/test_query_regression_runner_compaction_toctou.py @@ -201,7 +190,13 @@ jobs: CARGO_HOME: /home/runner/.cargo UV_CACHE_DIR: /home/runner/.cargo/uv-cache RUSTUP_HOME: /opt/rustup - RUSTUP_TOOLCHAIN: ${{ needs.test-tooling.outputs.rust_toolchain }} + # Deliberately NOT derived from rust-toolchain.toml: the benchmark + # also compiles the BASE checkout (the previous nightly build, which + # may predate the stable-toolchain migration and still uses + # #![feature] gates), so the runner must keep a nightly toolchain + # that can build historical revisions. This pin moves only via a + # deliberate runner-image rebuild. + RUSTUP_TOOLCHAIN: nightly-2026-03-21 RUSTUP_AUTO_INSTALL: "0" CARGO_TARGET_DIR: /home/runner/query-regression-target QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta @@ -515,16 +510,17 @@ jobs: require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' - # The toolchain pin flows from rust-toolchain.toml (resolved in the - # test-tooling job); escape it for the regex assertions below. - pin_regex="${RUSTUP_TOOLCHAIN//./\\.}" + # The runner toolchain is pinned to nightly-2026-03-21 (see the + # RUSTUP_TOOLCHAIN comment): it must compile historical base + # checkouts, so it intentionally does not follow rust-toolchain.toml. require_match cargo "$(capture cargo --version)" \ - "^cargo[[:space:]]${pin_regex}[[:space:]]\([0-9a-f]+[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$" + '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' require_match rustc "$(capture rustc --version)" \ - "^rustc[[:space:]]${pin_regex}[[:space:]]\([0-9a-f]+[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$" + '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' require_match active_toolchain "$(capture rustup show active-toolchain)" \ - "^${pin_regex}-x86_64-unknown-linux-gnu([[:space:]]|$)" + '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" + require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" require "readable /opt/rustup" test -r /opt/rustup require "executable /opt/rustup" test -x /opt/rustup diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index f9a55cb9985..2e1041c1ff4 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -71,9 +71,13 @@ jobs: files="$(git diff --name-only "${base}" "${head}")" dev_builder=false query_regression_runner=false + # rust-toolchain.toml only gates the dev-builder images: they bake + # the workspace toolchain. The query-regression runner toolchain is + # pinned inside its Dockerfile (nightly-2026-03-21) and must NOT + # follow the workspace pin, so its rebuild triggers only on changes + # under its own directory. if grep -qx 'rust-toolchain.toml' <<<"${files}"; then dev_builder=true - query_regression_runner=true fi if grep -q '^docker/dev-builder/' <<<"${files}"; then dev_builder=true @@ -335,18 +339,21 @@ jobs: # Rebuilds the query-regression ECS runner image via the ops tool in # .github/runner-scale-sets/query-regression/ecs-image/: boots a # temporary pay-as-you-go ECS builder, snapshots a new custom image, - # then completes the documented lockstep updates in order: + # then completes the documented lockstep updates: # 1. bumps RUNNER_IMAGE_EPOCH in query-regression.yml (target-cache - # invalidation) and pushes the commit to main, and only then + # invalidation) via a PR from a timestamped ci/ branch (main is + # branch-protected; mirrors update-dev-builder-version.sh), and # 2. points the repo variable QUERY_REGRESSION_ECS_IMAGE_ID at the - # new image, so the next regression run picks up image and epoch - # together. + # new image. The runner toolchain is pinned inside the Dockerfile + # (nightly-2026-03-21, independent of rust-toolchain.toml — the + # benchmark must compile historical base checkouts), so image + # rebuilds do not change the cache ABI. # # Required repository configuration (same names the provisioning job # uses): vars ALIYUN_ECS_REGION_ID, ALIYUN_ECS_VSWITCH_ID, # ALIYUN_ECS_SECURITY_GROUP_ID, optionally ALIYUN_ECS_RESOURCE_GROUP_ID # and ALIYUN_ECS_BASE_IMAGE_ID (deterministic base-image pin; otherwise - # the rebuild auto-resolves the latest public Ubuntu 24.04 image). + # the rebuild auto-resolves the latest public Ubuntu 26.04 image). # Secrets: ALICLOUD_ECS_ACCESS_KEY_ID, ALICLOUD_ECS_ACCESS_KEY_SECRET, # GH_PERSONAL_ACCESS_TOKEN (repo push + actions-variable write). name: Rebuild query-regression runner image @@ -392,7 +399,7 @@ jobs: echo "## Rebuilt runner image" >> "$GITHUB_STEP_SUMMARY" echo "New image: \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" - - name: Bump RUNNER_IMAGE_EPOCH and update the image id variable + - name: Open the epoch-bump PR and update the image id variable env: GH_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} run: | @@ -405,21 +412,42 @@ jobs: exit 1 fi new_epoch=$((old_epoch + 1)) + + # main is branch-protected (required reviews + status checks), so + # the epoch bump cannot be pushed there directly. Mirror + # .github/scripts/update-dev-builder-version.sh: timestamped bot + # branch, plain push, and a PR with reviewers. Do NOT use [skip ci] + # on the commit: the required checks must be able to run for the PR + # to become mergeable. + BRANCH="ci/update-query-regression-epoch-$(date +%Y%m%d%H%M%S)" + git config user.name "greptimedb-ci" + git config user.email "greptimedb-ci@greptime.com" + git fetch origin main + git checkout -b "${BRANCH}" origin/main sed -i "s/readonly RUNNER_IMAGE_EPOCH=\"${old_epoch}\"/readonly RUNNER_IMAGE_EPOCH=\"${new_epoch}\"/" \ .github/workflows/query-regression.yml - - git config user.name "greptimedb-ci" - git config user.email "greptimedb-ci@users.noreply.github.com" git add .github/workflows/query-regression.yml - git commit -m "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch} for image ${image_id} [skip ci]" - git pull --rebase origin main - git push origin HEAD:main + git commit -s -m "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch} for image ${image_id}" + git push origin "${BRANCH}" - # Only point the variable at the new image after the epoch commit - # landed, so the next run picks up image and epoch together. + gh pr create \ + --title "ci(query-regression): bump RUNNER_IMAGE_EPOCH to ${new_epoch}" \ + --body "Auto-generated by the \`Rebuild query-regression runner image\` job in ${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID} for image \`${image_id}\`. Merging this invalidates the query-regression target cache (RUNNER_IMAGE_EPOCH ${old_epoch} → ${new_epoch}). The \`QUERY_REGRESSION_ECS_IMAGE_ID\` repo variable already points at the new image; the epoch bump is belt-and-suspenders for image-content drift." \ + --base main \ + --head "${BRANCH}" \ + --reviewer sunng87 \ + --reviewer daviderli614 \ + --reviewer killme2008 \ + --reviewer evenyag \ + --reviewer fengjiachun \ + --reviewer WenyXu + + # Point the variable at the new image once the PR exists. Image + # rebuilds keep the same pinned toolchain, so the cache ABI is + # unchanged and not waiting for the PR merge is fine. gh api -X PATCH "repos/${{ github.repository }}/actions/variables/QUERY_REGRESSION_ECS_IMAGE_ID" \ -f value="${image_id}" echo "## Lockstep updates" >> "$GITHUB_STEP_SUMMARY" echo "- QUERY_REGRESSION_ECS_IMAGE_ID → \`${image_id}\`" >> "$GITHUB_STEP_SUMMARY" - echo "- RUNNER_IMAGE_EPOCH → \`${new_epoch}\` (committed to main)" >> "$GITHUB_STEP_SUMMARY" + echo "- RUNNER_IMAGE_EPOCH → \`${new_epoch}\` (PR on \`${BRANCH}\`, pending review/merge)" >> "$GITHUB_STEP_SUMMARY" From f8cc842468ba9b2e8031d804f81fe1a85bf3f70f Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 29 Sep 2026 02:37:42 +0000 Subject: [PATCH 069/133] chore: add mold back to nix flake (#9388) --- flake.nix | 1 + 1 file changed, 1 insertion(+) diff --git a/flake.nix b/flake.nix index 6adb69678b2..6e4151fb3e8 100644 --- a/flake.nix +++ b/flake.nix @@ -34,6 +34,7 @@ gcc protobuf gnumake + mold wild (rustToolchain.withComponents [ "cargo" From 133296f298c504b920772292c22f169cd46c54c9 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 29 Sep 2026 03:22:10 +0000 Subject: [PATCH 070/133] chore: bump version to 1.3.0-beta.1 (#9392) Signed-off-by: WenyXu --- Cargo.lock | 154 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 2 +- 2 files changed, 78 insertions(+), 78 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 96903faaa25..6997f0d782b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -230,7 +230,7 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" [[package]] name = "api" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arrow-schema 59.2.0", "common-base", @@ -959,7 +959,7 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "auth" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -1618,7 +1618,7 @@ dependencies = [ [[package]] name = "cache" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "catalog", "common-error", @@ -1684,7 +1684,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow 59.2.0", @@ -2040,7 +2040,7 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "cli" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-stream", "async-trait", @@ -2095,7 +2095,7 @@ dependencies = [ [[package]] name = "client" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arc-swap", @@ -2127,7 +2127,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.3.0-beta.1", "tokio", "tokio-stream", "tonic 0.14.2", @@ -2175,7 +2175,7 @@ dependencies = [ [[package]] name = "cmd" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -2311,7 +2311,7 @@ dependencies = [ [[package]] name = "common-base" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "anymap2", @@ -2332,7 +2332,7 @@ dependencies = [ [[package]] name = "common-batcher" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "dashmap", "futures", @@ -2341,14 +2341,14 @@ dependencies = [ [[package]] name = "common-catalog" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "const_format", ] [[package]] name = "common-config" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-base", "common-error", @@ -2371,7 +2371,7 @@ dependencies = [ [[package]] name = "common-datasource" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -2408,7 +2408,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "bigdecimal 0.4.8", "common-error", @@ -2421,7 +2421,7 @@ dependencies = [ [[package]] name = "common-error" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-macro", "http 1.5.0", @@ -2433,7 +2433,7 @@ dependencies = [ [[package]] name = "common-event-recorder" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -2458,7 +2458,7 @@ dependencies = [ [[package]] name = "common-frontend" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -2477,7 +2477,7 @@ dependencies = [ [[package]] name = "common-function" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -2539,7 +2539,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-trait", "common-runtime", @@ -2556,7 +2556,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow-flight", @@ -2591,7 +2591,7 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "common-catalog", @@ -2609,7 +2609,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "greptime-proto", "once_cell", @@ -2620,7 +2620,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "anyhow", "common-error", @@ -2636,7 +2636,7 @@ dependencies = [ [[package]] name = "common-memory-manager" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-error", "common-macro", @@ -2648,7 +2648,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "anymap2", "api", @@ -2719,7 +2719,7 @@ dependencies = [ [[package]] name = "common-options" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-base", "common-grpc", @@ -2730,11 +2730,11 @@ dependencies = [ [[package]] name = "common-plugins" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" [[package]] name = "common-pprof" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-error", "common-macro", @@ -2745,7 +2745,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-stream", @@ -2774,7 +2774,7 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-trait", "common-procedure", @@ -2783,7 +2783,7 @@ dependencies = [ [[package]] name = "common-query" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -2813,7 +2813,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arc-swap", "common-base", @@ -2838,7 +2838,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-trait", "catio", @@ -2866,7 +2866,7 @@ dependencies = [ [[package]] name = "common-session" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "serde", "strum 0.27.1", @@ -2874,7 +2874,7 @@ dependencies = [ [[package]] name = "common-sql" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arrow-schema 59.2.0", "common-base", @@ -2893,7 +2893,7 @@ dependencies = [ [[package]] name = "common-stat" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-base", "common-runtime", @@ -2907,7 +2907,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "backtrace", "common-base", @@ -2937,7 +2937,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-grpc", "common-query", @@ -2949,7 +2949,7 @@ dependencies = [ [[package]] name = "common-time" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arrow 59.2.0", "chrono", @@ -2967,7 +2967,7 @@ dependencies = [ [[package]] name = "common-version" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "cargo-manifest", "const_format", @@ -2977,7 +2977,7 @@ dependencies = [ [[package]] name = "common-wal" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-base", "common-error", @@ -3000,7 +3000,7 @@ dependencies = [ [[package]] name = "common-workload" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "common-telemetry", "serde", @@ -4473,7 +4473,7 @@ dependencies = [ [[package]] name = "datanode" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow-flight", @@ -4546,7 +4546,7 @@ checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" [[package]] name = "datatypes" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arrow 59.2.0", "arrow-array 59.2.0", @@ -5278,7 +5278,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "file-engine" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -5419,7 +5419,7 @@ checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" [[package]] name = "flow" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow 59.2.0", @@ -5471,7 +5471,7 @@ dependencies = [ "snafu 0.8.6", "sql", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.3.0-beta.1", "table", "tokio", "tokio-stream", @@ -5568,7 +5568,7 @@ dependencies = [ [[package]] name = "frontend" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow-flight", @@ -6946,7 +6946,7 @@ dependencies = [ [[package]] name = "index" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "async-trait", @@ -8097,7 +8097,7 @@ checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "log-query" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "chrono", "common-error", @@ -8109,7 +8109,7 @@ dependencies = [ [[package]] name = "log-store" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-stream", "async-trait", @@ -8461,7 +8461,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -8490,7 +8490,7 @@ dependencies = [ [[package]] name = "meta-srv" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -8581,7 +8581,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "aquamarine", @@ -8681,7 +8681,7 @@ dependencies = [ [[package]] name = "mito-codec" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "bytes", @@ -8706,7 +8706,7 @@ dependencies = [ [[package]] name = "mito2" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "aquamarine", @@ -9549,7 +9549,7 @@ dependencies = [ [[package]] name = "object-store" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "anyhow", "bytes", @@ -10094,7 +10094,7 @@ dependencies = [ [[package]] name = "operator" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -10156,7 +10156,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.3.0-beta.1", "table", "tokio", "tokio-util", @@ -10490,7 +10490,7 @@ checksum = "e3c406c9e2aa74554e662d2c2ee11cd3e73756988800be7e6f5eddb16fed4699" [[package]] name = "partition" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "common-base", @@ -10884,7 +10884,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -11039,7 +11039,7 @@ dependencies = [ [[package]] name = "plugins" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "auth", "catalog", @@ -11382,7 +11382,7 @@ dependencies = [ [[package]] name = "promql" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "async-trait", @@ -11799,7 +11799,7 @@ dependencies = [ [[package]] name = "puffin" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-compression", "async-trait", @@ -11851,7 +11851,7 @@ dependencies = [ [[package]] name = "query" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -11918,7 +11918,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.3.0-beta.1", "table", "tokio", "toml 0.8.23", @@ -13533,7 +13533,7 @@ dependencies = [ [[package]] name = "servers" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -13675,7 +13675,7 @@ dependencies = [ [[package]] name = "session" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "ahash 0.8.12", "api", @@ -14035,7 +14035,7 @@ dependencies = [ [[package]] name = "sql" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow-buffer 59.2.0", @@ -14089,7 +14089,7 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-trait", "clap", @@ -14368,7 +14368,7 @@ dependencies = [ [[package]] name = "standalone" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "async-trait", @@ -14418,7 +14418,7 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "store-api" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "aquamarine", @@ -14610,7 +14610,7 @@ dependencies = [ [[package]] name = "substrait" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "async-trait", "bytes", @@ -14746,7 +14746,7 @@ dependencies = [ [[package]] name = "table" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arc-swap", @@ -15026,7 +15026,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tests-fuzz" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "arbitrary", "async-trait", @@ -15074,7 +15074,7 @@ dependencies = [ [[package]] name = "tests-integration" -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" dependencies = [ "api", "arrow-flight", @@ -15153,7 +15153,7 @@ dependencies = [ "sqlx", "standalone", "store-api", - "substrait 1.3.0-alpha.1", + "substrait 1.3.0-beta.1", "table", "tempfile", "time", diff --git a/Cargo.toml b/Cargo.toml index 4f842313d8c..f6ca2d66562 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,7 +76,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.3.0-alpha.1" +version = "1.3.0-beta.1" edition = "2024" license = "Apache-2.0" From 886e02f0dcbf34498bbda099f363a6aa4bb5b98f Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:32:36 +0000 Subject: [PATCH 071/133] ci(query-regression): bump RUNNER_IMAGE_EPOCH to 7 (#9395) ci(query-regression): bump RUNNER_IMAGE_EPOCH to 7 for image m-0xihbfzm5xpbxolybo6i Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- .github/workflows/query-regression.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 47daa3f0435..6a368c455d6 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -549,7 +549,7 @@ jobs: readonly EXPECTED_SCCACHE_DIR="/home/runner/.cache/sccache" readonly EXPECTED_RUSTC_WRAPPER="/usr/local/bin/sccache" readonly RUNNER_IMAGE_DIGEST="sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11" - readonly RUNNER_IMAGE_EPOCH="6" + readonly RUNNER_IMAGE_EPOCH="7" require_expected_root() { local name="$1" From d6474b961cf7f7cd34479ae3338f88b6b557383c Mon Sep 17 00:00:00 2001 From: DeviousCardi <115358213+DeviousCardi@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:38:43 +0000 Subject: [PATCH 072/133] fix(promql): apply offset to subquery evaluation window (#9364) * fix(promql): apply offset to subquery evaluation window `prom_subquery_expr_to_plan` destructured `SubqueryExpr` without reading `offset`, so `[range:step] offset ` planned exactly the same window as the un-offset form and silently returned data for the wrong time range. The plain vector/matrix-selector paths already threaded the offset through `selector_to_series_normalize_plan` and `RangeManipulate`. Shift the inner evaluation window back by the offset and pass the offset to the subquery's `RangeManipulate`, which maps the inner samples forward onto the evaluation timeline before bucketing them into ranges. This matches Prometheus, whose `evaluator.subqueryTimeRange` evaluates the inner expression over `(start - offset - range, end - offset]` and whose `evalSubquery` then hands the samples to the outer range-vector function as a `MatrixSelector` that still carries the subquery offset. An offset on the inner selector composes additively, as `subqueryTimes` documents. `RangeManipulate`'s protobuf message has no offset field and recovers it on decode from an immediately underlying `SeriesNormalize`. Since `RangeManipulate` is commutative in `dist_plan` and can be pushed below a `MergeScan`, insert that carrier node so the offset survives distributed planning instead of decoding as zero. Known divergence, unchanged by this commit: Prometheus anchors subquery step points on absolute epoch multiples of the step, while GreptimeDB anchors them on the evaluation start. The two agree whenever the offset is a multiple of the subquery step; the added sqlness cases stay within that range. Closes #9330 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarav * docs(promql): correct the subquery-offset rationale and pin the histogram shift Follow-up to e132c267, which stated two things inaccurately and left one path untested. No behaviour change beyond a conditional clone. The `SeriesNormalize` carrier was described as existing "so the offset survives distributed planning", implying standalone is unaffected. That is wrong. `local_offset` is read on the substrait decode path (`range_manipulate.rs`, `instant_manipulate.rs`), and GreptimeDB routes PromQL plans through `MergeScan`/substrait in standalone too -- the standalone results `promql/encode_substrait.result` and `precisions.result` both show `MergeScan [is_placeholder=false, remote_input=[...]]`. Deleting the carrier therefore empties subquery-with-offset results in standalone as well, while `count_over_time_subquery_with_offset` keeps passing because it asserts the pre-serialization plan. Say so at the node, so the next reader does not remove it believing standalone is safe. The carrier is also not semantically inert: `SeriesNormalizeStream::normalize` shifts native histogram `start_timestamp` payloads by `offset`. On the subquery path that runs on a computed inner result, on top of any offset the inner selector already applied, and nothing pinned it. It is consistent -- `RangeManipulate` moves the millisecond time index forward by the same `offset_ms` -- so the distance from a histogram's start timestamp to the sample carrying it, which reset/rate detection reads, is invariant. Add `subquery_offset_shifts_histogram_start_and_time_index_together`, which executes the subquery's node stack over native histogram samples for a zero and a non-zero inner offset. It cannot be a sqlness case: native histograms are a struct column with no SQL type or literal (`sql_data_type_to_concrete_data_type` rejects structs) and the sqlness runner speaks only MySQL/Postgres, so such rows only arrive over gRPC/remote-write v2. The divergence note blamed sub-step offsets. Measured, that condition is too narrow: an unaligned evaluation timestamp alone diverges, with no offset at all -- `sum_over_time(fine[20s:10s])` at t=57 samples 47s and 57s here against Prometheus's 40s and 50s. The real condition is that `(start - offset)` is not a multiple of the subquery step, and it is pre-existing: the committed `tql eval (359, 359, '1s') sum_over_time(metric_total[60s:10s])` case is already an instance of it (`359 - 60 + 10 = 309`). Reword the comment and the sqlness header accordingly. A sub-step offset stays accepted: `foo[20s:10s] offset 5s` is valid PromQL, erroring on it would be a regression and would not close the gap anyway. This repo has no PromQL-compatibility documentation page, so there is nowhere else to record it. Also stop cloning the series key columns when there is no offset, and stop the `offset 0s` sqlness comment implying its parser error is evidence about this fix -- the same message comes from a `u32` overflow in `promql-parser`'s shared duration check (`offset 9999999999d`). Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Aarav * refactor(promql): trim subquery-offset comments and tests per review Shorten the planner and sqlness comments to the semantic contract, build `SeriesDivide` with a clone instead of an `Option`, and drop the normalize.rs histogram test and the `offset 0s` parser case. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bvab9gpXYjKyLN8Mo6JWWL Signed-off-by: Aarav * fix(promql): record the subquery offset as the range fold offset The subquery's `RangeManipulate` now shifts the payload timestamps by the offset, so `predict_linear` over an offset subquery must recover the evaluation time with that offset instead of 0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Bvab9gpXYjKyLN8Mo6JWWL Signed-off-by: Aarav --------- Signed-off-by: Aarav Co-authored-by: Claude Opus 5 (1M context) --- src/query/src/promql/planner.rs | 46 ++++++-- src/query/src/promql/planner/test.rs | 22 ++++ .../standalone/common/promql/subquery.result | 100 ++++++++++++++++++ .../standalone/common/promql/subquery.sql | 46 ++++++++ 4 files changed, 207 insertions(+), 7 deletions(-) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index be0c1364489..ffa99a52d0d 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -403,18 +403,34 @@ impl PromPlanner { subquery_expr: &SubqueryExpr, ) -> Result { let SubqueryExpr { - expr, range, step, .. + expr, + range, + step, + offset, + .. } = subquery_expr; + // Prometheus evaluates the inner expression over `(start - offset - range, end - offset]` + // (`subqueryTimeRange`). Shift the inner window back here; `RangeManipulate` maps the + // samples forward again by the same offset. + let offset_ms = match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + }; + let current_interval = self.ctx.interval; if let Some(step) = step { self.ctx.interval = step.as_millis() as _; } let current_start = self.ctx.start; - self.ctx.start -= range.as_millis() as i64 - self.ctx.interval; + let current_end = self.ctx.end; + self.ctx.start -= offset_ms + range.as_millis() as i64 - self.ctx.interval; + self.ctx.end -= offset_ms; let input = self.prom_expr_to_plan(expr, query_engine_state).await?; self.ctx.interval = current_interval; self.ctx.start = current_start; + self.ctx.end = current_end; ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); let range_ms = range.as_millis() as _; @@ -471,26 +487,42 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)?; let divide_plan = LogicalPlan::Extension(Extension { node: Arc::new(SeriesDivide::new( - series_key_columns, + series_key_columns.clone(), time_index_column.clone(), sort_plan, )), }); + // `RangeManipulate` has no offset in its protobuf message; decoding recovers it from the + // `SeriesNormalize` directly below. Stale markers are not filtered: the input is computed. + let divide_plan = if offset_ms != 0 { + LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + offset_ms, + time_index_column.clone(), + false, + series_key_columns, + divide_plan, + )), + }) + } else { + divide_plan + }; + let manipulate = RangeManipulate::new( self.ctx.start, self.ctx.end, self.ctx.interval, - 0, + offset_ms, range_ms, time_index_column, self.ctx.field_columns.clone(), divide_plan, ) .context(DataFusionPlanningSnafu)?; - // A subquery always folds with offset 0, so its payload timestamps are already on the - // evaluation timeline a function above it reads; see [`Self::create_range_eval_ts_expr`]. - self.ctx.range_fold_offset = Some(0); + // The payload timestamps are shifted by the subquery offset; see + // [`Self::create_range_eval_ts_expr`]. + self.ctx.range_fold_offset = Some(offset_ms); Ok(LogicalPlan::Extension(Extension { node: Arc::new(manipulate), diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index a83a42bae53..0393930b033 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -5109,6 +5109,28 @@ async fn count_over_time_subquery() { indie_query_plan_compare(query, expected).await; } +/// `offset` on a subquery must shift the inner evaluation window back and be +/// carried into the outer range manipulation. See +/// . +#[tokio::test] +async fn count_over_time_subquery_with_offset() { + let query = "count_over_time(some_metric[10m:1m] offset 5m)"; + let expected = String::from( + "Filter: prom_count_over_time(timestamp_range,field_0) IS NOT NULL [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n Projection: some_metric.timestamp, prom_count_over_time(timestamp_range, field_0) AS prom_count_over_time(timestamp_range,field_0), some_metric.tag_0 [timestamp:Timestamp(ms), prom_count_over_time(timestamp_range,field_0):Float64;N, tag_0:Utf8]\ + \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=[\"field_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))]\ + \n PromSeriesNormalize: offset=[300000], time index=[timestamp], filter NaN: [false] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromInstantManipulate: range=[-840000..99700000], lookback=[1000], interval=[60000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.timestamp >= TimestampMillisecond(-840999, None) AND some_metric.timestamp <= TimestampMillisecond(99700000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + ); + indie_query_plan_compare(query, expected).await; +} + #[tokio::test] async fn test_hash_join() { let mut eval_stmt = EvalStmt { diff --git a/tests/cases/standalone/common/promql/subquery.result b/tests/cases/standalone/common/promql/subquery.result index 12e65c4310c..61f09258a41 100644 --- a/tests/cases/standalone/common/promql/subquery.result +++ b/tests/cases/standalone/common/promql/subquery.result @@ -63,3 +63,103 @@ drop table metric_total; Affected Rows: 0 +-- Offset on a subquery shifts the subquery's own evaluation window back by the offset. +-- Reference: Prometheus `evaluator.subqueryTimeRange` (promql/engine.go). +-- The offset cases stay on the subquery step grid so they match Prometheus directly. +create table subquery_offset_total ( + ts timestamp time index, + host string primary key, + val double, +); + +Affected Rows: 0 + +insert into subquery_offset_total values + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7); + +Affected Rows: 7 + +-- baseline: no offset at t=60 covers the 10s subquery points in (40s, 60s] -> 6 + 7 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 13.0 | a | ++---------------------+----------------------------------+------+ + +-- the same subquery evaluated at t=30 -> 3 + 4 +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- `offset 30s` at t=60 must equal the un-offset subquery at t=30 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- a negative offset looks ahead of the evaluation time +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s] offset -30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 13.0 | a | ++---------------------+----------------------------------+------+ + +-- an offset on the inner selector composes additively with the subquery offset: Prometheus +-- `subqueryTimes` accumulates "the sum of offsets and ranges of all subqueries in the path", +-- and the inner selector subtracts its own offset from the already shifted step timestamps. +-- 20s + 10s therefore behaves like the un-offset subquery at t=30. +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 10s)[20s:10s] offset 20s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- ... and the inner offset alone accounts for the same total shift +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 30s)[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +-- `predict_linear` predicts from the evaluation time: 4 + 0.1 * 30 = 7 and 7 - 0.1 * 30 = 4 +tql eval (60, 60, '1s') predict_linear(subquery_offset_total[20s:10s] offset 30s, 0); + ++---------------------+----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(0)) | host | ++---------------------+----------------------------------------------+------+ +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------------------+------+ + +tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30s, 0); + ++---------------------+----------------------------------------------+------+ +| ts | prom_predict_linear(ts_range,val,Float64(0)) | host | ++---------------------+----------------------------------------------+------+ +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+----------------------------------------------+------+ + +drop table subquery_offset_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/subquery.sql b/tests/cases/standalone/common/promql/subquery.sql index 95215b8488b..598aa91a06e 100644 --- a/tests/cases/standalone/common/promql/subquery.sql +++ b/tests/cases/standalone/common/promql/subquery.sql @@ -20,3 +20,49 @@ tql eval (10, 10, '1s') rate(metric_total[20s:10s]); tql eval (20, 20, '1s') rate(metric_total[20s:5s]); drop table metric_total; + +-- Offset on a subquery shifts the subquery's own evaluation window back by the offset. +-- Reference: Prometheus `evaluator.subqueryTimeRange` (promql/engine.go). +-- The offset cases stay on the subquery step grid so they match Prometheus directly. +create table subquery_offset_total ( + ts timestamp time index, + host string primary key, + val double, +); + +insert into subquery_offset_total values + (0, 'a', 1), + (10000, 'a', 2), + (20000, 'a', 3), + (30000, 'a', 4), + (40000, 'a', 5), + (50000, 'a', 6), + (60000, 'a', 7); + +-- baseline: no offset at t=60 covers the 10s subquery points in (40s, 60s] -> 6 + 7 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s]); + +-- the same subquery evaluated at t=30 -> 3 + 4 +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s]); + +-- `offset 30s` at t=60 must equal the un-offset subquery at t=30 +tql eval (60, 60, '1s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + +-- a negative offset looks ahead of the evaluation time +tql eval (30, 30, '1s') sum_over_time(subquery_offset_total[20s:10s] offset -30s); + +-- an offset on the inner selector composes additively with the subquery offset: Prometheus +-- `subqueryTimes` accumulates "the sum of offsets and ranges of all subqueries in the path", +-- and the inner selector subtracts its own offset from the already shifted step timestamps. +-- 20s + 10s therefore behaves like the un-offset subquery at t=30. +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 10s)[20s:10s] offset 20s); + +-- ... and the inner offset alone accounts for the same total shift +tql eval (60, 60, '1s') sum_over_time((subquery_offset_total offset 30s)[20s:10s]); + +-- `predict_linear` predicts from the evaluation time: 4 + 0.1 * 30 = 7 and 7 - 0.1 * 30 = 4 +tql eval (60, 60, '1s') predict_linear(subquery_offset_total[20s:10s] offset 30s, 0); + +tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30s, 0); + +drop table subquery_offset_total; From 5ccfcd4644fd93d49b138a995f47f257864c292e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 29 Sep 2026 06:47:06 +0000 Subject: [PATCH 073/133] feat: support automatic column addition for Flight bulk inserts (#9285) * fix: auto-add columns when initializing bulk insert streams Signed-off-by: Lei, HUANG * fix: reject nested columns in bulk schema auto-add Signed-off-by: Lei, HUANG * fix: reject unknown columns in non-empty bulk inserts Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/frontend/AGENTS.md | 4 + src/frontend/src/instance/grpc.rs | 43 +-- src/operator/src/bulk_insert.rs | 18 ++ src/operator/src/insert.rs | 81 +++++- src/servers/src/grpc/flight.rs | 9 +- tests-integration/src/grpc/flight.rs | 375 ++++++++++++++++++++++++--- 6 files changed, 483 insertions(+), 47 deletions(-) diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 49a4b8fcd0f..20c7bc323a5 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -51,6 +51,10 @@ remote datanodes via `operator`/`client`. - Internal gRPC listeners mark requests with `Channel::Internal` in middleware (`server.rs`), including requests handled by Enterprise Flight wrappers. +- **Flight bulk insert** (`instance/grpc.rs`): initializes on the first batch after + the lazy schema handshake; checks permissions and reconciles missing columns + through `Inserter` once per stream, then reuses the refreshed table. + - **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. diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index e2435b65689..5cb546485bc 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -392,6 +392,7 @@ impl Instance { let catalog_manager = self.catalog_manager().clone(); let plugins = self.plugins.clone(); let inserter = self.inserter.clone(); + let statement_executor = self.statement_executor.clone(); let ctx = ctx.clone(); let mut table_ref: Option = None; let mut table_checked = false; @@ -422,25 +423,37 @@ impl Instance { .context(PermissionSnafu)?; // Resolve table reference - table_ref = Some( - catalog_manager - .table( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - None, - ) - .await - .context(CatalogSnafu)? - .with_context(|| TableNotFoundSnafu { - table_name: table_name.to_string(), - })?, - ); + let table = catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)? + .with_context(|| TableNotFoundSnafu { + table_name: table_name.to_string(), + })?; // Check permissions for the table let interceptor_ref = plugins.get::>(); let interceptor = interceptor_ref.as_ref(); - interceptor.pre_bulk_insert(table_ref.clone().unwrap(), ctx.clone())?; + interceptor.pre_bulk_insert(table.clone(), ctx.clone())?; + + // The client sends the schema lazily after the stream handshake. + // Reconcile it once, before writing the first batch. + table_ref = Some( + inserter + .ensure_bulk_insert_schema( + table, + request.record_batch.schema().as_ref(), + &ctx, + statement_executor.as_ref(), + ) + .await + .context(TableOperationSnafu)?, + ); table_checked = true; } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index f1f57a64479..f8d58f5d173 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -140,6 +140,24 @@ impl Inserter { return Ok(0); } + // Bulk storage paths may discard unknown columns instead of rejecting them. + for field in record_batch.schema_ref().fields() { + ensure!( + table_info + .meta + .schema + .column_schema_by_name(field.name()) + .is_some(), + error::InvalidInsertRequestSnafu { + reason: format!( + "Column '{}' not found in table '{}'", + field.name(), + table_info.full_table_name() + ), + } + ); + } + // The zero value is WCU, not bytes. Bulk writes have no WCU accounting; // preserve that behavior while admitting their rows before dispatch. write_meter!(MeterRecord::new( diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index a0ebe236767..e1fbe39a9f4 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::v1::alter_table_expr::Kind; -use api::v1::column_def::options_from_skipping; +use api::v1::column_def::{options_from_skipping, try_as_column_def}; use api::v1::region::{ InsertRequest as RegionInsertRequest, InsertRequests as RegionInsertRequests, RegionRequestHeader, @@ -27,6 +27,7 @@ use api::v1::{ AlterTableExpr, ColumnDataType, ColumnSchema, CreateTableExpr, InsertRequests, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; +use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ @@ -906,6 +907,84 @@ impl Inserter { ) } + /// Adds missing columns from a bulk stream's schema and returns the refreshed table. + /// Call once when initializing the stream, before writing its first batch. + /// Does not infer new nested or dictionary columns. + pub async fn ensure_bulk_insert_schema( + &self, + table: TableRef, + request_schema: &ArrowSchema, + ctx: &QueryContextRef, + statement_executor: &StatementExecutor, + ) -> Result { + let table_info = table.table_info(); + if self.auto_create_disabled_reason(ctx)?.is_some() + && !Self::is_auto_create_exempt_private_table(&table_info.schema_name, &table_info.name) + { + return Ok(table); + } + + let table_schema = table.schema(); + let schema = request_schema + .fields() + .iter() + .filter(|field| table_schema.column_schema_by_name(field.name()).is_none()) + .map(|field| { + let data_type = field.data_type(); + // Dictionary values can reach the same infallible child-type conversion + // as nested types, even when Arrow's is_nested() returns false. + ensure!( + !data_type.is_nested() && !matches!(data_type, ArrowDataType::Dictionary(..)), + crate::error::NotSupportedSnafu { + feat: format!( + "automatically adding bulk insert column '{}' with type {:?}", + field.name(), + data_type + ), + } + ); + let column = datatypes::schema::ColumnSchema::try_from(field.as_ref()) + .context(crate::error::ConvertSchemaSnafu)?; + // Arrow fields do not carry primary-key semantics. New columns are + // fields, unless explicitly marked as a time index. + let column_def = + try_as_column_def(&column, false).context(crate::error::ColumnDataTypeSnafu)?; + Ok(ColumnSchema { + column_name: column_def.name, + datatype: column_def.data_type, + semantic_type: column_def.semantic_type, + datatype_extension: column_def.datatype_extension, + options: column_def.options, + }) + }) + .collect::>>()?; + let mut request = RowInsertRequest { + table_name: table_info.name.clone(), + rows: Some(Rows { + schema, + rows: Vec::new(), + }), + }; + let Some(alter_expr) = + self.get_alter_table_expr_on_demand(&mut request, &table, ctx, false, false, true)? + else { + return Ok(table); + }; + + statement_executor + .alter_table_inner(alter_expr, ctx.clone(), TriggerReason::AutoAlter) + .await?; + self.get_table( + &table_info.catalog_name, + &table_info.schema_name, + &table_info.name, + ) + .await? + .with_context(|| TableNotFoundSnafu { + table_name: table_info.full_table_name(), + }) + } + /// Ensures a trace table has the request-global schema without requiring a /// padded data row to drive on-demand creation or alteration. When /// `alter_existing` is false, a table created after planning is left for the diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 874e98b7326..ca05ea91844 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -46,6 +46,7 @@ use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; +use table::requests::AUTO_CREATE_TABLE_KEY; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -254,8 +255,14 @@ impl FlightCraft for GreptimeRequestHandler { let limiter = extensions.get::().cloned(); - let query_ctx = + let mut query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?; + // Bulk streams use the same schema-on-write switch as row inserts. + for (key, value) in hint_headers::extract_hints(&headers) { + if key == AUTO_CREATE_TABLE_KEY { + Arc::make_mut(&mut query_ctx).set_extension(key, value); + } + } context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?; const MAX_PENDING_RESPONSES: usize = 32; diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 2cd0c44d504..ef3df79f558 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -47,6 +47,9 @@ mod test { RecordBatch, RecordBatchStreamWrapper, RecordBatches, SendableRecordBatchStream, }; use common_telemetry::tracing_context::TracingContext; + use datatypes::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit, + }; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; @@ -66,14 +69,17 @@ mod test { use session::context::QueryContextRef; use session::hints::INSERT_SKIP_WAL_HINT; use tokio::net::TcpListener; - use tokio_stream::wrappers::TcpListenerStream; + use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream}; use tonic::transport::Server as TonicServer; use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; - use crate::test_util::{MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server}; + use crate::test_util::{ + MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server, + setup_grpc_server_with_auto_create_table_disabled, + }; use crate::tests::test_util::MockInstance; struct SlowFlightCraft; @@ -467,6 +473,315 @@ mod test { db.shutdown().await; } + #[rstest] + #[case::standalone_single_region(false, false)] + #[case::standalone_partitioned(false, true)] + #[case::distributed_single_region(true, false)] + #[case::distributed_partitioned(true, true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns( + #[case] distributed: bool, + #[case] partitioned: bool, + ) { + let mut db = MockInstanceImpl::new( + &format!("flight_bulk_auto_add_columns_{partitioned}"), + distributed, + ) + .await; + let runtime = common_runtime::global_runtime().clone(); + let handler = GreptimeRequestHandler::new( + db.frontend(), + None, + Some(runtime.clone()), + FlightCompression::default(), + ); + let mut server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime) + .flight_handler(Arc::new(handler)) + .build(); + server.start("127.0.0.1:0".parse().unwrap()).await.unwrap(); + let client = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + let partition = if partitioned { + " PARTITION ON COLUMNS (a) (a < 0, a >= 0)" + } else { + "" + }; + client.sql(&format!( + "CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a)){partition}" + )).await.unwrap(); + // Existing rows require the inferred column to be nullable. + client + .sql("INSERT INTO foo VALUES (100, 100)") + .await + .unwrap(); + test_put_record_batches(&client, create_record_batches(-4)).await; + + query_and_expect( + db.frontend().as_ref(), + "SELECT a, \"B\" FROM foo ORDER BY a", + "\ ++-----+-----+ +| a | B | ++-----+-----+ +| -4 | s4 | +| -3 | s3 | +| -2 | s2 | +| -1 | s1 | +| 0 | s0 | +| 1 | s-1 | +| 2 | s-2 | +| 3 | s-3 | +| 4 | s-4 | +| 100 | | ++-----+-----+", + ) + .await; + server.shutdown().await.unwrap(); + db.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_once_per_stream() { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_auto_add_columns_once_per_stream", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batches = create_record_batches(1); + let schema = batches[0].schema.arrow_schema().clone(); + let (tx, rx) = tokio::sync::mpsc::channel(4); + // The client must be able to receive the handshake before sending its schema. + let mut responses = client + .do_put(ReceiverStream::new(rx).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + tx.send(encode_put_schema(schema.as_ref())).await.unwrap(); + + let empty_batch = common_recordbatch::DfRecordBatch::new_empty(schema); + for round in 0..2 { + for message in + FlightEncoder::default().encode(FlightMessage::RecordBatch(empty_batch.clone())) + { + tx.send(message).await.unwrap(); + } + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + if round == 0 { + // Even an empty first batch initializes the schema. Dropping the + // column makes a repeated reconciliation on the next batch observable. + client + .sql("ALTER TABLE foo DROP COLUMN \"B\"") + .await + .unwrap(); + } + } + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + drop(tx); + assert!(responses.next().await.is_none()); + + // A new stream reconciles again and persists the new column's values. + test_put_record_batches(&client, batches).await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(\"B\") AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 9 | ++---+", + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::hint(false)] + #[case::config(true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_auto_add_columns_disabled(#[case] disabled_by_config: bool) { + let name = "test_flight_bulk_auto_add_columns_disabled"; + let (db, server) = if disabled_by_config { + setup_grpc_server_with_auto_create_table_disabled(StorageType::File, name).await + } else { + setup_grpc_server(StorageType::File, name).await + }; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, PRIMARY KEY (a))") + .await + .unwrap(); + let batch = create_record_batches(1).remove(0).into_df_record_batch(); + let schema = batch.schema(); + let mut messages = vec![encode_put_schema(schema.as_ref())]; + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + messages.extend(FlightEncoder::default().encode(FlightMessage::RecordBatch(batch))); + // A request hint cannot override the server-side setting. + let hint = if disabled_by_config { "true" } else { "false" }; + let mut responses = client + .do_put_with_hints( + tokio_stream::iter(messages).boxed(), + &[("auto_create_table", hint)], + ) + .await + .unwrap(); + // The handshake and empty batch succeed, but the non-empty write must fail. + for _ in 0..2 { + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + } + let Some(Err(err)) = responses.next().await else { + panic!("expected the bulk write with an unknown column to fail"); + }; + assert_eq!(err.status_code(), StatusCode::InvalidArguments); + let message = err.to_string(); + assert!(message.contains("Column 'B' not found"), "{message}"); + assert!(responses.next().await.is_none()); + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| a | +| ts | ++-------------+", + ) + .await; + query_and_expect( + db.frontend().as_ref(), + "SELECT count(*) AS n FROM foo", + "\ ++---+ +| n | ++---+ +| 0 | ++---+", + ) + .await; + test_put_record_batches_with_hints( + &client, + create_record_batches_without_nullable_column(1), + &[("auto_create_table", hint)], + ) + .await; + server.shutdown().await.unwrap(); + } + + #[rstest] + #[case::list(ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Int32, true))))] + #[case::list_unsupported_child(ArrowDataType::List(Arc::new(Field::new( + "item", + ArrowDataType::FixedSizeBinary(16), + true, + ))))] + #[case::struct_unsupported_child(ArrowDataType::Struct( + vec![Field::new("item", ArrowDataType::FixedSizeBinary(16), true)].into(), + ))] + #[case::dictionary(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::Utf8), + ))] + #[case::dictionary_unsupported_value(ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int32), + Box::new(ArrowDataType::FixedSizeBinary(16)), + ))] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_rejects_nested_auto_add_without_altering_table( + #[case] data_type: ArrowDataType, + ) { + let (db, server) = setup_grpc_server( + StorageType::File, + "test_flight_bulk_rejects_nested_auto_add_without_altering_table", + ) + .await; + let client = Database::new_with_dbname( + "greptime-public", + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + client + .sql("CREATE TABLE foo (ts TIMESTAMP TIME INDEX)") + .await + .unwrap(); + + // Use raw Arrow fields: Greptime's schema conversion itself used to panic + // on unsupported child types. A preceding scalar must not be partially added. + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("new_scalar", ArrowDataType::Int32, true), + Field::new("new_nested", data_type.clone(), true), + ])); + let mut encoder = FlightEncoder::default(); + let mut schema_data = encoder.encode_schema(schema.as_ref()); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + let mut messages = vec![schema_data]; + // An empty first batch still performs schema reconciliation. + messages.extend(encoder.encode(FlightMessage::RecordBatch( + common_recordbatch::DfRecordBatch::new_empty(schema), + ))); + let mut responses = client + .do_put(tokio_stream::iter(messages).boxed()) + .await + .unwrap(); + assert_eq!(responses.next().await.unwrap().unwrap().affected_rows(), 0); + let Some(Err(err)) = responses.next().await else { + panic!("expected an unsupported bulk schema error"); + }; + assert_eq!(err.status_code(), StatusCode::Unsupported); + let message = err.to_string(); + assert!(message.contains("new_nested"), "{message}"); + assert!(message.contains(&format!("{data_type:?}")), "{message}"); + assert!(responses.next().await.is_none()); + + query_and_expect( + db.frontend().as_ref(), + "SELECT column_name FROM information_schema.columns WHERE table_name = 'foo' ORDER BY column_name", + "\ ++-------------+ +| column_name | ++-------------+ +| ts | ++-------------+", + ) + .await; + server.shutdown().await.unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -816,34 +1131,23 @@ mod test { let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone(); - let stream = futures::stream::once(async move { - let mut schema_data = FlightEncoder::default().encode_schema(schema.as_ref()); - let metadata = DoPutMetadata::new(0); - schema_data.app_metadata = serde_json::to_vec(&metadata).unwrap().into(); - // first message in "DoPut" stream should carry table name in flight descriptor - schema_data.flight_descriptor = Some(FlightDescriptor { - r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, - path: vec!["foo".to_string()], - ..Default::default() - }); - schema_data - }) - .chain( - tokio_stream::iter(record_batches) - .enumerate() - .flat_map(|(i, x)| { - let mut encoder = FlightEncoder::default(); - let message = FlightMessage::RecordBatch(x.into_df_record_batch()); - let mut data = encoder.encode(message); - let metadata = DoPutMetadata::new((i + 1) as i64); - data.iter_mut().for_each(|x| { - x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() - }); - tokio_stream::iter(data) - }) - .boxed(), - ) - .boxed(); + let stream = futures::stream::once(async move { encode_put_schema(schema.as_ref()) }) + .chain( + tokio_stream::iter(record_batches) + .enumerate() + .flat_map(|(i, x)| { + let mut encoder = FlightEncoder::default(); + let message = FlightMessage::RecordBatch(x.into_df_record_batch()); + let mut data = encoder.encode(message); + let metadata = DoPutMetadata::new((i + 1) as i64); + data.iter_mut().for_each(|x| { + x.app_metadata = serde_json::to_vec(&metadata).unwrap().into() + }); + tokio_stream::iter(data) + }) + .boxed(), + ) + .boxed(); let response_stream = client.do_put_with_hints(stream, hints).await.unwrap(); @@ -863,6 +1167,17 @@ mod test { assert_eq!(requests_count + 1, responses_count); } + fn encode_put_schema(schema: &datatypes::arrow::datatypes::Schema) -> FlightData { + let mut schema_data = FlightEncoder::default().encode_schema(schema); + schema_data.app_metadata = serde_json::to_vec(&DoPutMetadata::new(0)).unwrap().into(); + schema_data.flight_descriptor = Some(FlightDescriptor { + r#type: arrow_flight::flight_descriptor::DescriptorType::Path as i32, + path: vec!["foo".to_string()], + ..Default::default() + }); + schema_data + } + fn create_record_batches_without_nullable_column(start: i64) -> Vec { let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( From a7f99b34250123a1fd04cf64e19eb1643ccfe8fe Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 29 Sep 2026 07:48:29 +0000 Subject: [PATCH 074/133] ci: retry failed nightly release on following weekdays (#9394) Previously the nightly release was scheduled only on Mondays, so a failed run left users without nightly builds for a whole week. Now the schedule triggers every weekday at 00:00 UTC, but the release only proceeds when the latest published nightly release is older than NIGHTLY_RELEASE_MAX_AGE_DAYS (5) days. The check runs in allocate-runners before any EC2 runner is allocated, and gates all schedule-driven jobs (including the Slack notification, so a skipped nightly is silent). Tag pushes and manual dispatches are unaffected. Signed-off-by: Ning Sun --- .github/scripts/check-nightly-release.sh | 68 ++++++++++++++++++++++++ .github/workflows/release.yml | 57 +++++++++++++++----- 2 files changed, 111 insertions(+), 14 deletions(-) create mode 100755 .github/scripts/check-nightly-release.sh diff --git a/.github/scripts/check-nightly-release.sh b/.github/scripts/check-nightly-release.sh new file mode 100755 index 00000000000..82eb910d6fd --- /dev/null +++ b/.github/scripts/check-nightly-release.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +set -euo pipefail + +# Returns whether the current release run should proceed, 'true' or 'false'. +# +# The nightly schedule runs every weekday, but a scheduled release only proceeds +# when the latest published nightly release is older than +# NIGHTLY_RELEASE_MAX_AGE_DAYS days. This way, a failed nightly run is retried +# by the next weekday schedule, instead of leaving a whole week without nightly +# builds. Non-scheduled runs (tag pushes and manual dispatches) always proceed. +# +# Only the result is written to stdout, so that it can be used in GitHub Actions +# outputs; all diagnostics go to stderr. +# +# You can run as following examples: +# GITHUB_EVENT_NAME=push ./check-nightly-release.sh +# GITHUB_EVENT_NAME=schedule GITHUB_REPOSITORY=GreptimeTeam/greptimedb GH_TOKEN=ghp_xxx ./check-nightly-release.sh +# GITHUB_EVENT_NAME=schedule NIGHTLY_RELEASE_MAX_AGE_DAYS=5 GITHUB_REPOSITORY=GreptimeTeam/greptimedb GH_TOKEN=ghp_xxx ./check-nightly-release.sh +if [ -z "$GITHUB_EVENT_NAME" ]; then + echo "GITHUB_EVENT_NAME is empty" >&2 + exit 1 +fi + +if [ "$GITHUB_EVENT_NAME" != "schedule" ]; then + echo "Not a scheduled run, the nightly check is skipped." >&2 + echo "true" + exit 0 +fi + +MAX_AGE_DAYS="${NIGHTLY_RELEASE_MAX_AGE_DAYS:-5}" +if ! [[ "$MAX_AGE_DAYS" =~ ^[0-9]+$ ]]; then + echo "Invalid NIGHTLY_RELEASE_MAX_AGE_DAYS: ${MAX_AGE_DAYS}" >&2 + exit 1 +fi + +SECONDS_PER_DAY=86400 + +# The publish time of the latest published nightly release, like the tag +# 'v0.2.0-nightly-20230313'. Fail open on GitHub API errors: proceeding with +# the release is safer than skipping it. +if ! latest_nightly_published_at="$(gh api "repos/${GITHUB_REPOSITORY}/releases?per_page=100" \ + --jq '[.[] | select(.draft == false and (.tag_name | contains("nightly")))][0].published_at // empty')"; then + echo "::warning::Failed to query the latest nightly release from the GitHub API, will proceed with the nightly release." >&2 + echo "true" + exit 0 +fi + +if [ -z "$latest_nightly_published_at" ]; then + echo "No published nightly release found, will proceed with the nightly release." >&2 + echo "true" + exit 0 +fi + +published_at_epoch="$(date --date="$latest_nightly_published_at" +%s)" +now_epoch="$(date +%s)" +age_seconds=$(( now_epoch - published_at_epoch )) +age_days=$(( age_seconds / SECONDS_PER_DAY )) + +echo "The latest nightly release was published at ${latest_nightly_published_at} (${age_days} days ago)." >&2 + +if [ "$age_seconds" -gt $(( MAX_AGE_DAYS * SECONDS_PER_DAY )) ]; then + echo "The latest nightly release is older than ${MAX_AGE_DAYS} days, will proceed with the nightly release." >&2 + echo "true" +else + echo "The latest nightly release is not older than ${MAX_AGE_DAYS} days, will skip the nightly release." >&2 + echo "false" +fi diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6c0c36c895e..5a689910aee 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -4,15 +4,20 @@ name: Release # 1. The tag('v*.*.*') push release: the release workflow will be triggered by the tag push event. # 2. The scheduled release(the version will be '-nightly-YYYYMMDD', any pre-release # extension in the next release version like '-alpha.1' is stripped, so 'nightly' is the extension): -# the release workflow will be triggered by the schedule event. +# the release workflow will be triggered by the schedule event. The schedule runs every weekday, +# but the release only proceeds when the latest published nightly release is older than +# NIGHTLY_RELEASE_MAX_AGE_DAYS days, so a failed nightly run is retried by the next weekday +# schedule instead of leaving a whole week without nightly builds. # A failed tag push release can be manually dispatched on the same tag to recover it without changing the formal release version. on: push: tags: - "v*.*.*" schedule: - # At 00:00 on Monday. - - cron: '0 0 * * 1' + # At 00:00 on every weekday (Monday to Friday). The scheduled nightly release only + # proceeds when the latest published nightly release is older than + # NIGHTLY_RELEASE_MAX_AGE_DAYS days, see the 'Check nightly release necessity' step. + - cron: '0 0 * * 1-5' workflow_dispatch: # Allows you to run this workflow manually. # Notes: The GitHub Actions ONLY support 10 inputs; this final input is the tenth. inputs: @@ -100,6 +105,10 @@ env: # the pre-release extension of the next release version is stripped, e.g. v0.2.0-alpha.1 -> v0.2.0-nightly-20230313. NIGHTLY_RELEASE_PREFIX: nightly + # The maximum age (in days) of the latest published nightly release before a scheduled + # nightly release may run again; see the 'Check nightly release necessity' step. + NIGHTLY_RELEASE_MAX_AGE_DAYS: 5 + jobs: allocate-runners: name: Allocate runners @@ -123,6 +132,10 @@ jobs: # The 'is-current-version-latest' determines whether to update 'latest' Docker tags and downstream repositories. is-current-version-latest: ${{ steps.check-version.outputs.is-current-version-latest }} is-current-version-stable: ${{ steps.check-version.outputs.is-current-version-stable }} + + # Whether a scheduled nightly release should run: it only runs when the latest published + # nightly release is older than NIGHTLY_RELEASE_MAX_AGE_DAYS days. Always 'true' for other events. + nightly-required: ${{ steps.check-nightly.outputs.nightly-required }} steps: - name: Checkout uses: actions/checkout@v4 @@ -135,6 +148,19 @@ jobs: run: | ./scripts/check-builder-rust-version.sh + # A scheduled run only proceeds when the latest published nightly release is older than + # NIGHTLY_RELEASE_MAX_AGE_DAYS days, so a failed nightly run is retried by the next weekday + # schedule. The check runs before the runners are allocated to avoid wasting EC2 resources. + - name: Check nightly release necessity + id: check-nightly + run: | + echo "nightly-required=$(./.github/scripts/check-nightly-release.sh)" >> "$GITHUB_OUTPUT" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_EVENT_NAME: ${{ github.event_name }} + GITHUB_REPOSITORY: ${{ github.repository }} + NIGHTLY_RELEASE_MAX_AGE_DAYS: ${{ env.NIGHTLY_RELEASE_MAX_AGE_DAYS }} + # The create-version will create a global variable named 'version' in the global workflows. # - If it's a tag push release or manual dispatch on a tag with REUSE_EXISTING_RELEASE_TAG=true, the version is the tag name(${{ github.ref_name }}); # - If it's a scheduled release, the version is '-nightly-$buildTime', like v0.2.0-nightly-20230313 (pre-release extensions like '-alpha.1' are stripped); @@ -156,7 +182,7 @@ jobs: ./.github/scripts/check-version.sh "${{ steps.create-version.outputs.version }}" - name: Allocate linux-amd64 runner - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-amd64-runner with: @@ -170,7 +196,7 @@ jobs: subnet-id: ${{ vars.EC2_RUNNER_SUBNET_ID }} - name: Allocate linux-arm64 runner - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && steps.check-nightly.outputs.nightly-required == 'true' }} uses: ./.github/actions/start-runner id: start-linux-arm64-runner with: @@ -185,7 +211,7 @@ jobs: build-linux-amd64-artifacts: name: Build linux-amd64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -211,7 +237,7 @@ jobs: build-linux-arm64-artifacts: name: Build linux-arm64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -238,7 +264,7 @@ jobs: build-linux-riscv64-artifacts: name: Build linux-riscv64 artifacts - if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -302,7 +328,7 @@ jobs: needs: [ allocate-runners, ] - if: ${{ inputs.build_macos_artifacts || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_macos_artifacts || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} steps: - uses: actions/checkout@v4 with: @@ -340,7 +366,7 @@ jobs: needs: [ allocate-runners, ] - if: ${{ inputs.build_windows_artifacts || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.build_windows_artifacts || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} steps: - run: git config --global core.autocrlf false @@ -365,7 +391,7 @@ jobs: prepare-release-validation: name: Prepare release validation - if: ${{ inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ allocate-runners, ] @@ -528,6 +554,7 @@ jobs: if: | always() && (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && needs.allocate-runners.result == 'success' && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && @@ -571,7 +598,7 @@ jobs: release-cn-artifacts: name: Release artifacts to CN region - if: ${{ inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + if: ${{ (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} needs: [ # The job have to wait for all the artifacts are built. allocate-runners, build-linux-amd64-artifacts, @@ -618,6 +645,7 @@ jobs: if: | always() && (inputs.publish_github_release || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && needs.allocate-runners.result == 'success' && needs.prepare-release-validation.result == 'success' && (needs.compat-release.result == 'success' || @@ -726,7 +754,7 @@ jobs: bump-downstream-repo-versions: name: Bump downstream repo versions - if: ${{ github.event_name == 'schedule' || ((github.event_name == 'push' || github.event_name == 'workflow_dispatch') && github.ref_type == 'tag' && !contains(github.ref_name, 'nightly')) && (needs.allocate-runners.outputs.is-current-version-stable == 'true' && needs.allocate-runners.outputs.is-current-version-latest == 'true') }} + if: ${{ (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true') || ((github.event_name == 'push' || github.event_name == 'workflow_dispatch') && github.ref_type == 'tag' && !contains(github.ref_name, 'nightly')) && (needs.allocate-runners.outputs.is-current-version-stable == 'true' && needs.allocate-runners.outputs.is-current-version-latest == 'true') }} needs: [allocate-runners, publish-github-release] runs-on: ubuntu-latest # Permission reference: https://docs.github.com/en/actions/using-jobs/assigning-permissions-to-jobs @@ -793,9 +821,10 @@ jobs: ./.github/scripts/update-homebrew-greptme-version.sh notification: - if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'push' || github.event_name == 'schedule') && always() }} + if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name == 'push' || (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true')) && always() }} name: Send notification to Greptime team needs: [ + allocate-runners, release-images-to-dockerhub, build-macos-artifacts, build-windows-artifacts, From abf1396c285b102ba66fde72ccd775e155c3d75c Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 29 Sep 2026 08:05:50 +0000 Subject: [PATCH 075/133] fix(client): yield Flight batches and affected rows without waiting for next message (#8918) * fix(client): avoid Flight metrics lookahead stalls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(client): extract trailing Flight metrics task Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(client): synchronize trailing metrics and bound cancellation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/client/src/database.rs | 348 +++++++++++++++++++++++---- src/client/src/flight.rs | 41 +--- src/client/src/region.rs | 113 ++++----- src/flow/src/batching_mode/task.rs | 6 + tests-integration/src/grpc/flight.rs | 3 +- 5 files changed, 364 insertions(+), 147 deletions(-) diff --git a/src/client/src/database.rs b/src/client/src/database.rs index e97fff6db00..de3c722a86a 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use std::task::{Context, Poll}; use std::time::Duration; @@ -55,6 +55,7 @@ use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; use snafu::{IntoError, ResultExt}; +use tokio::sync::Notify; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; @@ -70,20 +71,47 @@ type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; const HINTS_METADATA_KEY: &str = "x-greptime-hints"; +/// Maximum time to wait for the optional trailing metrics message after +/// affected rows have already been delivered. +const FLIGHT_TRAILING_METRICS_TIMEOUT: Duration = Duration::from_secs(5); /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully -/// drained and [`Self::is_ready`] returns `true`. +/// drained and [`Self::is_ready`] returns `true`. Affected-row outputs may +/// briefly await a compatibility trailing metrics message. #[derive(Debug, Clone, Default)] pub struct OutputMetrics { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct OutputMetricsInner { metrics: RwLock>, + completion_error: RwLock>, ready: AtomicBool, + ready_notify: Notify, + compatibility_task: Mutex>, +} + +impl Default for OutputMetricsInner { + fn default() -> Self { + Self { + metrics: RwLock::new(None), + completion_error: RwLock::new(None), + ready: AtomicBool::new(false), + ready_notify: Notify::new(), + compatibility_task: Mutex::new(None), + } + } +} + +impl Drop for OutputMetricsInner { + fn drop(&mut self) { + if let Some(handle) = self.compatibility_task.get_mut().unwrap().take() { + handle.abort(); + } + } } impl OutputMetrics { @@ -98,10 +126,45 @@ impl OutputMetrics { /// Marks the terminal metrics as final for this output. pub fn mark_ready(&self) { - let _ = self + if self .inner .ready - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire); + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.inner.ready_notify.notify_waiters(); + } + } + + /// Waits until terminal metrics are final. + pub async fn wait_ready(&self) { + loop { + let notified = self.inner.ready_notify.notified(); + if self.is_ready() { + return; + } + notified.await; + } + } + + /// Returns an error encountered while completing the output, if any. + pub fn completion_error(&self) -> Option { + self.inner.completion_error.read().unwrap().clone() + } + + fn set_completion_error(&self, error: impl Into) { + *self.inner.completion_error.write().unwrap() = Some(error.into()); + } + + fn set_compatibility_task(&self, handle: tokio::task::AbortHandle) { + let mut task = self.inner.compatibility_task.lock().unwrap(); + if !self.is_ready() { + *task = Some(handle); + } + } + + fn take_compatibility_task(&self) -> Option { + self.inner.compatibility_task.lock().unwrap().take() } /// Returns whether terminal metrics are final. @@ -148,7 +211,8 @@ impl OutputMetrics { /// /// The contained [`OutputMetrics`] lets callers read stream terminal metrics /// after consuming `output`. For non-stream outputs, metrics are ready -/// immediately. +/// immediately. Flight affected-row outputs without inline metrics require +/// [`OutputMetrics::wait_ready`] before compatibility trailing metrics can be read. #[derive(Debug)] pub struct OutputWithMetrics { pub output: Output, @@ -194,6 +258,55 @@ fn parse_terminal_metrics(metrics_json: &str) -> Result { }) } +fn spawn_affected_rows_trailing_metrics_task( + terminal_metrics: &OutputMetrics, + mut reader: FlightMessageReader, +) where + S: Stream> + Send + Unpin + 'static, +{ + let metrics_ref = Arc::downgrade(&terminal_metrics.inner); + let remote_addr = reader.remote_addr().to_string(); + let task = common_runtime::spawn_global(async move { + let result = + tokio::time::timeout(FLIGHT_TRAILING_METRICS_TIMEOUT, reader.read_next()).await; + let Some(inner) = metrics_ref.upgrade() else { + return; + }; + let metrics = OutputMetrics { inner }; + match result { + Ok(Ok(Some(FlightMessage::Metrics(s)))) => match parse_terminal_metrics(&s) { + Ok(metrics_json) => metrics.update(Some(metrics_json)), + Err(error) => { + metrics.set_completion_error(error.to_string()); + warn!( + "Failed to decode trailing Flight metrics from {}: {}", + remote_addr, error + ); + } + }, + Ok(Ok(None)) => {} + Ok(Ok(Some(other))) => { + let error = format!("Unexpected trailing Flight message: {other:?}"); + metrics.set_completion_error(error.clone()); + warn!("{} from {}", error, remote_addr); + } + Ok(Err(error)) => { + let error = flight_stream_error(&remote_addr, error); + metrics.set_completion_error(error.to_string()); + warn!("{}", error); + } + Err(_) => { + let error = "Timed out waiting for trailing Flight metrics"; + metrics.set_completion_error(error); + warn!("{} from {}", error, remote_addr); + } + } + metrics.mark_ready(); + metrics.take_compatibility_task(); + }); + terminal_metrics.set_compatibility_task(task.abort_handle()); +} + struct StreamWithMetrics { stream: common_recordbatch::SendableRecordBatchStream, metrics: OutputMetrics, @@ -233,6 +346,9 @@ impl Stream for StreamWithMetrics { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let polled = Pin::new(&mut self.stream).poll_next(cx); + if let Poll::Ready(Some(Err(error))) = &polled { + self.metrics.set_completion_error(error.to_string()); + } if let Poll::Ready(None) = &polled { self.sync_terminal_metrics(); self.metrics.mark_ready(); @@ -280,32 +396,12 @@ where FlightMessage::AffectedRows { rows, metrics } => { let terminal_metrics = OutputMetrics::new(); if let Some(metrics) = metrics { + // Inline metrics are authoritative. They complete the output + // without touching the rest of the Flight stream. terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); - } - let next_message = reader - .read_next() - .await - .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; - match next_message { - None => terminal_metrics.mark_ready(), - Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { - terminal_metrics.update(Some(parse_terminal_metrics(&s)?)); - terminal_metrics.mark_ready(); - } - Some(FlightMessage::Metrics(_)) => { - return IllegalFlightMessagesSnafu { - reason: "'AffectedRows' Flight metadata already carries Metrics and cannot be followed by another Metrics message", - } - .fail(); - } - Some(other) => { - return IllegalFlightMessagesSnafu { - reason: format!( - "'AffectedRows' Flight message can only be followed by a Metrics message, got {other:?}" - ), - } - .fail(); - } + terminal_metrics.mark_ready(); + } else { + spawn_affected_rows_trailing_metrics_task(&terminal_metrics, reader); } Ok(OutputWithMetrics { output: Output::new_with_affected_rows(rows), @@ -749,7 +845,9 @@ impl Database { /// Executes a SQL query and returns the output with terminal metrics. /// /// For stream outputs, callers must consume the stream before reading final - /// terminal metrics from [`OutputWithMetrics::metrics`]. + /// terminal metrics from [`OutputWithMetrics::metrics`]. For affected-row + /// outputs without inline metrics, call [`OutputMetrics::wait_ready`] when + /// compatibility trailing metrics are required. pub async fn sql_with_terminal_metrics( &self, sql: S, @@ -986,6 +1084,7 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; + use tokio::sync::oneshot; use tonic::codegen::http::{HeaderMap, HeaderValue}; use tonic::metadata::MetadataMap; use tonic::{Code, Status}; @@ -1296,11 +1395,6 @@ mod tests { ); } - #[test] - fn test_parse_terminal_metrics_rejects_invalid_json() { - assert!(parse_terminal_metrics("{not-json}").is_err()); - } - #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { let output = output_from_flight_message_stream( @@ -1322,27 +1416,191 @@ mod tests { } #[tokio::test] - async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { + async fn test_affected_rows_inline_metrics_do_not_poll_trailer() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream( + let output = output_from_flight_message_stream( "test-peer".to_string(), futures_util::stream::iter(vec![ Ok(FlightMessage::AffectedRows { rows: 3, - metrics: Some(metrics_json.clone()), + metrics: Some(metrics_json), }), - Ok(FlightMessage::Metrics(metrics_json)), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(99))), ] as Vec>), ) .await - .unwrap_err(); + .unwrap(); - assert!( - err.to_string().contains("already carries Metrics"), - "unexpected error: {err:?}" + assert!(output.metrics.is_ready()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) ); } + #[tokio::test] + async fn test_affected_rows_without_inline_metrics_becomes_ready_after_trailer() { + // Hold the trailer until the not-ready state has been observed. + let (trailer_tx, trailer_rx) = oneshot::channel(); + let trailer = futures_util::stream::once(trailer_rx).map(|ready| { + ready.unwrap(); + Ok(FlightMessage::Metrics(terminal_metrics_json())) + }); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })]) + .chain(trailer), + ) + .await + .unwrap(); + + assert!(!output.metrics.is_ready()); + trailer_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), output.metrics.wait_ready()) + .await + .expect("terminal metrics must become ready once the trailer arrives"); + assert!(output.metrics.completion_error().is_none()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) + ); + } + + #[tokio::test] + async fn test_affected_rows_malformed_trailer_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("Invalid terminal metrics message")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_transport_trailer_error_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Err(Status::unavailable("trailer read failed").into()), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("trailer read failed")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_compatibility_reader_is_cancelled_after_second_poll_begins() { + struct DropProbe { + first: Option>, + polled: Option>, + dropped: Option>, + } + + impl Stream for DropProbe { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if let Some(message) = self.first.take() { + return Poll::Ready(Some(message)); + } + if let Some(polled) = self.polled.take() { + let _ = polled.send(()); + } + Poll::Pending + } + } + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } + } + + let (polled_tx, polled_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + DropProbe { + first: Some(Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })), + polled: Some(polled_tx), + dropped: Some(dropped_tx), + }, + ) + .await + .unwrap(); + polled_rx.await.unwrap(); + drop(output); + // Require cancellation before the five-second compatibility timeout. + tokio::time::timeout(Duration::from_secs(1), dropped_rx) + .await + .expect("the compatibility reader must be dropped once the output is dropped") + .unwrap(); + } + + #[tokio::test] + async fn test_schema_record_batch_yields_before_pending_message_stream() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "v", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let messages = futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ] as Vec>) + .chain(futures_util::stream::pending()); + let output = output_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + let OutputData::Stream(mut stream) = output.output.data else { + panic!("expected stream output"); + }; + + let batch = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(!output.metrics.is_ready()); + } + #[tokio::test] async fn test_invalid_terminal_metrics_after_record_batch_yields_batch_then_error() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 1753fdd5246..3a21e6daa4e 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,40 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::pin::Pin; - use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use futures_util::stream::Peekable; use futures_util::{Stream, StreamExt}; use snafu::{OptionExt, ResultExt}; use crate::Result; use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum FlightMessageKind { - Schema, - RecordBatch, - AffectedRows, - Metrics, -} - -impl From<&FlightMessage> for FlightMessageKind { - fn from(message: &FlightMessage) -> Self { - match message { - FlightMessage::Schema(_) => Self::Schema, - FlightMessage::RecordBatch(_) => Self::RecordBatch, - FlightMessage::AffectedRows { .. } => Self::AffectedRows, - FlightMessage::Metrics(_) => Self::Metrics, - } - } -} - pub(crate) struct FlightMessageReader { /// Remote Flight peer associated with this response stream. remote_addr: String, - messages: Peekable, + messages: S, } impl FlightMessageReader @@ -55,7 +33,7 @@ where pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { Self { remote_addr: remote_addr.into(), - messages: messages.peekable(), + messages, } } @@ -72,21 +50,6 @@ where pub(crate) async fn read_next(&mut self) -> Result> { self.messages.next().await.transpose() } - - pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { - match Pin::new(&mut self.messages).peek().await { - Some(Ok(message)) => Ok(Some(message.into())), - None => Ok(None), - Some(Err(_)) => match self.read_next().await { - // `peek` only borrows the error; consume it to preserve the source error. - Err(error) => Err(error), - Ok(_) => IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek".to_string(), - } - .fail(), - }, - } - } } pub(crate) fn decode_flight_data( diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f0c683ec155..fb2e4d10ca1 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -48,7 +48,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -247,9 +247,7 @@ where "poll_flight_data_stream" )); - let mut stream_ended = false; - - while !stream_ended { + loop { let flight_message = match reader.read_next().await { Ok(Some(message)) => message, Ok(None) => break, @@ -262,59 +260,27 @@ where match flight_message { FlightMessage::RecordBatch(record_batch) => { - let result_to_yield = - RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - - // Metrics follow a batch so MergeScan can observe them before yielding it. - match reader.peek_next_message_kind().await { - Ok(Some(FlightMessageKind::Metrics)) => { - let metrics_message = match reader.read_next().await { - Ok(Some(FlightMessage::Metrics(metrics))) => metrics, - Ok(Some(_) | None) => { - yield IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek", - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(error) => { - yield Err(BoxedError::new(flight_stream_error( - &stream_addr, - error, - ))) - .context(ExternalSnafu); - break; - } - }; - let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); - metrics_ref.swap(metrics); - } - Ok(Some(FlightMessageKind::RecordBatch)) => {} - Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Ok(None) => stream_ended = true, - Err(error) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) - .context(ExternalSnafu); - break; - } - } - - yield Ok(result_to_yield); + // Deliver each batch immediately. In particular, do not + // wait for a possible following Metrics message; it is + // consumed on the next poll of this stream. + yield Ok(RecordBatch::from_df_record_batch( + schema_cloned.clone(), + record_batch, + )); } FlightMessage::Metrics(s) => { // Metrics may arrive before the next RecordBatch. - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); + match serde_json::from_str(&s) { + Ok(metrics) => { + metrics_ref.swap(Some(Arc::new(metrics))); + } + Err(error) => { + common_telemetry::warn!( + "Failed to decode region Flight metrics: {}", + error + ); + } + } continue; } _ => { @@ -719,7 +685,34 @@ mod test { } #[tokio::test] - async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + async fn test_record_batch_is_yielded_without_waiting_for_next_message() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let messages = stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]) + .chain(stream::pending::>()); + let mut recordbatches = + recordbatches_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + + let batch = tokio::time::timeout(Duration::from_secs(1), recordbatches.next()) + .await + .expect("the first batch must not wait for lookahead") + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + + #[tokio::test] + async fn test_malformed_region_metrics_are_non_fatal() { let schema = test_schema(); let batch = RecordBatch::new( schema.clone(), @@ -730,23 +723,19 @@ mod test { "test-peer".to_string(), stream::iter(vec![ Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics(test_metrics_json())), ]), ) .await .unwrap(); - let batch = recordbatches.next().await.unwrap().unwrap(); - assert_eq!(batch.num_rows(), 1); - - let metrics = recordbatches.metrics().unwrap(); - assert_eq!(metrics.elapsed_compute, 7); + assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1); assert!(recordbatches.next().await.is_none()); } #[tokio::test] - async fn test_record_batch_stream_preserves_peeked_record_batch() { + async fn test_record_batch_stream_preserves_following_record_batch() { let schema = test_schema(); let first_batch = RecordBatch::new( schema.clone(), diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 40020f639c1..c2c3082995a 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -1400,6 +1400,12 @@ impl BatchingTask { match res { Ok(res) => { let (affected_rows, _) = res.output.extract_rows_and_cost(); + if matches!(&res.output.data, common_query::OutputData::AffectedRows(_)) { + res.metrics.wait_ready().await; + } + if let Some(error) = res.metrics.completion_error() { + warn!("Flow {flow_id} completed with terminal metrics error: {error}"); + } debug!( "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}", elapsed, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index ef3df79f558..e74d4d5bf70 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -990,7 +990,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); - assert!(result.metrics.is_ready()); + result.metrics.wait_ready().await; assert!(result.region_watermark_map().is_none()); let err = client @@ -1016,6 +1016,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); + result.metrics.wait_ready().await; assert_eq!( result.region_watermark_map(), Some(std::collections::HashMap::from([previous_watermark])) From 218000e21b4d4020fd3233beaf021cb940755822 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 29 Sep 2026 08:13:28 +0000 Subject: [PATCH 076/133] fix(promql): resolve dotted column names as unqualified columns (#9391) * fix(promql): resolve dotted column names as unqualified columns col(), From<&str>/From for Column and string join keys go through Column::from_qualified_name, which splits `service.name` into relation `service` and column `name` and lowercases unquoted identifiers. Build PromQL column references with Column::from_name / ident() instead. Signed-off-by: Dennis Zhuang * fix(servers): resolve remote read matcher labels as unqualified columns Signed-off-by: Dennis Zhuang * test(promql): cover same-name columns differing in case and without() Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/promql/src/extension_plan/absent.rs | 4 +- src/promql/src/extension_plan/empty_metric.rs | 4 +- .../src/extension_plan/histogram_fold.rs | 10 +- .../src/extension_plan/instant_manipulate.rs | 6 +- src/promql/src/extension_plan/normalize.rs | 6 +- .../src/extension_plan/range_manipulate.rs | 6 +- .../src/extension_plan/scalar_calculate.rs | 8 +- .../src/extension_plan/series_divide.rs | 6 +- .../src/extension_plan/union_distinct_on.rs | 6 +- src/query/src/promql/planner.rs | 38 +- .../src/promql/planner/function_plans.rs | 14 +- src/query/src/promql/planner/set_operator.rs | 1 + src/query/src/promql/planner/test.rs | 8 +- src/servers/src/prom_store.rs | 31 +- .../common/promql/dotted_column_names.result | 505 ++++++++++++++++++ .../common/promql/dotted_column_names.sql | 139 +++++ 16 files changed, 728 insertions(+), 64 deletions(-) create mode 100644 tests/cases/standalone/common/promql/dotted_column_names.result create mode 100644 tests/cases/standalone/common/promql/dotted_column_names.sql diff --git a/src/promql/src/extension_plan/absent.rs b/src/promql/src/extension_plan/absent.rs index fa811b3df7d..9afe85873c5 100644 --- a/src/promql/src/extension_plan/absent.rs +++ b/src/promql/src/extension_plan/absent.rs @@ -34,7 +34,7 @@ use datafusion::physical_plan::{ Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use datafusion_common::DFSchema; -use datafusion_expr::{EmptyRelation, col}; +use datafusion_expr::{EmptyRelation, ident}; use datatypes::arrow; use datatypes::arrow::array::{ArrayRef, Float64Array, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit}; @@ -108,7 +108,7 @@ impl UserDefinedLogicalNodeCore for Absent { return vec![]; } - vec![col(&self.time_index_column)] + vec![ident(&self.time_index_column)] } fn necessary_children_exprs(&self, _output_columns: &[usize]) -> Option>> { diff --git a/src/promql/src/extension_plan/empty_metric.rs b/src/promql/src/extension_plan/empty_metric.rs index 91c8a1666a2..4d2e029417a 100644 --- a/src/promql/src/extension_plan/empty_metric.rs +++ b/src/promql/src/extension_plan/empty_metric.rs @@ -40,7 +40,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::physical_planner::PhysicalPlanner; -use datafusion::prelude::{Expr, col, lit}; +use datafusion::prelude::{Expr, ident, lit}; use datafusion_expr::LogicalPlanBuilder; use datatypes::arrow::array::TimestampMillisecondArray; use datatypes::arrow::datatypes::SchemaRef; @@ -409,7 +409,7 @@ fn build_ts_only_schema(column_name: &str) -> DFSchema { pub fn build_special_time_expr(time_index_column_name: &str) -> Expr { let input_schema = build_ts_only_schema(time_index_column_name); // safety: should not failed (UT covers this) - col(time_index_column_name) + ident(time_index_column_name) .cast_to(&DataType::Int64, &input_schema) .unwrap() .cast_to(&DataType::Float64, &input_schema) diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 94769df28a4..7fdab474f5d 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -41,7 +41,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::prelude::{Column, Expr}; -use datafusion_expr::{EmptyRelation, col}; +use datafusion_expr::{EmptyRelation, ident}; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::{ConcreteDataType, DataType as GtDataType}; use datatypes::value::{OrderedF64, Value, ValueRef}; @@ -141,14 +141,14 @@ impl UserDefinedLogicalNodeCore for HistogramFold { } let mut exprs = vec![ - col(&self.le_column), - col(&self.ts_column), - col(&self.field_column), + ident(&self.le_column), + ident(&self.ts_column), + ident(&self.field_column), ]; exprs.extend(self.input.schema().fields().iter().filter_map(|f| { let name = f.name(); if name != &self.le_column && name != &self.ts_column && name != &self.field_column { - Some(col(name)) + Some(ident(name)) } else { None } diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 45857879d73..4085d382c37 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -37,7 +37,7 @@ use datafusion::physical_plan::{ PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; @@ -140,8 +140,8 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { return vec![]; } - let mut exprs = vec![col(&self.time_index_column)]; - exprs.extend(self.staleness_field_columns().map(col)); + let mut exprs = vec![ident(&self.time_index_column)]; + exprs.extend(self.staleness_field_columns().map(ident)); exprs } diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index 1ca015b10f1..b1ea6ca3e7c 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::array::TimestampMillisecondArray; use datatypes::arrow::datatypes::{SchemaRef, TimestampMillisecondType}; use datatypes::arrow::record_batch::RecordBatch; @@ -93,8 +93,8 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index_column_name))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index_column_name))) .collect() } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 96a0c97778d..8ea7ba6b191 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -38,7 +38,7 @@ use datafusion::physical_plan::{ InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, StatisticsArgs, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; @@ -318,8 +318,8 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { } let mut exprs = Vec::with_capacity(1 + self.field_columns.len()); - exprs.push(col(&self.time_index)); - exprs.extend(self.field_columns.iter().map(col)); + exprs.push(ident(&self.time_index)); + exprs.extend(self.field_columns.iter().map(ident)); exprs } diff --git a/src/promql/src/extension_plan/scalar_calculate.rs b/src/promql/src/extension_plan/scalar_calculate.rs index e71333837f3..4ad188a1d4f 100644 --- a/src/promql/src/extension_plan/scalar_calculate.rs +++ b/src/promql/src/extension_plan/scalar_calculate.rs @@ -33,7 +33,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, }; use datafusion::prelude::Expr; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::array::{Array, ArrayRef, Float64Array, TimestampMillisecondArray}; use datatypes::arrow::compute::{CastOptions, cast_with_options, concat_batches}; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; @@ -282,9 +282,9 @@ impl UserDefinedLogicalNodeCore for ScalarCalculate { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index))) - .chain(std::iter::once(col(&self.field_column))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index))) + .chain(std::iter::once(ident(&self.field_column))) .collect() } diff --git a/src/promql/src/extension_plan/series_divide.rs b/src/promql/src/extension_plan/series_divide.rs index bf62f2c3332..bdd90696d26 100644 --- a/src/promql/src/extension_plan/series_divide.rs +++ b/src/promql/src/extension_plan/series_divide.rs @@ -38,7 +38,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use datatypes::arrow_array::string_array_value_at_index; use datatypes::compute::SortOptions; @@ -186,8 +186,8 @@ impl UserDefinedLogicalNodeCore for SeriesDivide { self.tag_columns .iter() - .map(col) - .chain(std::iter::once(col(&self.time_index_column))) + .map(ident) + .chain(std::iter::once(ident(&self.time_index_column))) .collect() } diff --git a/src/promql/src/extension_plan/union_distinct_on.rs b/src/promql/src/extension_plan/union_distinct_on.rs index 774904a27aa..90d9dda637e 100644 --- a/src/promql/src/extension_plan/union_distinct_on.rs +++ b/src/promql/src/extension_plan/union_distinct_on.rs @@ -34,7 +34,7 @@ use datafusion::physical_plan::{ Partitioning, PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, hash_utils, }; -use datafusion_expr::col; +use datafusion_expr::ident; use datatypes::arrow::compute; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; @@ -278,12 +278,12 @@ impl UserDefinedLogicalNodeCore for UnionDistinctOn { let mut exprs = self .compare_key_indices .iter() - .filter_map(|index| fields.get(*index).map(|field| col(field.name()))) + .filter_map(|index| fields.get(*index).map(|field| ident(field.name()))) .collect::>(); if !self.compare_key_indices.contains(&self.ts_col_idx) && let Some(field) = fields.get(self.ts_col_idx) { - exprs.push(col(field.name())); + exprs.push(ident(field.name())); } exprs } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index ffa99a52d0d..dc023762b00 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -58,7 +58,7 @@ use datafusion_common::{DFSchema, NullEquality, TableReference}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; use datafusion_expr::utils::{conjunction, disjunction}; -use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, col, lit}; +use datafusion_expr::{ExprSchemable, Literal, SortExpr, TableSource, ident, lit}; use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; @@ -678,7 +678,7 @@ impl PromPlanner { .cloned() .chain(prev_field_exprs.clone()) .collect::>(); - group_exprs.push(col(label)); + group_exprs.push(ident(label)); let project_fields = self .create_field_column_exprs()? .into_iter() @@ -798,10 +798,10 @@ impl PromPlanner { ), self.promql_annotations.clone(), )), - args: vec![col(&histogram_column)], + args: vec![ident(&histogram_column)], }); - let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram) - .otherwise(col(&float_column).is_not_null()) + let keep_float = when(ident(&histogram_column).is_not_null(), drop_histogram) + .otherwise(ident(&float_column).is_not_null()) .context(DataFusionPlanningSnafu)?; input = LogicalPlanBuilder::from(input) .filter(keep_float) @@ -852,7 +852,7 @@ impl PromPlanner { .iter() .fold(None, |expr, rank| { let predicate = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(col(rank)), + left: Box::new(ident(rank)), op: Operator::LtEq, right: Box::new(val.clone()), }); @@ -868,7 +868,7 @@ impl PromPlanner { }) .unwrap(); - let rank_columns: Vec<_> = rank_columns.into_iter().map(col).collect(); + let rank_columns: Vec<_> = rank_columns.into_iter().map(ident).collect(); let mut new_group_exprs = group_exprs.clone(); // Order by ranks @@ -921,10 +921,12 @@ impl PromPlanner { if Self::field_column_is_native_histogram(&input_schema, col) { Ok(DfExpr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramNeg::scalar_udf()), - args: vec![DfExpr::Column(col.into())], + args: vec![DfExpr::Column(Column::from_name(col))], })) } else { - Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + Ok(DfExpr::Negative(Box::new(DfExpr::Column( + Column::from_name(col), + )))) } }) } @@ -1034,7 +1036,7 @@ impl PromPlanner { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; let rhs_is_histogram = Self::field_column_is_native_histogram(&input_schema, col); - let rhs = DfExpr::Column(col.into()); + let rhs = DfExpr::Column(Column::from_name(col)); let mut binary_expr = match Self::native_histogram_binary_expr( *op, expr.clone(), @@ -1101,7 +1103,7 @@ impl PromPlanner { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; let lhs_is_histogram = Self::field_column_is_native_histogram(&input_schema, col); - let lhs = DfExpr::Column(col.into()); + let lhs = DfExpr::Column(Column::from_name(col)); let mut binary_expr = match Self::native_histogram_binary_expr( *op, lhs.clone(), @@ -1753,7 +1755,7 @@ impl PromPlanner { // `timestamp()` preserves the shifted selector timeline even though // SeriesNormalize now retains raw native timestamp storage. Decimal // arithmetic shifts before truncating to milliseconds. - let unit_factor = match col(&time_index_column) + let unit_factor = match ident(&time_index_column) .get_type(normalize.schema()) .context(DataFusionPlanningSnafu)? { @@ -1763,7 +1765,7 @@ impl PromPlanner { ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6), _ => unreachable!("time index is a timestamp"), }; - let sample_time = col(&time_index_column) + let sample_time = ident(&time_index_column) .cast_to(&ArrowDataType::Int64, normalize.schema()) .context(DataFusionPlanningSnafu)? .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema()) @@ -1877,7 +1879,7 @@ impl PromPlanner { input: LogicalPlan, timestamp_value_column: &str, ) -> Result { - let time_expr = col(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN); + let time_expr = ident(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN); self.ctx.field_columns = vec![time_expr.schema_name().to_string()]; let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2); project_exprs.push(self.create_time_index_column_expr()?); @@ -2507,7 +2509,7 @@ impl PromPlanner { // collect remaining fields and convert to col expr let mut exprs = all_fields .into_iter() - .map(|c| DfExpr::Column(Column::from(c))) + .map(|c| DfExpr::Column(Column::from_name(c))) .collect::>(); // add timestamp column @@ -4774,7 +4776,7 @@ impl PromPlanner { .map(|col| { let mut sort_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 1); // Order by value in the specific order - sort_exprs.push(DfExpr::Column(Column::from(col)).sort(asc, true)); + sort_exprs.push(DfExpr::Column(Column::from_name(col)).sort(asc, true)); // Then tags if the values are equal, // Try to ensure the relative stability of the output results. sort_exprs.extend(tag_sort_exprs.clone()); @@ -5547,7 +5549,7 @@ impl PromPlanner { .collect::>(); let assert_expr = DfExpr::ScalarFunction(ScalarFunction { func: Arc::new(UniqueMatchGroup::scalar_udf(group_labels, violation)), - args: std::iter::once(col(count_column.as_str())) + args: std::iter::once(ident(count_column.as_str())) .chain(group_exprs) .collect(), }); @@ -6186,7 +6188,7 @@ impl PromPlanner { /// Generate an expr like `date_part("hour", )`. Caller should ensure the /// time index column in context is set fn date_part_on_time_index(&self, date_part: &str) -> Result { - let input_expr = datafusion::logical_expr::col( + let input_expr = datafusion::logical_expr::ident( self.ctx .time_index_column .as_ref() diff --git a/src/query/src/promql/planner/function_plans.rs b/src/query/src/promql/planner/function_plans.rs index 68a254dd2bc..4723fe35232 100644 --- a/src/query/src/promql/planner/function_plans.rs +++ b/src/query/src/promql/planner/function_plans.rs @@ -28,7 +28,7 @@ use datafusion::scalar::ScalarValue; use datafusion_common::DFSchema; use datafusion_expr::expr_fn::when; use datafusion_expr::simplify::SimplifyContext; -use datafusion_expr::{col, lit}; +use datafusion_expr::{ident, lit}; use datafusion_functions::core::coalesce; use datatypes::arrow::datatypes::DataType as ArrowDataType; use promql::extension_plan::{ @@ -253,12 +253,12 @@ impl PromPlanner { "vector contains a mix of classic and native histograms".to_string(), self.promql_annotations.clone(), )), - args: vec![col(&float_field), col(&histogram_field)], + args: vec![ident(&float_field), ident(&histogram_field)], }); let keep = when( - col(&float_field) + ident(&float_field) .is_not_null() - .and(col(&histogram_field).is_not_null()), + .and(ident(&histogram_field).is_not_null()), record_collision, ) .otherwise(lit(true)) @@ -268,7 +268,7 @@ impl PromPlanner { let output_field = native_expr.schema_name().to_string(); let value = DfExpr::ScalarFunction(ScalarFunction { func: coalesce(), - args: vec![col(&float_field), native_expr], + args: vec![ident(&float_field), native_expr], }); self.ctx.field_columns = vec![output_field.clone()]; LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { @@ -277,9 +277,9 @@ impl PromPlanner { .filter(keep) .context(DataFusionPlanningSnafu)? .project( - std::iter::once(col(&time_index_column)) + std::iter::once(ident(&time_index_column)) .chain(std::iter::once(value.alias(output_field))) - .chain(tag_columns.iter().map(col)), + .chain(tag_columns.iter().map(ident)), ) .context(DataFusionPlanningSnafu)? .build() diff --git a/src/query/src/promql/planner/set_operator.rs b/src/query/src/promql/planner/set_operator.rs index 51cc100408f..7a14ff38ad4 100644 --- a/src/query/src/promql/planner/set_operator.rs +++ b/src/query/src/promql/planner/set_operator.rs @@ -764,6 +764,7 @@ impl PromPlanner { let join_keys = left_tag_col_set .into_iter() .chain([left_time_index]) + .map(Column::from_name) .collect::>(); ensure!( diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 0393930b033..382f4791e28 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -36,7 +36,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, provider_as_source}; use datafusion::execution::context::SessionContext; -use datafusion::logical_expr::Extension; +use datafusion::logical_expr::{Extension, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql::extension_plan::HistogramFold; @@ -1210,8 +1210,10 @@ fn classic_and_native_histogram_table_provider( let table_name = "mixed_histogram"; let catalog = MemoryCatalogManager::with_default_setup(); let schema = Arc::new(Schema::new(vec![ + // A dotted tag name guards the mixed histogram_quantile projection + // against qualified-name parsing (#9390). ColumnSchema::new( - "tag".to_string(), + "service.name".to_string(), ConcreteDataType::string_datatype(), false, ), @@ -4386,7 +4388,7 @@ async fn mixed_histogram_helpers_execute_classic_and_native_samples() { .iter() .flat_map(|batch| { let tags = batch - .column_by_name("tag") + .column_by_name("service.name") .unwrap() .as_any() .downcast_ref::() diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index e202070d98d..fcf177c1e5f 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -36,7 +36,7 @@ use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{tracing, warn}; use datafusion::dataframe::DataFrame; -use datafusion::prelude::{Expr, col, lit, regexp_match}; +use datafusion::prelude::{Expr, ident, lit, regexp_match}; use datafusion_common::ScalarValue; use datafusion_expr::LogicalPlan; use snafu::{OptionExt, ResultExt, ensure}; @@ -150,8 +150,9 @@ pub fn query_to_plan( let mut conditions = Vec::with_capacity(label_matches.len() + 1); conditions - .push(col(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); - conditions.push(col(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); + .push(ident(timestamp_column_name).gt_eq(lit_timestamp_millisecond(start_timestamp_ms))); + conditions + .push(ident(timestamp_column_name).lt_eq(lit_timestamp_millisecond(end_timestamp_ms))); for m in label_matches { let name = &m.name; @@ -170,18 +171,18 @@ pub fn query_to_plan( match m_type { MatcherType::Eq => { - conditions.push(col(name).eq(lit(value))); + conditions.push(ident(name).eq(lit(value))); } MatcherType::Neq => { - conditions.push(col(name).not_eq(lit(value))); + conditions.push(ident(name).not_eq(lit(value))); } // Case sensitive regexp match MatcherType::Re => { - conditions.push(regexp_match(col(name), lit(value), None).is_not_null()); + conditions.push(regexp_match(ident(name), lit(value), None).is_not_null()); } // Case sensitive regexp not match MatcherType::Nre => { - conditions.push(regexp_match(col(name), lit(value), None).is_null()); + conditions.push(regexp_match(ident(name), lit(value), None).is_null()); } } } @@ -890,6 +891,8 @@ mod tests { ColumnSchema::new(greptime_value(), ConcreteDataType::float64_datatype(), true), ColumnSchema::new("instance", ConcreteDataType::string_datatype(), true), ColumnSchema::new("job", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("service.name", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("Region", ConcreteDataType::string_datatype(), true), ])); let recordbatch = RecordBatch::new( schema, @@ -898,6 +901,8 @@ mod tests { Arc::new(Float64Vector::from_vec(vec![3.0])) as _, Arc::new(StringVector::from(vec!["host1"])) as _, Arc::new(StringVector::from(vec!["job"])) as _, + Arc::new(StringVector::from(vec!["api"])) as _, + Arc::new(StringVector::from(vec!["us"])) as _, ], ) .unwrap(); @@ -936,6 +941,16 @@ mod tests { value: "localhost".to_string(), r#type: NEQ_TYPE, }, + LabelMatcher { + name: "service.name".to_string(), + value: "api".to_string(), + r#type: EQ_TYPE, + }, + LabelMatcher { + name: "Region".to_string(), + value: "us".to_string(), + r#type: EQ_TYPE, + }, ], ..Default::default() }; @@ -946,7 +961,7 @@ mod tests { let ts_col = greptime_timestamp(); let expected = format!( - "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"*prom*\")) IS NOT NULL AND ?table?.instance != Utf8(\"localhost\")\n TableScan: ?table?", + "Filter: ?table?.{} >= TimestampMillisecond(1000, None) AND ?table?.{} <= TimestampMillisecond(2000, None) AND regexp_match(?table?.job, Utf8(\"*prom*\")) IS NOT NULL AND ?table?.instance != Utf8(\"localhost\") AND ?table?.service.name = Utf8(\"api\") AND ?table?.Region = Utf8(\"us\")\n TableScan: ?table?", ts_col, ts_col ); assert_eq!(expected, display_string); diff --git a/tests/cases/standalone/common/promql/dotted_column_names.result b/tests/cases/standalone/common/promql/dotted_column_names.result new file mode 100644 index 00000000000..c855a5ba5c7 --- /dev/null +++ b/tests/cases/standalone/common/promql/dotted_column_names.result @@ -0,0 +1,505 @@ +-- https://github.com/GreptimeTeam/greptimedb/issues/9390 +-- Tag, field and time index names containing dots or upper case letters must +-- be resolved as plain column names, not as `relation.column`. `Host` and `host` +-- hold different values, so reading one for the other changes the output. +CREATE TABLE "otel.m" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + "Host" STRING, + host STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", "Host", host) +) PARTITION ON COLUMNS ("service.name") ( + "service.name" < 'b', + "service.name" >= 'b' +); + +Affected Rows: 0 + +INSERT INTO "otel.m" VALUES + (0, 'a', 'h1', 'x', 1), + (0, 'a', 'h2', 'x', 2), + (0, 'b', 'h1', 'y', 3), + (5000, 'a', 'h1', 'x', 4), + (5000, 'a', 'h2', 'x', 5), + (5000, 'b', 'h1', 'y', 6), + (10000, 'a', 'h1', 'x', 8), + (10000, 'a', 'h2', 'x', 9), + (10000, 'b', 'h1', 'y', 8); + +Affected Rows: 9 + +CREATE TABLE "otel.h" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + le STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", le) +); + +Affected Rows: 0 + +INSERT INTO "otel.h" VALUES + (0, 'a', '0.1', 1), + (0, 'a', '1', 3), + (0, 'a', '+Inf', 4), + (10000, 'a', '0.1', 2), + (10000, 'a', '1', 6), + (10000, 'a', '+Inf', 10); + +Affected Rows: 6 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h1 | x | 1.0 | +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | a | h1 | x | 4.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | a | h1 | x | 8.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on("service.name") {"otel.m", "service.name"="b"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on(host) {"otel.m", host="y"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | b | h1 | y | 3.0 | +| 1970-01-01T00:00:05 | b | h1 | y | 6.0 | +| 1970-01-01T00:00:10 | b | h1 | y | 8.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless {"otel.m", "Host"="h1"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h2 | x | 2.0 | +| 1970-01-01T00:00:05 | a | h2 | x | 5.0 | +| 1970-01-01T00:00:10 | a | h2 | x | 9.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless on("service.name") {"otel.m", "service.name"="b"}; + ++---------------------+--------------+------+------+-------+ +| ts.time | service.name | Host | host | v.val | ++---------------------+--------------+------+------+-------+ +| 1970-01-01T00:00:00 | a | h1 | x | 1.0 | +| 1970-01-01T00:00:00 | a | h2 | x | 2.0 | +| 1970-01-01T00:00:05 | a | h1 | x | 4.0 | +| 1970-01-01T00:00:05 | a | h2 | x | 5.0 | +| 1970-01-01T00:00:10 | a | h1 | x | 8.0 | +| 1970-01-01T00:00:10 | a | h2 | x | 9.0 | ++---------------------+--------------+------+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} or {"otel.m", "Host"="h1"}; + ++---------------------+------+------+--------------+-------+ +| ts.time | Host | host | service.name | v.val | ++---------------------+------+------+--------------+-------+ +| 1970-01-01T00:00:00 | h1 | x | a | 1.0 | +| 1970-01-01T00:00:00 | h1 | y | b | 3.0 | +| 1970-01-01T00:00:00 | h2 | x | a | 2.0 | +| 1970-01-01T00:00:05 | h1 | x | a | 4.0 | +| 1970-01-01T00:00:05 | h1 | y | b | 6.0 | +| 1970-01-01T00:00:05 | h2 | x | a | 5.0 | +| 1970-01-01T00:00:10 | h1 | x | a | 8.0 | +| 1970-01-01T00:00:10 | h1 | y | b | 8.0 | +| 1970-01-01T00:00:10 | h2 | x | a | 9.0 | ++---------------------+------+------+--------------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values("v.name", {"otel.m"}); + ++---------------------+---------------------+--------+ +| count(otel.m.v.val) | ts.time | v.name | ++---------------------+---------------------+--------+ +| 1 | 1970-01-01T00:00:00 | 1 | +| 1 | 1970-01-01T00:00:00 | 2 | +| 1 | 1970-01-01T00:00:00 | 3 | +| 1 | 1970-01-01T00:00:05 | 4 | +| 1 | 1970-01-01T00:00:05 | 5 | +| 1 | 1970-01-01T00:00:05 | 6 | +| 1 | 1970-01-01T00:00:10 | 9 | +| 2 | 1970-01-01T00:00:10 | 8 | ++---------------------+---------------------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values by ("service.name") ("v.name", {"otel.m"}); + ++---------------------+--------------+---------------------+--------+ +| count(otel.m.v.val) | service.name | ts.time | v.name | ++---------------------+--------------+---------------------+--------+ +| 1 | a | 1970-01-01T00:00:00 | 1 | +| 1 | a | 1970-01-01T00:00:00 | 2 | +| 1 | a | 1970-01-01T00:00:05 | 4 | +| 1 | a | 1970-01-01T00:00:05 | 5 | +| 1 | a | 1970-01-01T00:00:10 | 8 | +| 1 | a | 1970-01-01T00:00:10 | 9 | +| 1 | b | 1970-01-01T00:00:00 | 3 | +| 1 | b | 1970-01-01T00:00:05 | 6 | +| 1 | b | 1970-01-01T00:00:10 | 8 | ++---------------------+--------------+---------------------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') topk(1, {"otel.m"}); + ++-------+--------------+------+------+---------------------+ +| v.val | service.name | Host | host | ts.time | ++-------+--------------+------+------+---------------------+ +| 3.0 | b | h1 | y | 1970-01-01T00:00:00 | +| 6.0 | b | h1 | y | 1970-01-01T00:00:05 | +| 9.0 | a | h2 | x | 1970-01-01T00:00:10 | ++-------+--------------+------+------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') bottomk by ("service.name") (1, {"otel.m"}); + ++-------+--------------+------+------+---------------------+ +| v.val | service.name | Host | host | ts.time | ++-------+--------------+------+------+---------------------+ +| 1.0 | a | h1 | x | 1970-01-01T00:00:00 | +| 3.0 | b | h1 | y | 1970-01-01T00:00:00 | +| 4.0 | a | h1 | x | 1970-01-01T00:00:05 | +| 6.0 | b | h1 | y | 1970-01-01T00:00:05 | +| 8.0 | a | h1 | x | 1970-01-01T00:00:10 | +| 8.0 | b | h1 | y | 1970-01-01T00:00:10 | ++-------+--------------+------+------+---------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("service.name") ({"otel.m"}); + ++--------------+---------------------+-------------------+ +| service.name | ts.time | sum(otel.m.v.val) | ++--------------+---------------------+-------------------+ +| a | 1970-01-01T00:00:00 | 3.0 | +| a | 1970-01-01T00:00:05 | 9.0 | +| a | 1970-01-01T00:00:10 | 17.0 | +| b | 1970-01-01T00:00:00 | 3.0 | +| b | 1970-01-01T00:00:05 | 6.0 | +| b | 1970-01-01T00:00:10 | 8.0 | ++--------------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("Host") ({"otel.m"}); + ++------+---------------------+-------------------+ +| Host | ts.time | sum(otel.m.v.val) | ++------+---------------------+-------------------+ +| h1 | 1970-01-01T00:00:00 | 4.0 | +| h1 | 1970-01-01T00:00:05 | 10.0 | +| h1 | 1970-01-01T00:00:10 | 16.0 | +| h2 | 1970-01-01T00:00:00 | 2.0 | +| h2 | 1970-01-01T00:00:05 | 5.0 | +| h2 | 1970-01-01T00:00:10 | 9.0 | ++------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by (host) ({"otel.m"}); + ++------+---------------------+-------------------+ +| host | ts.time | sum(otel.m.v.val) | ++------+---------------------+-------------------+ +| x | 1970-01-01T00:00:00 | 3.0 | +| x | 1970-01-01T00:00:05 | 9.0 | +| x | 1970-01-01T00:00:10 | 17.0 | +| y | 1970-01-01T00:00:00 | 3.0 | +| y | 1970-01-01T00:00:05 | 6.0 | +| y | 1970-01-01T00:00:10 | 8.0 | ++------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum without ("Host") ({"otel.m"}); + ++------+--------------+---------------------+-------------------+ +| host | service.name | ts.time | sum(otel.m.v.val) | ++------+--------------+---------------------+-------------------+ +| x | a | 1970-01-01T00:00:00 | 3.0 | +| x | a | 1970-01-01T00:00:05 | 9.0 | +| x | a | 1970-01-01T00:00:10 | 17.0 | +| y | b | 1970-01-01T00:00:00 | 3.0 | +| y | b | 1970-01-01T00:00:05 | 6.0 | +| y | b | 1970-01-01T00:00:10 | 8.0 | ++------+--------------+---------------------+-------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') -{"otel.m"}; + ++--------------+------+------+---------------------+-----------+ +| service.name | Host | host | ts.time | (- v.val) | ++--------------+------+------+---------------------+-----------+ +| a | h1 | x | 1970-01-01T00:00:00 | -1.0 | +| a | h1 | x | 1970-01-01T00:00:05 | -4.0 | +| a | h1 | x | 1970-01-01T00:00:10 | -8.0 | +| a | h2 | x | 1970-01-01T00:00:00 | -2.0 | +| a | h2 | x | 1970-01-01T00:00:05 | -5.0 | +| a | h2 | x | 1970-01-01T00:00:10 | -9.0 | +| b | h1 | y | 1970-01-01T00:00:00 | -3.0 | +| b | h1 | y | 1970-01-01T00:00:05 | -6.0 | +| b | h1 | y | 1970-01-01T00:00:10 | -8.0 | ++--------------+------+------+---------------------+-----------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') 2 * {"otel.m"}; + ++--------------+------+------+---------------------+--------------------+ +| service.name | Host | host | ts.time | Float64(2) * v.val | ++--------------+------+------+---------------------+--------------------+ +| a | h1 | x | 1970-01-01T00:00:00 | 2.0 | +| a | h1 | x | 1970-01-01T00:00:05 | 8.0 | +| a | h1 | x | 1970-01-01T00:00:10 | 16.0 | +| a | h2 | x | 1970-01-01T00:00:00 | 4.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 10.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 18.0 | +| b | h1 | y | 1970-01-01T00:00:00 | 6.0 | +| b | h1 | y | 1970-01-01T00:00:05 | 12.0 | +| b | h1 | y | 1970-01-01T00:00:10 | 16.0 | ++--------------+------+------+---------------------+--------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} - 1; + ++--------------+------+------+---------------------+--------------------+ +| service.name | Host | host | ts.time | v.val - Float64(1) | ++--------------+------+------+---------------------+--------------------+ +| a | h1 | x | 1970-01-01T00:00:00 | 0.0 | +| a | h1 | x | 1970-01-01T00:00:05 | 3.0 | +| a | h1 | x | 1970-01-01T00:00:10 | 7.0 | +| a | h2 | x | 1970-01-01T00:00:00 | 1.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 4.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 8.0 | +| b | h1 | y | 1970-01-01T00:00:00 | 2.0 | +| b | h1 | y | 1970-01-01T00:00:05 | 5.0 | +| b | h1 | y | 1970-01-01T00:00:10 | 7.0 | ++--------------+------+------+---------------------+--------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') timestamp({"otel.m"}); + ++---------------------+-------+--------------+------+------+ +| ts.time | value | service.name | Host | host | ++---------------------+-------+--------------+------+------+ +| 1970-01-01T00:00:00 | 0.0 | a | h1 | x | +| 1970-01-01T00:00:00 | 0.0 | a | h2 | x | +| 1970-01-01T00:00:00 | 0.0 | b | h1 | y | +| 1970-01-01T00:00:05 | 5.0 | a | h1 | x | +| 1970-01-01T00:00:05 | 5.0 | a | h2 | x | +| 1970-01-01T00:00:05 | 5.0 | b | h1 | y | +| 1970-01-01T00:00:10 | 10.0 | a | h1 | x | +| 1970-01-01T00:00:10 | 10.0 | a | h2 | x | +| 1970-01-01T00:00:10 | 10.0 | b | h1 | y | ++---------------------+-------+--------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') minute({"otel.m"}); + ++---------------------+-----------------------------------+--------------+------+------+ +| ts.time | date_part(Utf8("minute"),ts.time) | service.name | Host | host | ++---------------------+-----------------------------------+--------------+------+------+ +| 1970-01-01T00:00:00 | 0 | a | h1 | x | +| 1970-01-01T00:00:00 | 0 | a | h2 | x | +| 1970-01-01T00:00:00 | 0 | b | h1 | y | +| 1970-01-01T00:00:05 | 0 | a | h1 | x | +| 1970-01-01T00:00:05 | 0 | a | h2 | x | +| 1970-01-01T00:00:05 | 0 | b | h1 | y | +| 1970-01-01T00:00:10 | 0 | a | h1 | x | +| 1970-01-01T00:00:10 | 0 | a | h2 | x | +| 1970-01-01T00:00:10 | 0 | b | h1 | y | ++---------------------+-----------------------------------+--------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + ++--------------+---------------------+--------------------------------------------------------------+ +| service.name | ts.time | sum(prom_increase(ts.time_range,v.val,ts.time,Int64(10000))) | ++--------------+---------------------+--------------------------------------------------------------+ +| a | 1970-01-01T00:00:05 | 9.0 | +| a | 1970-01-01T00:00:10 | 16.0 | +| b | 1970-01-01T00:00:05 | 6.0 | +| b | 1970-01-01T00:00:10 | 4.0 | ++--------------+---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') scalar(sum({"otel.m"})) * {"otel.m", "Host"="h2"}; + ++--------------+------+------+---------------------+-------------------------------------------+ +| service.name | Host | host | ts.time | lhs.scalar(sum(otel.m.v.val)) * rhs.v.val | ++--------------+------+------+---------------------+-------------------------------------------+ +| a | h2 | x | 1970-01-01T00:00:00 | 12.0 | +| a | h2 | x | 1970-01-01T00:00:05 | 75.0 | +| a | h2 | x | 1970-01-01T00:00:10 | 225.0 | ++--------------+------+------+---------------------+-------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') absent({"otel.m", "service.name"="c"}); + ++---------------------+-------+--------------+ +| ts.time | v.val | service.name | ++---------------------+-------+--------------+ +| 1970-01-01T00:00:00 | 1.0 | c | +| 1970-01-01T00:00:05 | 1.0 | c | +| 1970-01-01T00:00:10 | 1.0 | c | ++---------------------+-------+--------------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + ++---------------------+--------------+-------+ +| ts.time | service.name | v.val | ++---------------------+--------------+-------+ +| 1970-01-01T00:00:00 | a | 0.55 | +| 1970-01-01T00:00:05 | a | 0.55 | +| 1970-01-01T00:00:10 | a | 0.775 | ++---------------------+--------------+-------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| LeftSemi Join: otel.m.Host = otel.m.Host, otel.m.host = otel.m.host, otel.m.service.name = otel.m.service.name, otel.m.ts.time = otel.m.ts.time_| +|_|_Aggregate: groupBy=[[otel.m.ts.time, otel.m.service.name, otel.m.Host, otel.m.host, otel.m.v.val]], aggr=[[]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Distinct:_| +|_|_PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +|_|_Projection: otel.m.ts.time, otel.m.service.name, otel.m.Host, otel.m.host_| +|_|_PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Filter: otel.m.Host = Utf8("h1") AND otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.Host = Utf8("h1"), otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)] | +|_| ]]_| +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(Host@2, Host@2), (host@3, host@3), (service.name@1, service.name@1), (ts.time@0, ts.time@0)], NullsEqual: true_| +|_|_CoalescePartitionsExec_| +|_|_ProjectionExec: expr=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host]_| +|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivideExec: tags=["service.name", "Host", "host"]_| +|_|_SortExec: expr=[service.name@1 ASC, Host@2 ASC, host@3 ASC, ts.time@0 ASC], preserve_REDACTED +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_AggregateExec: mode=FinalPartitioned, gby=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host, v.val@4 as v.val], aggr=[]_| +|_|_RepartitionExec: REDACTED +|_|_AggregateExec: mode=Partial, gby=[ts.time@0 as ts.time, service.name@1 as service.name, Host@2 as Host, host@3 as host, v.val@4 as v.val], aggr=[]_| +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Projection: otel.m.v.val, otel.m.service.name, otel.m.Host, otel.m.host, otel.m.ts.time_| +|_|_Sort: otel.m.ts.time ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST_| +|_|_Filter: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS Float64) <= Float64(1)_| +|_|_WindowAggr: windowExpr=[[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-299999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-299999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +| physical_plan | ProjectionExec: expr=[v.val@4 as v.val, service.name@1 as service.name, Host@2 as Host, host@3 as host, ts.time@0 as ts.time]_| +|_|_SortPreservingMergeExec: [ts.time@0 ASC NULLS LAST, row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 ASC NULLS LAST]_| +|_|_FilterExec: CAST(row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@5 AS Float64) <= 1_| +|_|_BoundedWindowAggExec: wdw=[row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [otel.m.ts.time] ORDER BY [otel.m.v.val DESC NULLS FIRST, otel.m.service.name DESC NULLS FIRST, otel.m.Host DESC NULLS FIRST, otel.m.host DESC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | +|_|_SortExec: expr=[ts.time@0 ASC NULLS LAST, v.val@4 DESC, service.name@1 DESC, Host@2 DESC, host@3 DESC], preserve_REDACTED +|_|_RepartitionExec: REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| MergeSort: otel.m.service.name ASC NULLS LAST, otel.m.ts.time ASC NULLS LAST_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Sort: otel.m.service.name ASC NULLS LAST, otel.m.ts.time ASC NULLS LAST_| +|_|_Aggregate: groupBy=[[otel.m.service.name, otel.m.ts.time]], aggr=[[sum(prom_increase(ts.time_range,v.val,ts.time,Int64(10000)))]]_| +|_|_Filter: prom_increase(ts.time_range,v.val,ts.time,Int64(10000)) IS NOT NULL_| +|_|_Projection: otel.m.ts.time, prom_increase(ts.time_range, v.val, otel.m.ts.time, Int64(10000)) AS prom_increase(ts.time_range,v.val,ts.time,Int64(10000)), otel.m.service.name, otel.m.Host, otel.m.host | +|_|_PromRangeManipulate: req range=[5000..10000], interval=[5000], eval range=[10000], time index=[ts.time], values=["v.val"]_| +|_|_PromSeriesNormalize: offset=[0], time index=[ts.time], filter NaN: [true]_| +|_|_PromSeriesDivide: tags=["service.name", "Host", "host"]_| +|_|_Sort: otel.m.service.name ASC NULLS FIRST, otel.m.Host ASC NULLS FIRST, otel.m.host ASC NULLS FIRST, otel.m.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.m.ts.time >= TimestampMillisecond(-4999, None) AND otel.m.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.m, partial_filters=[otel.m.ts.time >= TimestampMillisecond(-4999, None), otel.m.ts.time <= TimestampMillisecond(10000, None)]_| +|_| ]]_| +| physical_plan | MergeSortExec: [service.name@0 ASC NULLS LAST, ts.time@1 ASC NULLS LAST]_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| HistogramFold: le=le, field=v.val, quantile=0.5_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts.time]_| +|_|_PromSeriesDivide: tags=["service.name", "le"]_| +|_|_Sort: otel.h.service.name ASC NULLS FIRST, otel.h.le ASC NULLS FIRST, otel.h.ts.time ASC NULLS FIRST_| +|_|_Filter: otel.h.ts.time >= TimestampMillisecond(-299999, None) AND otel.h.ts.time <= TimestampMillisecond(10000, None)_| +|_|_TableScan: otel.h, partial_filters=[otel.h.ts.time >= TimestampMillisecond(-299999, None), otel.h.ts.time <= TimestampMillisecond(10000, None)] | +|_| ]]_| +| physical_plan | HistogramFoldExec: le=@2, field=@3, quantile=0.5_| +|_|_RepartitionExec: REDACTED +|_|_SortExec: expr=[service.name@1 ASC NULLS LAST, ts.time@0 ASC NULLS LAST, TRY_CAST(le@2 AS Float64) ASC NULLS LAST], preserve_REDACTED +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +DROP TABLE "otel.m"; + +Affected Rows: 0 + +DROP TABLE "otel.h"; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/dotted_column_names.sql b/tests/cases/standalone/common/promql/dotted_column_names.sql new file mode 100644 index 00000000000..e6edc71094b --- /dev/null +++ b/tests/cases/standalone/common/promql/dotted_column_names.sql @@ -0,0 +1,139 @@ +-- https://github.com/GreptimeTeam/greptimedb/issues/9390 +-- Tag, field and time index names containing dots or upper case letters must +-- be resolved as plain column names, not as `relation.column`. `Host` and `host` +-- hold different values, so reading one for the other changes the output. +CREATE TABLE "otel.m" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + "Host" STRING, + host STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", "Host", host) +) PARTITION ON COLUMNS ("service.name") ( + "service.name" < 'b', + "service.name" >= 'b' +); + +INSERT INTO "otel.m" VALUES + (0, 'a', 'h1', 'x', 1), + (0, 'a', 'h2', 'x', 2), + (0, 'b', 'h1', 'y', 3), + (5000, 'a', 'h1', 'x', 4), + (5000, 'a', 'h2', 'x', 5), + (5000, 'b', 'h1', 'y', 6), + (10000, 'a', 'h1', 'x', 8), + (10000, 'a', 'h2', 'x', 9), + (10000, 'b', 'h1', 'y', 8); + +CREATE TABLE "otel.h" ( + "ts.time" TIMESTAMP(3) TIME INDEX, + "service.name" STRING, + le STRING, + "v.val" DOUBLE, + PRIMARY KEY ("service.name", le) +); + +INSERT INTO "otel.h" VALUES + (0, 'a', '0.1', 1), + (0, 'a', '1', 3), + (0, 'a', '+Inf', 4), + (10000, 'a', '0.1', 2), + (10000, 'a', '1', 6), + (10000, 'a', '+Inf', 10); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on("service.name") {"otel.m", "service.name"="b"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} and on(host) {"otel.m", host="y"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} unless on("service.name") {"otel.m", "service.name"="b"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} or {"otel.m", "Host"="h1"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values("v.name", {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') count_values by ("service.name") ("v.name", {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') topk(1, {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') bottomk by ("service.name") (1, {"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("service.name") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by ("Host") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum by (host) ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') sum without ("Host") ({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') -{"otel.m"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') 2 * {"otel.m"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') {"otel.m"} - 1; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') timestamp({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') minute({"otel.m"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') scalar(sum({"otel.m"})) * {"otel.m", "Host"="h2"}; + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') absent({"otel.m", "service.name"="c"}); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') {"otel.m"} and {"otel.m", "Host"="h1"}; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') topk(1, {"otel.m"}); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (5, 10, '5s') sum by ("service.name") (increase({"otel.m"}[10s])); + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +TQL EXPLAIN (0, 10, '5s') histogram_quantile(0.5, {"otel.h"}); + +DROP TABLE "otel.m"; + +DROP TABLE "otel.h"; From 3131cdbcb6fc05b52897426ca4b8c21405ff3668 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 29 Sep 2026 10:07:18 +0000 Subject: [PATCH 077/133] fix: bound decompressed request size and close memory-admission gaps for compressed requests (#9264) * fix: bound decompressed request size and close memory-admission gaps for compressed requests The request-memory accounting only bounds and charges the encoded bytes on the wire, but a compressed body can expand far beyond that during decompression, so tiny requests could allocate disproportionate frontend memory before any protobuf validation or quota charge. - Handler-level decompression (Prometheus remote read/write v1+v2, Loki) now enforces a hard 512 MiB decoded-size cap, checked before any output buffer is allocated, and charges the decoded bytes to the shared ServerMemoryLimiter, holding the permits for the lifetime of the decompressed buffer. - gRPC requests with transport compression reserve the configured max_recv_message_size before tonic decompresses, so the decoding phase is admitted against max_in_flight_write_bytes; the later per-message charge is skipped to avoid double accounting. - The HTTP memory-limit middleware keeps its upfront Content-Length charge but now also accounts the bytes actually streamed beyond it, so chunked requests and understated Content-Length headers no longer bypass the aggregate quota. - Routes that decompress request bodies via RequestDecompressionLayer (InfluxDB, OTLP, Loki, Splunk, Elasticsearch, pipelines, dashboards) now charge the decompressed bytes as handlers consume them: the global middleware marks Content-Encoding requests, and a route-local accounting layer inside the decompression layer charges the decoded stream. Plain requests are skipped to avoid double-counting. Signed-off-by: Ning Sun * fix: address review comments on memory admission guard lifetimes and 429 mapping - Retain the gRPC pre-decode reservation for the whole request: the extensions holding the guard were dropped when the request was consumed (into_inner), releasing the reservation while per-message charges stayed skipped. Both the unary and streaming handlers now clone and hold the reservation marker for the duration of request handling. - Retain the HTTP body permits across the handler: the AccountedBody wrapper and the request extensions are dropped once the extractors finish collecting the body, before the handler is done with the decoded data. Both accounting middlewares now keep their own accounting handle alive across next.run(req).await. - Return a ChargedBuffer from the Loki snappy decompressor so the reservation outlives the decompressed bytes through the caller's protobuf decoding, matching the Prometheus path. - Map mid-stream quota exhaustion to 429 instead of the generic body error (400): the accounting flags exhaustion and the middlewares rewrite the extractor rejection, so clients can distinguish backpressure from malformed input. Signed-off-by: Ning Sun * fix: resolve the in-flight body acquisition before trying a new one A parked acquisition in AccountedBody::charge always belongs to the currently buffered frame, so it must be resolved before any new acquisition is tried. Letting the fast-path try_acquire succeed while a waiter is parked left the waiter alive, and its late completion would then be credited with a later frame's byte count, under-reserving memory relative to what was marked charged. Adds a regression test that reproduces the misattribution: with the buggy ordering the test delivers a body the quota cannot cover; with the fix the over-quota frame waits for its own acquisition and times out. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- Cargo.lock | 1 + src/servers/Cargo.toml | 1 + src/servers/src/error.rs | 13 + src/servers/src/grpc.rs | 18 +- src/servers/src/grpc/builder.rs | 19 +- src/servers/src/grpc/database.rs | 33 +- src/servers/src/grpc/memory_limit.rs | 240 +++++++- src/servers/src/http.rs | 81 ++- src/servers/src/http/loki.rs | 149 ++++- src/servers/src/http/memory_limit.rs | 600 ++++++++++++++++++- src/servers/src/http/prom_store.rs | 38 +- src/servers/src/prom_remote_write/mod.rs | 83 ++- src/servers/src/prom_remote_write/v2.rs | 72 ++- src/servers/src/prom_store.rs | 270 ++++++++- src/servers/src/request_memory_limiter.rs | 5 + tests-integration/src/test_util.rs | 51 +- tests-integration/src/tests/instance_test.rs | 2 +- tests-integration/tests/grpc.rs | 99 +++ tests-integration/tests/http.rs | 188 +++++- 19 files changed, 1827 insertions(+), 136 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6997f0d782b..834b0b541bb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13589,6 +13589,7 @@ dependencies = [ "hex", "hostname 0.3.1", "http 1.5.0", + "http-body 1.0.1", "humantime", "humantime-serde", "hyper 1.6.0", diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index c1af56a6281..f9c3b45832f 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -70,6 +70,7 @@ headers = "0.4" hex.workspace = true hostname = "0.3" http.workspace = true +http-body = "1" humantime.workspace = true humantime-serde.workspace = true hyper = { workspace = true, features = ["full"] } diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index 47f9953a187..af6b2499e7f 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -319,6 +319,18 @@ pub enum Error { error: std::io::Error, }, + #[snafu(display( + "Decompressed request body is too large: {} bytes exceeds the limit {} bytes", + size, + limit + ))] + DecompressedBodyTooLarge { + size: u64, + limit: u64, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to compress prometheus remote request"))] CompressPromRemoteRequest { #[snafu(implicit)] @@ -785,6 +797,7 @@ impl ErrorExt for Error { | DecompressSnappyPromRemoteRequest { .. } | DecompressSnappyLokiRequest { .. } | DecompressZstdPromRemoteRequest { .. } + | DecompressedBodyTooLarge { .. } | InvalidPromRemoteRequest { .. } | InvalidFlightTicket { .. } | InvalidPrepareStatement { .. } diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs index 50dd0b69c4a..1be1f8215b5 100644 --- a/src/servers/src/grpc.rs +++ b/src/servers/src/grpc.rs @@ -50,6 +50,7 @@ use tonic::{Request, Response, Status}; use tonic_reflection::server::v1::{ServerReflection, ServerReflectionServer}; use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu}; +use crate::grpc::memory_limit::MemoryLimiterExtensionService; use crate::install_default_crypto_provider; use crate::metrics::MetricsMiddlewareLayer; use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler}; @@ -213,6 +214,14 @@ impl FlightCompression { } } +/// The wrapped OTLP Arrow service type used by [`GrpcServer`]. +type OtelArrowService = MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, +>; + pub struct GrpcServer { // states shutdown_tx: Mutex>>, @@ -224,14 +233,7 @@ pub struct GrpcServer { // tls config tls_config: Option, // Otel arrow service - otel_arrow_service: Mutex< - Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, - >, - >, - >, + otel_arrow_service: Mutex>, bind_addr: Option, name: Option, config: GrpcServerConfig, diff --git a/src/servers/src/grpc/builder.rs b/src/servers/src/grpc/builder.rs index d240cde3f1b..99513dcceb3 100644 --- a/src/servers/src/grpc/builder.rs +++ b/src/servers/src/grpc/builder.rs @@ -34,12 +34,13 @@ use tonic::codegen::Service; use tonic::service::RoutesBuilder; use tonic::service::interceptor::InterceptedService; use tonic::transport::{Identity, ServerTlsConfig}; -use tower::Layer; +use tower::{Layer, ServiceBuilder}; use crate::grpc::database::DatabaseService; use crate::grpc::flight::{FlightCraftRef, FlightCraftWrapper}; use crate::grpc::frontend_grpc_handler::FrontendGrpcHandler; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::{MemoryLimiterExtensionLayer, MemoryLimiterExtensionService}; use crate::grpc::prom_query_gateway::PrometheusGatewayService; use crate::grpc::region_server::{RegionServerHandlerRef, RegionServerRequestHandler}; use crate::grpc::{GrpcServer, GrpcServerConfig}; @@ -71,6 +72,7 @@ macro_rules! add_service { let service_with_limiter = $crate::tower::ServiceBuilder::new() .layer(MemoryLimiterExtensionLayer::new( $builder.memory_limiter().clone(), + max_recv_message_size, )) .service(service_builder); @@ -87,9 +89,11 @@ pub struct GrpcServerBuilder { routes_builder: RoutesBuilder, tls_config: Option, otel_arrow_service: Option< - InterceptedService< - ArrowMetricsServiceServer>, - HeaderInterceptor, + MemoryLimiterExtensionService< + InterceptedService< + ArrowMetricsServiceServer>, + HeaderInterceptor, + >, >, >, memory_limiter: ServerMemoryLimiter, @@ -185,6 +189,13 @@ impl GrpcServerBuilder { .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Zstd); let svc = InterceptedService::new(server, HeaderInterceptor {}); + // Same pre-decode memory admission as `add_service!`. + let svc = ServiceBuilder::new() + .layer(MemoryLimiterExtensionLayer::new( + self.memory_limiter.clone(), + self.config.max_recv_message_size, + )) + .service(svc); self.otel_arrow_service = Some(svc); self } diff --git a/src/servers/src/grpc/database.rs b/src/servers/src/grpc/database.rs index 2a63a982bef..237d760269f 100644 --- a/src/servers/src/grpc/database.rs +++ b/src/servers/src/grpc/database.rs @@ -25,6 +25,7 @@ use session::context::Channel; use tonic::{Request, Response, Status, Streaming}; use crate::grpc::greptime_handler::GreptimeRequestHandler; +use crate::grpc::memory_limit::PreDecodeMemoryReservation; use crate::grpc::{TonicResult, cancellation}; use crate::hint_headers; use crate::request_memory_limiter::ServerMemoryLimiter; @@ -57,7 +58,20 @@ impl GreptimeDatabase for DatabaseService { remote_addr, hints ); - let _guard = if let Some(limiter) = request.extensions().get::() { + // Retain the pre-decode reservation for the whole request: the + // extension holding the guard would be dropped when the request is + // consumed below, but the post-decode charge is skipped while it is + // active. + let _pre_reservation = request + .extensions() + .get::() + .cloned(); + let _guard = if _pre_reservation.is_some() { + // Compressed requests already reserved the worst-case decoded + // size before tonic decompressed the message; skip the exact + // post-decode charge to avoid double accounting. + None + } else if let Some(limiter) = request.extensions().get::() { let message_size = request.get_ref().encoded_len() as u64; Some(limiter.acquire(message_size).await?) } else { @@ -117,16 +131,31 @@ impl GreptimeDatabase for DatabaseService { ); let limiter = request.extensions().get::().cloned(); + // For compressed streams the whole stream's decoding memory was + // reserved before tonic started decompressing; messages are decoded + // one at a time, so the reservation covers each of them. The + // reservation is retained below for the whole stream: the extension + // holding the guard would otherwise be dropped when the request is + // consumed, while per-message charges stay skipped. + let reservation = request + .extensions() + .get::() + .cloned(); + let pre_reserved = reservation.is_some(); let handler = self.handler.clone(); let request_future = async move { let mut affected_rows = 0; + // Hold the pre-decode reservation until the stream is exhausted. + let _reservation = reservation; let mut stream = request.into_inner(); while let Some(request) = stream.next().await { let request = request?; - let _guard = if let Some(limiter_ref) = &limiter { + let _guard = if pre_reserved { + None + } else if let Some(limiter_ref) = &limiter { let message_size = request.encoded_len() as u64; Some(limiter_ref.acquire(message_size).await?) } else { diff --git a/src/servers/src/grpc/memory_limit.rs b/src/servers/src/grpc/memory_limit.rs index 722a2c68c08..c9c62e88dd1 100644 --- a/src/servers/src/grpc/memory_limit.rs +++ b/src/servers/src/grpc/memory_limit.rs @@ -12,22 +12,69 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Aggregate memory admission for gRPC services. +//! +//! Tonic decompresses and decodes a complete request message (up to +//! `max_recv_message_size`, 512 MiB by default) *before* the typed handler is +//! invoked, so a handler that charges the [`ServerMemoryLimiter`] only sees +//! the request after the memory is already allocated. For transport-compressed +//! messages (`grpc-encoding: gzip/zstd`) a few KiB on the wire can expand to +//! hundreds of MiB during that window, which would bypass any aggregate quota. +//! +//! [`MemoryLimiterExtensionLayer`] therefore reserves the configured maximum +//! decoded message size for every compressed request *before* the inner +//! (tonic) service runs, and holds the reservation for the whole request. +//! Handlers can detect the reservation via [`PreDecodeMemoryReservation`] in +//! the request extensions and skip their own post-decode charge to avoid +//! double accounting. +//! +//! The reservation is intentionally conservative (it upper-bounds the decoded +//! size before it is known); with the default unlimited limiter it is a no-op. + +use std::convert::Infallible; +use std::sync::Arc; use std::task::{Context, Poll}; +use axum::response::IntoResponse; +use common_memory_manager::MemoryGuard; use futures::future::BoxFuture; +use http::Request; +use tonic::Status; +use tonic::codegen::Service; use tonic::server::NamedService; -use tower::{Layer, Service}; +use tower::Layer; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// The gRPC request compression selector header. +const GRPC_ENCODING_HEADER: &str = "grpc-encoding"; +/// The "no compression" encoding value. +const IDENTITY_ENCODING: &str = "identity"; + +/// Present in the request extensions when memory for the (compressed) +/// request's decoded message was reserved before tonic decompression. +/// +/// Handlers that charge the [`ServerMemoryLimiter`] after decoding should skip +/// that charge when this marker is present: the reservation already covers the +/// peak decoding memory and stays alive for the whole request. +#[derive(Clone)] +pub(crate) struct PreDecodeMemoryReservation { + _guard: Arc>, +} #[derive(Clone)] pub struct MemoryLimiterExtensionLayer { limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl MemoryLimiterExtensionLayer { - pub fn new(limiter: ServerMemoryLimiter) -> Self { - Self { limiter } + pub fn new(limiter: ServerMemoryLimiter, max_decoding_message_size: usize) -> Self { + Self { + limiter, + max_decoding_message_size, + } } } @@ -38,6 +85,7 @@ impl Layer for MemoryLimiterExtensionLayer { MemoryLimiterExtensionService { inner: service, limiter: self.limiter.clone(), + max_decoding_message_size: self.max_decoding_message_size, } } } @@ -46,27 +94,199 @@ impl Layer for MemoryLimiterExtensionLayer { pub struct MemoryLimiterExtensionService { inner: S, limiter: ServerMemoryLimiter, + max_decoding_message_size: usize, } impl NamedService for MemoryLimiterExtensionService { const NAME: &'static str = S::NAME; } -impl Service> for MemoryLimiterExtensionService +impl Service> for MemoryLimiterExtensionService where - S: Service>, + S: Service, Error = Infallible> + Clone + Send + 'static, + S::Response: axum::response::IntoResponse, S::Future: Send + 'static, + ReqBody: Send + 'static, { - type Response = S::Response; - type Error = S::Error; + type Response = axum::response::Response; + type Error = Infallible; type Future = BoxFuture<'static, Result>; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } - fn call(&mut self, mut req: http::Request) -> Self::Future { - req.extensions_mut().insert(self.limiter.clone()); - Box::pin(self.inner.call(req)) + fn call(&mut self, mut req: Request) -> Self::Future { + // Own a clone of the service so the returned future does not borrow + // `self` (tower's load-shedding pattern). + let mut this = self.clone(); + Box::pin(async move { + req.extensions_mut().insert(this.limiter.clone()); + + let compressed = req + .headers() + .get(GRPC_ENCODING_HEADER) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| !value.eq_ignore_ascii_case(IDENTITY_ENCODING)); + + if compressed { + // A compressed message can expand up to the decoding limit + // inside tonic, before any handler runs. Reserve that worst + // case against the aggregate quota so the decoding phase is + // admitted as well. No-op for an unlimited limiter. + let reservation = this.max_decoding_message_size as u64; + match this.limiter.acquire(reservation).await { + Ok(guard) => { + req.extensions_mut().insert(PreDecodeMemoryReservation { + _guard: Arc::new(guard), + }); + } + Err(e) => { + return Ok(Status::resource_exhausted(format!( + "request memory limit exceeded: {e}" + )) + .into_http::() + .into_response()); + } + } + } + + match this.inner.call(req).await { + Ok(response) => Ok(response.into_response()), + Err(e) => match e {}, + } + }) + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use axum::body::Body; + use axum::response::IntoResponse; + use common_memory_manager::OnExhaustedPolicy; + use futures_util::future::BoxFuture; + use http::{HeaderValue, StatusCode}; + use tower::ServiceExt; + + use super::*; + use crate::grpc::memory_limit::MemoryLimiterExtensionLayer; + + #[derive(Clone)] + struct EchoService; + + impl Service> for EchoService { + type Response = axum::response::Response; + type Error = Infallible; + type Future = BoxFuture<'static, Result>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + let saw_reservation = req + .extensions() + .get::() + .is_some(); + let saw_limiter = req.extensions().get::().is_some(); + Box::pin(async move { + Ok(( + StatusCode::OK, + format!("reservation={saw_reservation} limiter={saw_limiter}"), + ) + .into_response()) + }) + } + } + + impl NamedService for EchoService { + const NAME: &'static str = "test.Echo"; + } + + #[tokio::test] + async fn test_inserts_limiter_for_uncompressed_requests() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder().body(Body::empty()).unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=false limiter=true"[..], + "uncompressed requests must not be pre-reserved" + ); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_reserves_max_message_size_for_compressed_requests() { + let max_size = 512 * 1024 * 1024; + let limiter = ServerMemoryLimiter::new(1024 * 1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), max_size).layer(EchoService); + + // The inner service observes the reservation; use a oneshot call and + // check the limiter was charged while the call is in flight via the + // response extensions observed by the echo service. + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("gzip")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!( + &body[..], + &b"reservation=true limiter=true"[..], + "compressed requests must carry a pre-decode reservation" + ); + // The reservation is released once the inner call (request) finishes. + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_rejects_compressed_request_when_quota_exhausted() { + // Quota smaller than the decoding limit: the reservation cannot fit. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("zstd")) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + // gRPC errors are HTTP 200 with a grpc-status header. + assert_eq!( + res.headers() + .get("grpc-status") + .and_then(|v| v.to_str().ok()), + Some("8"), + "expected RESOURCE_EXHAUSTED grpc status" + ); + } + + #[tokio::test] + async fn test_identity_encoding_is_not_reserved() { + let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail); + let mut svc = + MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService); + + let req = Request::builder() + .header( + GRPC_ENCODING_HEADER, + HeaderValue::from_static(IDENTITY_ENCODING), + ) + .body(Body::empty()) + .unwrap(); + let res = svc.ready().await.unwrap().call(req).await.unwrap(); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], &b"reservation=false limiter=true"[..]); } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index b9fcd0a9420..a33e52a1d40 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -716,7 +716,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/influxdb"), - HttpServer::route_influxdb(handler), + HttpServer::route_influxdb(handler, self.memory_limiter.clone()), ), ..self } @@ -736,6 +736,7 @@ impl HttpServerBuilder { prom_store_with_metric_engine, prom_validation_mode, pending_rows_batcher, + memory_limiter: self.memory_limiter.clone(), }; Self { @@ -765,7 +766,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp(handler, with_metric_engine), + HttpServer::route_otlp(handler, with_metric_engine, self.memory_limiter.clone()), ), ..self } @@ -799,23 +800,23 @@ impl HttpServerBuilder { let router = self.router.nest( &format!("/{HTTP_API_VERSION}"), - HttpServer::route_pipelines(log_state.clone()), + HttpServer::route_pipelines(log_state.clone(), self.memory_limiter.clone()), ); // deprecated since v0.11.0. Use `/logs` and `/pipelines` instead. let router = router.nest( &format!("/{HTTP_API_VERSION}/events"), #[allow(deprecated)] - HttpServer::route_log_deprecated(log_state.clone()), + HttpServer::route_log_deprecated(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/loki"), - HttpServer::route_loki(log_state.clone()), + HttpServer::route_loki(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( &format!("/{HTTP_API_VERSION}/elasticsearch"), - HttpServer::route_elasticsearch(log_state.clone()), + HttpServer::route_elasticsearch(log_state.clone(), self.memory_limiter.clone()), ); let router = router.nest( @@ -827,7 +828,7 @@ impl HttpServerBuilder { let router = router.nest( &format!("/{HTTP_API_VERSION}/splunk"), - HttpServer::route_splunk(log_state), + HttpServer::route_splunk(log_state, self.memory_limiter.clone()), ); Self { router, ..self } @@ -858,7 +859,7 @@ impl HttpServerBuilder { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/dashboards"), - HttpServer::route_dashboard(handler), + HttpServer::route_dashboard(handler, self.memory_limiter.clone()), ), ..self } @@ -1184,9 +1185,13 @@ impl HttpServer { .with_state(metrics_handler) } - fn route_loki(log_state: LogState) -> Router { + fn route_loki(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/api/v1/push", routing::post(loki::loki_ingest)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1198,7 +1203,7 @@ impl HttpServer { .with_state(log_state) } - fn route_splunk(log_state: LogState) -> Router { + fn route_splunk(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route( "/services/collector/health", @@ -1226,6 +1231,10 @@ impl HttpServer { "/services/collector/raw/1.0", routing::post(splunk::handle_raw), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1237,7 +1246,10 @@ impl HttpServer { .with_state(log_state) } - fn route_elasticsearch(log_state: LogState) -> Router { + fn route_elasticsearch( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() // Return fake responsefor HEAD '/' request. .route( @@ -1309,6 +1321,10 @@ impl HttpServer { axum::Json(serde_json::json!({})), )), ) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer(ServiceBuilder::new().layer(RequestDecompressionLayer::new())) .layer(middleware::from_fn_with_state( BatchingProtocol::Elasticsearch, @@ -1318,7 +1334,10 @@ impl HttpServer { } #[deprecated(since = "0.11.0", note = "Use `route_pipelines()` instead.")] - fn route_log_deprecated(log_state: LogState) -> Router { + fn route_log_deprecated( + log_state: LogState, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/logs", routing::post(event::log_ingester)) .route( @@ -1334,6 +1353,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1345,7 +1368,7 @@ impl HttpServer { .with_state(log_state) } - fn route_pipelines(log_state: LogState) -> Router { + fn route_pipelines(log_state: LogState, memory_limiter: ServerMemoryLimiter) -> Router { Router::new() .route("/ingest", routing::post(event::log_ingester)) .route( @@ -1365,6 +1388,10 @@ impl HttpServer { routing::delete(event::delete_pipeline), ) .route("/pipelines/_dryrun", routing::post(event::pipeline_dryrun)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1465,10 +1492,17 @@ impl HttpServer { .with_state(state) } - fn route_influxdb(influxdb_handler: InfluxdbLineProtocolHandlerRef) -> Router { + fn route_influxdb( + influxdb_handler: InfluxdbLineProtocolHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new() .route("/write", routing::post(influxdb_write_v1)) .route("/api/v2/write", routing::post(influxdb_write_v2)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1495,11 +1529,16 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + memory_limiter: ServerMemoryLimiter, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) .route("/v1/traces", routing::post(otlp::traces)) .route("/v1/logs", routing::post(otlp::logs)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1540,13 +1579,20 @@ impl HttpServer { } #[cfg(feature = "dashboard")] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + memory_limiter: ServerMemoryLimiter, + ) -> Router { use crate::http::dashboard::{add_dashboard, delete_dashboard, list_dashboards}; Router::new() .route("/", routing::get(list_dashboards)) .route("/{dashboard_name}", routing::post(add_dashboard)) .route("/{dashboard_name}", routing::delete(delete_dashboard)) + .layer(middleware::from_fn_with_state( + memory_limiter, + memory_limit::decoded_body_accounting_middleware, + )) .layer( ServiceBuilder::new() .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)), @@ -1555,7 +1601,10 @@ impl HttpServer { } #[cfg(not(feature = "dashboard"))] - fn route_dashboard(handler: DashboardHandlerRef) -> Router { + fn route_dashboard( + handler: DashboardHandlerRef, + _memory_limiter: ServerMemoryLimiter, + ) -> Router { Router::new().with_state(DashboardState { handler }) } } diff --git a/src/servers/src/http/loki.rs b/src/servers/src/http/loki.rs index 4e12eeb931e..161664a2ce1 100644 --- a/src/servers/src/http/loki.rs +++ b/src/servers/src/http/loki.rs @@ -108,6 +108,7 @@ pub async fn loki_ingest( TypedHeader(content_type): TypedHeader, LogTableName(table_name): LogTableName, pipeline_info: PipelineInfo, + Extension(memory_limiter): Extension, bytes: Bytes, ) -> Result { ctx.set_channel(Channel::Loki); @@ -127,6 +128,7 @@ pub async fn loki_ingest( pipeline_info, bytes, &ctx, + &memory_limiter, ) .await?; @@ -182,6 +184,7 @@ async fn build_loki_context_req( pipeline_info: PipelineInfo, bytes: Bytes, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { // A pipeline header switches Loki into the generic pipeline path; without // it, Loki writes directly to the target log table. @@ -195,10 +198,12 @@ async fn build_loki_context_req( pipeline_params: pipeline_info.pipeline_params, bytes, }; - build_loki_pipeline_context_req(handler, pipeline_req, ctx).await + build_loki_pipeline_context_req(handler, pipeline_req, ctx, memory_limiter).await } None => { - let req = build_loki_raw_insert_request(content_type, table_name, bytes)?; + let req = + build_loki_raw_insert_request(content_type, table_name, bytes, memory_limiter) + .await?; Ok(ContextReq::default_opt_with_reqs(vec![req])) } } @@ -208,6 +213,7 @@ async fn build_loki_pipeline_context_req( handler: &PipelineHandlerRef, pipeline_req: LokiPipelineContextReq, ctx: &QueryContextRef, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let LokiPipelineContextReq { content_type, @@ -223,7 +229,8 @@ async fn build_loki_pipeline_context_req( PipelineDefinition::from_name(&pipeline_name, version, None).context(PipelineSnafu)?; let pipeline_ctx = PipelineContext::new(&def, &pipeline_params, Channel::Loki); - let values = extract_item::(content_type, bytes)? + let values = extract_item::(content_type, bytes, memory_limiter) + .await? .map(|item| item.map) .collect::>(); @@ -235,14 +242,16 @@ async fn build_loki_pipeline_context_req( run_pipeline(handler, &pipeline_ctx, req, ctx, true).await } -fn build_loki_raw_insert_request( +async fn build_loki_raw_insert_request( content_type: ContentType, table_name: String, bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, ) -> Result { let mut schema_info = SchemaInfo::from_schema_list(LOKI_INIT_SCHEMAS.clone()); let mut rows = Vec::with_capacity(256); - for loki_row in extract_item::(content_type, bytes)? { + let loki_rows = extract_item::(content_type, bytes, memory_limiter).await?; + for loki_row in loki_rows { let mut row = init_row( schema_info.schema.len(), loki_row.ts, @@ -275,7 +284,11 @@ fn build_loki_raw_insert_request( /// JSON push bodies become `LokiMiddleItem`, protobuf push bodies /// become `LokiMiddleItem>`, and the generic `Into` /// conversion selects either direct-write `LokiRawItem` or pipeline `LokiPipeline`. -fn extract_item(content_type: ContentType, bytes: Bytes) -> Result>> +async fn extract_item( + content_type: ContentType, + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result>> where LokiMiddleItem: Into, LokiMiddleItem>: Into, @@ -285,7 +298,9 @@ where LokiJsonParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), )), x if x == *PB_CONTENT_TYPE => Ok(Box::new( - LokiPbParser::from_bytes(bytes)?.flat_map(|item| item.into_iter().map(|i| i.into())), + LokiPbParser::from_bytes(bytes, memory_limiter) + .await? + .flat_map(|item| item.into_iter().map(|i| i.into())), )), _ => UnsupportedContentTypeSnafu { content_type }.fail(), } @@ -554,8 +569,13 @@ pub struct LokiPbParser { } impl LokiPbParser { - pub fn from_bytes(bytes: Bytes) -> Result { - let decompressed = snappy_decompress_loki_request(&bytes)?; + pub async fn from_bytes( + bytes: Bytes, + memory_limiter: &crate::request_memory_limiter::ServerMemoryLimiter, + ) -> Result { + // `decompressed` carries the memory permits for the decoded bytes and + // keeps them alive until the protobuf decoding below is finished. + let decompressed = snappy_decompress_loki_request(&bytes, memory_limiter).await?; let req = loki_proto::logproto::PushRequest::decode(&decompressed[..]) .context(DecodeLokiRequestSnafu)?; @@ -565,13 +585,34 @@ impl LokiPbParser { } } -fn snappy_decompress_loki_request(buf: &[u8]) -> Result> { +async fn snappy_decompress_loki_request( + buf: &[u8], + limiter: &crate::request_memory_limiter::ServerMemoryLimiter, +) -> Result { // Loki's protobuf push body is Snappy-compressed independent of HTTP // content-encoding, so keep this decode step explicit. + // + // The decoded size is validated before the output buffer is allocated: + // the raw Snappy header carries the decoded length, and without this cap + // a tiny body could declare multiple GiB of output. The declared size is + // also charged to the aggregate request-memory limiter before the + // allocation, so concurrent Loki pushes cannot bypass the quota. The + // returned buffer carries the permits, keeping the reservation alive + // until the caller is done with the decoded bytes. + let decoded_len = snap::raw::decompress_len(buf).context(DecompressSnappyLokiRequestSnafu)?; + ensure!( + decoded_len <= crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE, + crate::error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: crate::prom_store::MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); + let guard = limiter.acquire(decoded_len as u64).await?; let mut decoder = Decoder::new(); - decoder + let data = decoder .decompress_vec(buf) - .context(DecompressSnappyLokiRequestSnafu) + .context(DecompressSnappyLokiRequestSnafu)?; + Ok(crate::prom_store::ChargedBuffer::new(data, vec![guard])) } impl Iterator for LokiPbParser { @@ -819,6 +860,7 @@ mod tests { use super::*; use crate::error::Error::{DecompressSnappyLokiRequest, InvalidLokiLabels}; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; const JSON_PAYLOAD: &[u8] = br#"{ "streams": [ @@ -869,13 +911,15 @@ mod tests { assert_eq!(prost_ts_to_nano(&ts), 1731748568804293888); } - #[test] - fn test_json_direct_ingest_builds_schema_and_pads_rows() { + #[tokio::test] + async fn test_json_direct_ingest_builds_schema_and_pads_rows() { let request = build_loki_raw_insert_request( JSON_CONTENT_TYPE.clone(), "custom_loki".to_string(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap(); assert_eq!(request.table_name, "custom_loki"); @@ -916,12 +960,14 @@ mod tests { assert_eq!(row_string_value(second, 5), Some("worker-0")); } - #[test] - fn test_json_pipeline_conversion_names_loki_fields() { + #[tokio::test] + async fn test_json_pipeline_conversion_names_loki_fields() { let items = extract_item::( JSON_CONTENT_TYPE.clone(), Bytes::from_static(JSON_PAYLOAD), + &ServerMemoryLimiter::default(), ) + .await .unwrap() .collect::>(); @@ -951,8 +997,8 @@ mod tests { ); } - #[test] - fn test_protobuf_parser_decodes_snappy_push_request() { + #[tokio::test] + async fn test_protobuf_parser_decodes_snappy_push_request() { let request = PushRequest { streams: vec![StreamAdapter { labels: r#"{job="api"}"#.to_string(), @@ -973,9 +1019,14 @@ mod tests { }; let bytes = snappy_compress(&request.encode_to_vec()).unwrap(); - let items = extract_item::(PB_CONTENT_TYPE.clone(), Bytes::from(bytes)) - .unwrap() - .collect::>(); + let items = extract_item::( + PB_CONTENT_TYPE.clone(), + Bytes::from(bytes), + &ServerMemoryLimiter::default(), + ) + .await + .unwrap() + .collect::>(); assert_eq!(items.len(), 1); assert_eq!(items[0].ts, 1731748568804293888); @@ -987,9 +1038,14 @@ mod tests { assert!(!items[0].structured_metadata.is_empty()); } - #[test] - fn test_protobuf_parser_rejects_invalid_snappy_payload() { - let err = match LokiPbParser::from_bytes(Bytes::from_static(b"not-snappy")) { + #[tokio::test] + async fn test_protobuf_parser_rejects_invalid_snappy_payload() { + let err = match LokiPbParser::from_bytes( + Bytes::from_static(b"not-snappy"), + &ServerMemoryLimiter::default(), + ) + .await + { Ok(_) => panic!("expected invalid snappy payload to fail"), Err(err) => err, }; @@ -997,6 +1053,51 @@ mod tests { assert!(matches!(err, DecompressSnappyLokiRequest { .. })); } + #[tokio::test] + async fn test_protobuf_push_charges_decoded_size() { + use common_memory_manager::OnExhaustedPolicy; + + // A valid snappy payload whose decoded size exceeds the 1 KiB quota. + let request = loki_proto::logproto::PushRequest { + streams: vec![loki_proto::logproto::StreamAdapter { + labels: "{job=\"quota\"}".to_string(), + entries: std::iter::repeat_with(|| loki_proto::logproto::EntryAdapter { + timestamp: Some(Timestamp { + seconds: 1, + nanos: 0, + }), + line: "x".repeat(256), + ..Default::default() + }) + .take(8) + .collect(), + ..Default::default() + }], + }; + let bytes = Bytes::from(snappy_compress(&request.encode_to_vec()).unwrap()); + + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let err = match LokiPbParser::from_bytes(bytes.clone(), &limiter).await { + Ok(_) => panic!("expected quota exhaustion to fail"), + Err(err) => err, + }; + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + + // The same payload fits an adequately sized quota. + let limiter = ServerMemoryLimiter::new(64 * 1024, OnExhaustedPolicy::Fail); + let parser = LokiPbParser::from_bytes(bytes, &limiter).await.unwrap(); + assert_eq!(1, parser.streams.len()); + assert_eq!( + 0, + limiter.used_bytes(), + "guards must be released after decode" + ); + } + #[test] fn test_parse_loki_labels() { let mut expected = BTreeMap::new(); diff --git a/src/servers/src/http/memory_limit.rs b/src/servers/src/http/memory_limit.rs index 50b4c68f03e..8d7b99c15fa 100644 --- a/src/servers/src/http/memory_limit.rs +++ b/src/servers/src/http/memory_limit.rs @@ -13,13 +13,47 @@ // limitations under the License. //! Middleware for limiting total memory usage of concurrent HTTP request bodies. +//! +//! Admission works in two stages: +//! +//! 1. Upfront, the request's `Content-Length` (if present) is charged to the +//! shared [`ServerMemoryLimiter`], preserving the configured wait/fail +//! policy for regular requests. +//! 2. While the body is streamed, [`AccountedBody`] charges every byte beyond +//! the upfront reservation. This closes two gaps of a header-only +//! reservation: HTTP/1.1 chunked requests carry no `Content-Length` (and +//! would otherwise be admitted for free), and a client can understate the +//! header while sending a much larger body. +//! +//! Permits acquired while streaming are held by [`BodyMemoryAccounting`], +//! which is also inserted into the request extensions so they stay alive until +//! the whole request is finished (extractors collect the body into owned +//! buffers that outlive the body stream itself). +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use axum::body::Body; use axum::extract::{Request, State}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use common_memory_manager::MemoryGuard; use http::StatusCode; +use http_body::{Body as HttpBody, Frame}; +use crate::error::Result; use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; + +/// Marker inserted by [`memory_limit_middleware`] when the request arrives +/// with a non-identity `Content-Encoding` header. Route-local +/// [`decoded_body_accounting_middleware`] uses it to account the decompressed +/// body: only such requests expand during request decompression. +#[derive(Clone, Copy)] +pub(crate) struct ContentEncoded; pub async fn memory_limit_middleware( State(limiter): State, @@ -44,5 +78,569 @@ pub async fn memory_limit_middleware( } }; - next.run(req).await + let content_encoded = req + .headers() + .get(http::header::CONTENT_ENCODING) + .and_then(|v| v.to_str().ok()) + .is_some_and(|v| !v.eq_ignore_ascii_case("identity")); + + // Account whatever is actually streamed beyond the upfront reservation + // (chunked transfer has none). + let accounting = BodyMemoryAccounting::default(); + // The middleware keeps its own handle so the acquired permits stay held + // for the whole request: the body wrapper and the extensions are dropped + // once the extractors finish collecting the body, before the handler is + // done with the decoded data. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + // Expose the limiter to handlers that decompress request bodies on their + // own (e.g. the Loki protobuf push) so they can charge the decoded size. + parts.extensions.insert(limiter.clone()); + parts.extensions.insert(accounting.clone()); + if content_encoded { + parts.extensions.insert(ContentEncoded); + } + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Rewrites the response of a request whose body accounting hit the quota: +/// the failure aborts the body mid-stream, so the extractors reject the +/// request with a generic body error (mapped to 400). Quota exhaustion must +/// surface as 429, matching the upfront admission path, so clients can tell +/// backpressure apart from malformed input. +fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response) -> Response { + // The flag can only be set when the body failed before the handler ran, + // so the response is the extractor rejection and can be replaced. + if accounting.take_quota_exceeded() { + return ( + StatusCode::TOO_MANY_REQUESTS, + "Request body memory limit exceeded", + ) + .into_response(); + } + response +} + +/// Route-local counterpart of [`memory_limit_middleware`] for routes that +/// install `RequestDecompressionLayer`: it must be layered *inside* the +/// decompression layer, so the body it wraps is the **decompressed** stream. +/// +/// Only requests marked [`ContentEncoded`] by the global middleware are +/// accounted; plain (uncompressed) requests are already covered by the global +/// wire-byte accounting, and charging them here as well would double-count. +/// For decompressed requests the `Content-Length` header has been stripped by +/// the decompression layer, so the full decoded size is charged as it streams. +pub(crate) async fn decoded_body_accounting_middleware( + State(limiter): State, + req: Request, + next: Next, +) -> Response { + if req.extensions().get::().is_none() { + return next.run(req).await; + } + + let content_length = req + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .unwrap_or(0); + + let accounting = BodyMemoryAccounting::default(); + // Retain the permits for the whole request, like the global middleware. + let retained_accounting = accounting.clone(); + let (mut parts, body) = req.into_parts(); + parts.extensions.insert(accounting.clone()); + let accounted = AccountedBody::new(body, limiter, content_length, accounting); + let req = Request::from_parts(parts, Body::new(accounted)); + + let response = next.run(req).await; + // `retained_accounting` stays alive until this function returns, holding + // the permits across the handler's use of the collected body. + quota_exceeded_response(&retained_accounting, response) +} + +/// Holds the memory guards acquired while a request body is streamed. Shared +/// between the [`AccountedBody`] wrapper, the request extensions and the +/// middleware itself, so the permits are only released when the request +/// (including any collected body buffers and the handler still using them) +/// is finished. +#[derive(Clone, Default)] +struct BodyMemoryAccounting { + guards: Arc>>>, + /// Set when an incremental charge hit the quota, so the rejection can be + /// rewritten into a 429 after the middleware regains control. + quota_exceeded: Arc, +} + +impl BodyMemoryAccounting { + fn hold(&self, guard: MemoryGuard) { + self.guards.lock().unwrap().push(guard); + } + + fn mark_quota_exceeded(&self) { + self.quota_exceeded.store(true, Ordering::Release); + } + + fn take_quota_exceeded(&self) -> bool { + self.quota_exceeded.swap(false, Ordering::AcqRel) + } +} + +type AcquireFuture = + Pin>> + Send>>; + +enum ChargeOutcome { + Charged, + Pending, + Failed, +} + +/// A request body wrapper that charges the shared limiter for every byte +/// streamed beyond `pre_charged` (the upfront `Content-Length` reservation). +struct AccountedBody { + inner: Body, + limiter: ServerMemoryLimiter, + accounting: BodyMemoryAccounting, + /// Bytes of body data streamed so far. + streamed: u64, + /// Bytes of the streamed body already covered by incremental permits. + incrementally_charged: u64, + /// Bytes covered by the upfront `Content-Length` reservation. + pre_charged: u64, + /// A data frame held back while its permit is being acquired, together + /// with the number of extra bytes to charge for it. + pending: Option<(Frame, u64)>, + pending_acquire: Option, +} + +impl AccountedBody { + fn new( + inner: Body, + limiter: ServerMemoryLimiter, + pre_charged: u64, + accounting: BodyMemoryAccounting, + ) -> Self { + Self { + inner, + limiter, + accounting, + streamed: 0, + incrementally_charged: 0, + pre_charged, + pending: None, + pending_acquire: None, + } + } + + /// Bytes that still need an incremental permit for the body streamed so far. + fn uncharged(&self) -> u64 { + self.streamed + .saturating_sub(self.pre_charged) + .saturating_sub(self.incrementally_charged) + } + + /// Acquires a permit for `bytes`; only one acquisition may be in flight. + /// + /// A parked acquisition always belongs to the currently buffered frame, so + /// it must be resolved *before* any new acquisition is tried: letting the + /// fast path succeed while a waiter is parked would leave the waiter alive, + /// and its late completion would then be credited with another frame's + /// byte count (under-reserving memory relative to what is marked charged). + fn charge(&mut self, bytes: u64, cx: &mut Context<'_>) -> ChargeOutcome { + debug_assert!(bytes > 0); + if let Some(fut) = self.pending_acquire.as_mut() { + return match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.pending_acquire = None; + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => { + self.pending_acquire = None; + ChargeOutcome::Failed + } + Poll::Pending => ChargeOutcome::Pending, + }; + } + // Fast path: the permit is immediately available. + if let Some(guard) = self.limiter.try_acquire(bytes) { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + return ChargeOutcome::Charged; + } + let limiter = self.limiter.clone(); + let mut fut = Box::pin(async move { limiter.acquire(bytes).await }); + // Poll immediately so a ready result is not missed (returning + // `Pending` without registering progress would park the task forever). + match fut.as_mut().poll(cx) { + Poll::Ready(Ok(guard)) => { + self.accounting.hold(guard); + self.incrementally_charged += bytes; + ChargeOutcome::Charged + } + Poll::Ready(Err(_)) => ChargeOutcome::Failed, + Poll::Pending => { + self.pending_acquire = Some(fut); + ChargeOutcome::Pending + } + } + } + + /// Hands out `frame` once `bytes` of it are charged, buffering the frame + /// while the acquisition is pending. + fn hand_out( + &mut self, + frame: Frame, + bytes: u64, + cx: &mut Context<'_>, + ) -> Poll, axum::Error>>> { + if bytes == 0 { + return Poll::Ready(Some(Ok(frame))); + } + match self.charge(bytes, cx) { + ChargeOutcome::Charged => Poll::Ready(Some(Ok(frame))), + ChargeOutcome::Pending => { + self.pending = Some((frame, bytes)); + Poll::Pending + } + ChargeOutcome::Failed => { + self.accounting.mark_quota_exceeded(); + Poll::Ready(Some(Err(limit_exceeded_error()))) + } + } + } +} + +impl HttpBody for AccountedBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + + // Finish a pending acquisition for a buffered frame first. + if let Some((frame, bytes)) = this.pending.take() { + return this.hand_out(frame, bytes, cx); + } + + match Pin::new(&mut this.inner).poll_frame(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(None) => Poll::Ready(None), + Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), + Poll::Ready(Some(Ok(frame))) => { + if let Some(data) = frame.data_ref() { + this.streamed += data.len() as u64; + let uncharged = this.uncharged(); + this.hand_out(frame, uncharged, cx) + } else { + // Trailers pass through uncharged. + Poll::Ready(Some(Ok(frame))) + } + } + } + } + + fn is_end_stream(&self) -> bool { + self.pending.is_none() && self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +fn limit_exceeded_error() -> axum::Error { + axum::Error::new(std::io::Error::other( + "request body exceeded the aggregate memory limit", + )) +} + +#[cfg(test)] +mod tests { + use axum::body::{Body, Bytes}; + use axum::http::{Request, StatusCode, header}; + use axum::routing::post; + use axum::{Router, middleware}; + use common_memory_manager::OnExhaustedPolicy; + use tower::ServiceExt; + + use super::memory_limit_middleware; + use crate::request_memory_limiter::ServerMemoryLimiter; + + /// Streams `body` in `chunk_size` frames without a Content-Length header + /// (like HTTP/1.1 chunked transfer). + fn chunked_request(uri: &str, body: Vec, chunk_size: usize) -> Request { + let stream = futures_util::stream::iter( + body.chunks(chunk_size) + .map(|c| Ok::<_, std::io::Error>(Bytes::copy_from_slice(c))) + .collect::>(), + ); + Request::builder() + .uri(uri) + .method("POST") + .body(Body::from_stream(stream)) + .unwrap() + } + + fn counted_body_app(limiter: ServerMemoryLimiter) -> Router { + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + memory_limit_middleware, + )) + } + + #[tokio::test] + async fn test_chunked_body_larger_than_quota_is_rejected() { + // The quota cannot cover a fully streamed 64 KiB chunked body. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 64 * 1024], 1024)) + .await + .unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "a chunked body larger than the quota must be rejected with 429, not 400" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_chunked_body_fits() { + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let res = app + .oneshot(chunked_request("/echo", vec![b'x'; 1024], 512)) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_content_length_still_admitted_upfront() { + // A Content-Length larger than the quota is rejected before the body + // is read (429), unchanged from the previous behavior. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn test_understated_content_length_is_charged_the_difference() { + // Claim 1 KiB in Content-Length but stream 16 KiB: the extra 15 KiB + // must be charged while streaming and rejected by the small quota. + let limiter = ServerMemoryLimiter::new(2 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let stream = futures_util::stream::iter(vec![Ok::<_, std::io::Error>( + Bytes::copy_from_slice(&vec![b'x'; 16 * 1024]), + )]); + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "1024") + .body(Body::from_stream(stream)) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_honest_content_length_is_not_double_charged() { + // With a quota large enough for the whole body, an accurate + // Content-Length must admit the request without extra charges. + let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail); + let app = counted_body_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header(header::CONTENT_LENGTH, "4096") + .body(Body::from(vec![b'x'; 4096])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(0, limiter.used_bytes()); + } + + /// Mirrors the route-local stack: decoded-body accounting *inside* a + /// request-decompression layer, like the ingestion routes in `http.rs`. + fn decompressing_app(limiter: ServerMemoryLimiter) -> Router { + use tower_http::decompression::RequestDecompressionLayer; + + Router::new() + .route( + "/echo", + post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }), + ) + .layer(middleware::from_fn_with_state( + limiter, + super::decoded_body_accounting_middleware, + )) + .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true)) + } + + fn zstd_body(decoded: &[u8]) -> Body { + let compressed = zstd::stream::encode_all(decoded, 3).unwrap(); + Body::from(compressed) + } + + #[tokio::test] + async fn test_decompressed_body_is_charged_for_content_encoded_requests() { + // The decoded size (64 KiB) exceeds the 16 KiB quota: the accounting + // layer must see the decompressed bytes and reject mid-stream. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 64 * 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the quota must be rejected with 429" + ); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_small_decompressed_body_fits() { + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .extension(super::ContentEncoded) + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap(); + assert_eq!(&body[..], b"1024", "the handler must see the decoded body"); + assert_eq!(0, limiter.used_bytes(), "guards must be released"); + } + + #[tokio::test] + async fn test_pending_acquisition_is_attributed_to_its_own_bytes() { + use std::time::Duration; + + use tokio_stream::wrappers::ReceiverStream; + + // Regression test: a parked acquisition must be resolved before any + // new acquisition is tried; otherwise a late waiter completion can be + // credited with a later frame's byte count. + // + // Scenario: quota 8 KiB fully held externally; frame A (4 KiB) parks + // its acquisition; the external guard is then released so A drains; + // frame B (8 KiB) arrives. Correct behavior: B waits for its own 8 KiB + // acquisition and is rejected when it times out (only 4 KiB are free). + // Buggy behavior: the stale waiter (4 KiB) completes and is credited + // with B's 8 KiB, so B is delivered although the quota cannot cover it. + let limiter = ServerMemoryLimiter::new( + 8 * 1024, + OnExhaustedPolicy::Wait { + timeout: Duration::from_millis(500), + }, + ); + let app = counted_body_app(limiter.clone()); + + let (tx, rx) = tokio::sync::mpsc::channel::>(2); + let driver_limiter = limiter.clone(); + let driver = tokio::spawn(async move { + // Hold the whole quota so frame A's acquisition parks. + let external = driver_limiter.acquire(8 * 1024).await.unwrap(); + tx.send(Ok(Bytes::from(vec![b'a'; 4096]))).await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + // Free the quota: frame A's waiter drains and A is delivered. + drop(external); + tokio::time::sleep(Duration::from_millis(300)).await; + tx.send(Ok(Bytes::from(vec![b'b'; 8192]))).await.unwrap(); + // Keep the channel open; frame B's own acquisition must time out. + tokio::time::sleep(Duration::from_millis(1500)).await; + }); + + let res = app + .oneshot( + Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from_stream(ReceiverStream::new(rx))) + .unwrap(), + ) + .await + .unwrap(); + driver.await.unwrap(); + + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "frame B must wait for its own 8 KiB and time out, not borrow A's reservation" + ); + } + + #[tokio::test] + async fn test_plain_requests_are_not_charged_by_the_decoded_layer() { + // Without the ContentEncoded marker the decoded layer must be a + // no-op: plain bodies are accounted by the global wire accounting. + let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail); + let app = decompressing_app(limiter.clone()); + + // Plain body, no marker: passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .body(Body::from(vec![b'x'; 1024])) + .unwrap(); + let res = app.clone().oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + // Compressed body WITHOUT the marker (e.g. a route that has the + // decompression layer but no accounting): still passes through. + let req = Request::builder() + .uri("/echo") + .method("POST") + .header("content-encoding", "zstd") + .body(zstd_body(&vec![b'x'; 1024])) + .unwrap(); + let res = app.oneshot(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } } diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 68542a6f642..05d9725f947 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -52,8 +52,9 @@ use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_remote_write::{ REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, }; -use crate::prom_store::snappy_decompress; +use crate::prom_store::{MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited}; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; +use crate::request_memory_limiter::ServerMemoryLimiter; pub const PHYSICAL_TABLE_PARAM: &str = "physical_table"; pub const DEFAULT_ENCODING: &str = "snappy"; @@ -75,6 +76,9 @@ pub struct PromStoreState { pub prom_store_with_metric_engine: bool, pub prom_validation_mode: PromValidationMode, pub pending_rows_batcher: Option>, + /// Shared request-memory limiter used to charge decompressed remote + /// read/write bodies against the aggregate quota. + pub memory_limiter: ServerMemoryLimiter, } #[derive(Debug, Serialize, Deserialize)] @@ -146,6 +150,7 @@ async fn remote_write_v1( prom_store_with_metric_engine, prom_validation_mode, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -173,13 +178,22 @@ async fn remote_write_v1( processor.set_pipeline(pipeline_handler, query_ctx.clone(), pipeline_def); } - let mut decoded = - decode_remote_write_request(is_zstd, body, prom_validation_mode, &mut processor)?; + // Keep the decode guards alive until the batches are built: the + // `TablesBuilder` retains the decompressed buffer as its raw data. + let (mut decoded, decode_guards) = decode_remote_write_request( + is_zstd, + body, + prom_validation_mode, + &mut processor, + &memory_limiter, + ) + .await?; // Parsing borrows the decode buffer, but row building copies out of it: tag // values through `decode_string`, column names through `to_owned`, and the // borrowing `col_indexes` dies inside `as_insert_requests`. Nothing below - // references the buffer, so it need not span the write. + // references the buffer, so it need not span the write — the same goes for + // the decode guards charged against the aggregate memory quota. let req = if processor.use_pipeline { drop(decoded); processor.exec_pipeline().await? @@ -188,6 +202,7 @@ async fn remote_write_v1( drop(decoded); req }; + drop(decode_guards); let batches = into_prom_write_batches(req, query_ctx); let outcome = match write_prometheus_rows_with_progress( @@ -227,6 +242,7 @@ async fn remote_write_v2( prom_store_with_metric_engine, prom_validation_mode: _, pending_rows_batcher, + memory_limiter, } = state; if let Some(response) = vm_proto_version_response(¶ms) { @@ -240,7 +256,7 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let req = match decode_remote_write_v2(is_zstd, body) { + let req = match decode_remote_write_v2(is_zstd, body, &memory_limiter).await { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -840,7 +856,7 @@ pub async fn remote_read( let db = params.db.clone().unwrap_or_default(); query_ctx.set_channel(Channel::Prometheus); - let request = decode_remote_read_request(body).await?; + let request = decode_remote_read_request(body, &state.memory_limiter).await?; let query_ctx = Arc::new(query_ctx); let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_READ_ELAPSED @@ -850,8 +866,13 @@ pub async fn remote_read( state.prom_store_handler.read(request, query_ctx).await } -async fn decode_remote_read_request(body: Bytes) -> Result { - let buf = snappy_decompress(&body[..])?; +async fn decode_remote_read_request( + body: Bytes, + limiter: &ServerMemoryLimiter, +) -> Result { + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding is finished. + let buf = snappy_decompress_limited(&body[..], MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await?; ReadRequest::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } @@ -1109,6 +1130,7 @@ mod tests { prom_store_with_metric_engine: false, prom_validation_mode: PromValidationMode::Strict, pending_rows_batcher: None, + memory_limiter: ServerMemoryLimiter::default(), } } diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index c9aabfb2745..06729795b5b 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -27,6 +27,7 @@ pub(crate) mod v2; pub mod validation; use bytes::Bytes; +use common_memory_manager::MemoryGuard; use lazy_static::lazy_static; use object_pool::Pool; use snafu::ResultExt; @@ -35,7 +36,12 @@ use crate::error; use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; use crate::prom_remote_write::row_builder::TablesBuilder; use crate::prom_remote_write::validation::PromValidationMode; -use crate::prom_store::{snappy_decompress, zstd_decompress}; +use crate::prom_store::{ + ChargedBuffer, MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited, + zstd_decompress_limited, +}; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; /// Prometheus remote write protocol versions, also used as the `version` label /// of the remote write metrics. @@ -47,43 +53,74 @@ lazy_static! { Pool::new(256, PromWriteRequest::default); } -pub fn try_decompress(is_zstd: bool, body: &[u8]) -> crate::error::Result> { - if is_zstd { - zstd_decompress(body) - } else { - snappy_decompress(body) +/// Decompresses a remote-write body with the codec selected by `is_zstd`, +/// bounded by the hard decoded-size cap and charged to the aggregate +/// request-memory limiter. +/// +/// Due to vmagent's limitation there is a chance that vmagent sends the +/// content-encoding header wrong, so decoding falls back to the other codec +/// when the first attempt fails on malformed input. Size-limit and quota +/// errors are final: the body either expanded beyond the cap under the +/// declared codec, or it could not be admitted, and the other codec cannot +/// produce a different outcome. +/// +/// see +/// see +pub(crate) async fn decompress_remote_write_body( + is_zstd: bool, + body: &[u8], + limiter: &ServerMemoryLimiter, +) -> crate::error::Result { + let decompress = |is_zstd: bool| async move { + if is_zstd { + zstd_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } else { + snappy_decompress_limited(body, MAX_DECOMPRESSED_REQUEST_SIZE, limiter).await + } + }; + + match decompress(is_zstd).await { + Ok(buf) => Ok(buf), + Err(e) => { + if matches!( + e, + crate::error::Error::DecompressSnappyPromRemoteRequest { .. } + | crate::error::Error::DecompressZstdPromRemoteRequest { .. } + ) { + decompress(!is_zstd).await + } else { + Err(e) + } + } } } -pub fn decode_remote_write_request( +pub async fn decode_remote_write_request( is_zstd: bool, body: Bytes, prom_validation_mode: PromValidationMode, processor: &mut PromSeriesProcessor, -) -> crate::error::Result> { + limiter: &ServerMemoryLimiter, +) -> crate::error::Result<( + TablesBuilder<'static>, + Vec>, +)> { let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION]) .start_timer(); - // due to vmagent's limitation, there is a chance that vmagent is - // sending content type wrong so we have to apply a fallback with decoding - // the content in another method. - // - // see https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301 - // see https://github.com/GreptimeTeam/greptimedb/issues/3929 - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - // fallback to the other compression method - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes; the caller must + // keep them alive as long as the returned `TablesBuilder`, which retains + // the decompressed buffer as its raw data. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); - + let (data, guards) = buf.into_parts(); let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default); request - .decode(buf, prom_validation_mode, processor) + .decode(data, prom_validation_mode, processor) .context(error::DecodePromRemoteRequestSnafu)?; - Ok(std::mem::take(&mut request.table_data)) + let tables = std::mem::take(&mut request.table_data); + Ok((tables, guards)) } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 77ebe5a7299..b9e34dc7f24 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -50,12 +50,13 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; -use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; +use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, decompress_remote_write_body}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; +use crate::request_memory_limiter::ServerMemoryLimiter; use crate::row_writer::{self, TableData}; use crate::semantic::{ METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, @@ -127,21 +128,18 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub semantic_index: SemanticIndexes, } -pub(crate) fn decode_remote_write_v2( +pub(crate) async fn decode_remote_write_v2( is_zstd: bool, body: Bytes, + limiter: &ServerMemoryLimiter, ) -> Result { let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) .start_timer(); - // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a - // mismatched content-encoding header, so try the other compression on failure. - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + // Holds the memory permits for the decompressed bytes until the protobuf + // decoding and conversion below are finished. + let buf = decompress_remote_write_body(is_zstd, &body[..], limiter).await?; // Decompression copied the payload out, so the compressed body is no longer needed. drop(body); let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; @@ -789,8 +787,9 @@ pub mod test_util { use snafu::ResultExt; use crate::error::{self, Result}; - use crate::prom_remote_write::try_decompress; + use crate::prom_remote_write::decompress_remote_write_body; use crate::prom_store::snappy_compress; + use crate::request_memory_limiter::ServerMemoryLimiter; pub fn request_with_labels_and_samples( labels: Vec<(&str, &str)>, @@ -807,14 +806,24 @@ pub mod test_util { } pub fn decode_request(is_zstd: bool, body: Bytes) -> Result { - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; + let buf = block_on_decode(is_zstd, body)?; Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } + /// Runs the charged decompression on a throwaway runtime so plain `#[test]` + /// callers can stay synchronous. + fn block_on_decode(is_zstd: bool, body: Bytes) -> Result> { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async move { + decompress_remote_write_body(is_zstd, &body[..], &ServerMemoryLimiter::default()) + .await + .map(|buf| buf.to_vec()) + }) + } + pub fn write_requests( request: Request, ) -> Result<(Vec, Vec, u64, u64)> { @@ -826,7 +835,15 @@ pub mod test_util { is_zstd: bool, body: Bytes, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::decode_remote_write_v2(is_zstd, body)?; + let requests = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(super::decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + ))?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -939,7 +956,10 @@ mod tests { assert_eq!(decoded.timeseries[0].samples.len(), 1); assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); - assert_eq!(decode_remote_write_v2(true, body).unwrap().sample_count, 1); + assert_eq!( + decode_v2_on_test_runtime(true, body).unwrap().sample_count, + 1 + ); } #[test] @@ -2021,7 +2041,21 @@ mod tests { fn decode_wire(wire: &[u8]) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); - decode_remote_write_v2(false, body) + decode_v2_on_test_runtime(false, body) + } + + /// Runs the async (charged) v2 decoder on a throwaway runtime so plain + /// `#[test]` callers can stay synchronous. + fn decode_v2_on_test_runtime(is_zstd: bool, body: Bytes) -> Result { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(decode_remote_write_v2( + is_zstd, + body, + &ServerMemoryLimiter::default(), + )) } fn decode_wire_error(wire: &[u8], name: &str) -> error::Error { @@ -2096,7 +2130,7 @@ mod tests { fn decode_test_request(request: Request) -> Result { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - decode_remote_write_v2(false, body) + decode_v2_on_test_runtime(false, body) } fn assert_invalid(name: &str, request: Request, expected: &str) { diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index fcf177c1e5f..ff82f8f933b 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -32,6 +32,7 @@ use arrow::datatypes::{ TimestampNanosecondType, UInt32Type, }; use common_grpc::precision::Precision; +use common_memory_manager::MemoryGuard; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{tracing, warn}; @@ -44,6 +45,8 @@ use snap::raw::{Decoder, Encoder}; use crate::error::{self, Result}; use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION; +use crate::request_memory_limiter::ServerMemoryLimiter; +use crate::request_memory_metrics::RequestMemoryMetrics; use crate::row_writer::{self, MultiTableData}; pub const METRIC_NAME_LABEL: &str = "__name__"; @@ -603,8 +606,147 @@ pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertR Ok(multi_table_data.into_row_insert_requests()) } +/// Hard cap on the decompressed size of a handler-decompressed request body +/// (Prometheus remote read/write, Loki push). +/// +/// The HTTP body limit only bounds the *encoded* bytes on the wire. Without +/// this cap a tiny compressed payload (e.g. a five-byte raw-Snappy varint or a +/// few KiB of zstd RLE blocks) could declare/expand to multiple GiB of heap +/// before any protobuf validation runs, so the decompressed size must be +/// bounded on its own. +pub const MAX_DECOMPRESSED_REQUEST_SIZE: usize = 512 * 1024 * 1024; + +/// Size of the chunks a zstd body is decompressed and charged in. +const ZSTD_DECOMPRESS_CHUNK: usize = 256 * 1024; + +/// Decompressed bytes together with the memory permits charged for them. +/// The permits are released when the buffer is dropped, so the aggregate +/// memory accounting covers the whole lifetime of the decompressed data. +pub(crate) struct ChargedBuffer { + pub data: Vec, + guards: Vec>, +} + +impl std::ops::Deref for ChargedBuffer { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + &self.data + } +} + +impl std::fmt::Debug for ChargedBuffer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChargedBuffer") + .field("data_len", &self.data.len()) + .field("guards", &self.guards.len()) + .finish() + } +} + +impl ChargedBuffer { + /// Creates a buffer with the permits charged for it. + pub(crate) fn new(data: Vec, guards: Vec>) -> Self { + Self { data, guards } + } + + /// Splits the buffer into the raw bytes and the memory permits charged for + /// them. Hold the permits as long as the bytes are kept. + pub(crate) fn into_parts(self) -> (Vec, Vec>) { + (self.data, self.guards) + } +} + +/// Decompresses a raw Snappy body after enforcing a hard cap on the decoded +/// size and charging the aggregate request-memory limiter. +/// +/// The raw Snappy format prefixes the stream with a varint of the decoded +/// length, so the output size is known (and validated) before the output +/// buffer is allocated by `decompress_vec`. +pub(crate) async fn snappy_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: max_decompressed as u64, + } + ); + // Reserve the declared size before the allocation happens so concurrent + // decompressions cannot exceed the aggregate quota. + let guard = limiter.acquire(decoded_len as u64).await?; + let mut decoder = Decoder::new(); + let data = decoder + .decompress_vec(buf) + .context(error::DecompressSnappyPromRemoteRequestSnafu)?; + Ok(ChargedBuffer { + data, + guards: vec![guard], + }) +} + +/// Decompresses a zstd body with a hard cap on the decoded size, charging the +/// aggregate request-memory limiter incrementally while the output is +/// materialized (zstd streams don't declare their decoded size up front). +pub(crate) async fn zstd_decompress_limited( + buf: &[u8], + max_decompressed: usize, + limiter: &ServerMemoryLimiter, +) -> Result { + use std::io::Read; + + let decoder = zstd::stream::read::Decoder::new(buf) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + // Allow reading one byte past the limit so oversized streams are detected + // instead of being silently truncated. + let mut limited = decoder.take(max_decompressed as u64 + 1); + let mut data = Vec::new(); + let mut guards = Vec::new(); + let mut chunk = vec![0u8; ZSTD_DECOMPRESS_CHUNK]; + loop { + let n = limited + .read(&mut chunk) + .context(error::DecompressZstdPromRemoteRequestSnafu)?; + if n == 0 { + break; + } + ensure!( + data.len() + n <= max_decompressed, + error::DecompressedBodyTooLargeSnafu { + size: (data.len() + n) as u64, + limit: max_decompressed as u64, + } + ); + // Charge each chunk before it is appended so the aggregate quota + // bounds the memory while it is being materialized. + guards.push(limiter.acquire(n as u64).await?); + data.extend_from_slice(&chunk[..n]); + } + Ok(ChargedBuffer { data, guards }) +} + +/// Decompresses a raw Snappy body with the hard decoded-size cap enforced, +/// without charging the aggregate request-memory limiter. +/// +/// Used for response bodies (remote read replies) and tests; request bodies +/// must use [`snappy_decompress_limited`] instead so their decoded size is +/// also charged to the aggregate quota. #[inline] pub fn snappy_decompress(buf: &[u8]) -> Result> { + let decoded_len = + snap::raw::decompress_len(buf).context(error::DecompressSnappyPromRemoteRequestSnafu)?; + ensure!( + decoded_len <= MAX_DECOMPRESSED_REQUEST_SIZE, + error::DecompressedBodyTooLargeSnafu { + size: decoded_len as u64, + limit: MAX_DECOMPRESSED_REQUEST_SIZE as u64, + } + ); let mut decoder = Decoder::new(); decoder .decompress_vec(buf) @@ -619,11 +761,6 @@ pub fn snappy_compress(buf: &[u8]) -> Result> { .context(error::CompressPromRemoteRequestSnafu) } -#[inline] -pub fn zstd_decompress(buf: &[u8]) -> Result> { - zstd::stream::decode_all(buf).context(error::DecompressZstdPromRemoteRequestSnafu) -} - /// Mock timeseries for test, it is both used in servers and frontend crate /// So we present it here pub fn mock_timeseries() -> Vec { @@ -763,6 +900,7 @@ mod tests { DictionaryArray, Float64Array, StringArray, TimestampMillisecondArray, UInt32Array, }; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, UInt32Type}; + use common_memory_manager::OnExhaustedPolicy; use common_recordbatch::DfRecordBatch; use datafusion::prelude::SessionContext; use datatypes::data_type::ConcreteDataType; @@ -1521,4 +1659,126 @@ mod tests { timeseries[1].samples ); } + + #[tokio::test] + async fn test_snappy_varint_bomb_is_rejected_before_allocation() { + // A five-byte raw Snappy header declaring a ~4 GiB decoded length. + // Regression test for the compression-bomb finding: the declared size + // must be rejected by the cap check before any buffer is allocated. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let limiter = ServerMemoryLimiter::default(); + + let err = snappy_decompress_limited(&bomb, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_decompress_limited_roundtrip() { + let raw = vec![b'a'; 1024 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::default(); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(&decoded[..], &raw[..]); + } + + #[tokio::test] + async fn test_snappy_charges_and_releases_quota() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + let limiter = ServerMemoryLimiter::new(128 * 1024, OnExhaustedPolicy::Fail); + + let decoded = + snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + // Rounding to the 1 KiB permit granularity. + assert!(limiter.used_bytes() >= 64 * 1024); + assert!(limiter.used_bytes() <= 65 * 1024); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[tokio::test] + async fn test_snappy_rejects_when_quota_exhausted() { + let raw = vec![b'a'; 64 * 1024]; + let mut encoder = snap::raw::Encoder::new(); + let compressed = encoder.compress_vec(&raw).unwrap(); + // Quota smaller than the decoded size. + let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail); + + let err = snappy_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::MemoryLimitExceeded { .. } + )); + assert_eq!(0, limiter.used_bytes(), "failed charge must be released"); + } + + #[tokio::test] + async fn test_zstd_at_and_over_the_cap() { + let limiter = ServerMemoryLimiter::default(); + let small_cap = 8 * 1024; + + // Exactly at the cap: accepted. + let raw = vec![0u8; small_cap]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let decoded = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), small_cap); + + // One byte past the cap: rejected. + let raw = vec![0u8; small_cap + 1]; + let compressed = zstd::stream::encode_all(&raw[..], 3).unwrap(); + let err = zstd_decompress_limited(&compressed, small_cap, &limiter) + .await + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } + + #[tokio::test] + async fn test_zstd_charges_incrementally() { + // 64 KiB wire that expands to 2 MiB must be charged chunk by chunk + // against a quota that only fits the expanded size. + let raw = vec![0u8; 2 * 1024 * 1024]; + let compressed = zstd::stream::encode_all(&raw[..], 19).unwrap(); + assert!(compressed.len() < 16 * 1024); + let limiter = ServerMemoryLimiter::new(4 * 1024 * 1024, OnExhaustedPolicy::Fail); + + let decoded = zstd_decompress_limited(&compressed, MAX_DECOMPRESSED_REQUEST_SIZE, &limiter) + .await + .unwrap(); + assert_eq!(decoded.len(), raw.len()); + assert!(limiter.used_bytes() >= raw.len() as u64); + drop(decoded); + assert_eq!(0, limiter.used_bytes()); + } + + #[test] + fn test_snappy_response_decompress_is_capped() { + // The response-side (uncharged) helper must enforce the cap too. + let bomb = [0xffu8, 0xff, 0xff, 0xff, 0x0f]; + let err = snappy_decompress(&bomb).unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DecompressedBodyTooLarge { .. } + )); + } } diff --git a/src/servers/src/request_memory_limiter.rs b/src/servers/src/request_memory_limiter.rs index 8c3d106d9c0..ac072f51b2e 100644 --- a/src/servers/src/request_memory_limiter.rs +++ b/src/servers/src/request_memory_limiter.rs @@ -73,4 +73,9 @@ impl ServerMemoryLimiter { pub fn limit_bytes(&self) -> u64 { self.manager.limit_bytes() } + + /// Returns currently reserved (in-flight) bytes. + pub fn used_bytes(&self) -> u64 { + self.manager.used_bytes() + } } diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 62c32091d76..64f9f4c6af1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -747,7 +747,13 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( ..Default::default() }); - let mut http_server = HttpServerBuilder::new(http_opts) + // The memory limiter must be wired before the handlers: routes capture + // the builder's limiter at wiring time (mirrors `frontend::server`). + let mut http_server = HttpServerBuilder::new(http_opts); + if let Some(limiter) = memory_limiter { + http_server = http_server.with_memory_limiter(limiter); + } + let mut http_server = http_server .with_sql_handler(instance.fe_instance().clone()) .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) @@ -762,10 +768,6 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( http_server = http_server.with_user_provider(user_provider); } - if let Some(limiter) = memory_limiter { - http_server = http_server.with_memory_limiter(limiter); - } - let http_server = http_server.build(); let app = http_server.build(http_server.make_app()).unwrap(); @@ -784,7 +786,7 @@ pub async fn setup_test_prom_app_with_frontend( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, false).await + setup_test_prom_app_with_frontend_inner(store_type, name, false, None).await } /// Like [`setup_test_prom_app_with_frontend`] but enables the pending-rows batcher, @@ -794,13 +796,25 @@ pub async fn setup_test_prom_app_with_frontend_batched( store_type: StorageType, name: &str, ) -> (Router, TestGuard) { - setup_test_prom_app_with_frontend_inner(store_type, name, true).await + setup_test_prom_app_with_frontend_inner(store_type, name, true, None).await +} + +/// Like [`setup_test_prom_app_with_frontend`] but wires a shared request-memory +/// limiter into the server, mirroring production deployments that configure +/// `max_in_flight_write_bytes`. +pub async fn setup_test_prom_app_with_frontend_and_memory_limiter( + store_type: StorageType, + name: &str, + memory_limiter: Option, +) -> (Router, TestGuard) { + setup_test_prom_app_with_frontend_inner(store_type, name, false, memory_limiter).await } async fn setup_test_prom_app_with_frontend_inner( store_type: StorageType, name: &str, enable_batcher: bool, + memory_limiter: Option, ) -> (Router, TestGuard) { unsafe { std::env::set_var("TZ", "UTC"); @@ -849,17 +863,24 @@ async fn setup_test_prom_app_with_frontend_inner( let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)"; run_sql(sql, &instance).await; - let http_server = build_test_prom_server(instance.fe_instance().clone(), enable_batcher) - .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) - .build(); + let http_server = build_test_prom_server( + instance.fe_instance().clone(), + enable_batcher, + memory_limiter, + ) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); let app = http_server.build(http_server.make_app()).unwrap(); (app, instance.guard) } /// Builds Prometheus HTTP routes for either a standalone or distributed frontend. +/// `memory_limiter` must be set before the handlers are wired (handlers capture +/// the limiter at wiring time, mirroring production in `frontend::server`). pub fn build_test_prom_server( frontend_ref: Arc, enable_batcher: bool, + memory_limiter: Option, ) -> HttpServerBuilder { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), @@ -886,9 +907,13 @@ pub fn build_test_prom_server( None }; assert_eq!(pending_rows_batcher.is_some(), enable_batcher); - HttpServerBuilder::new(http_opts) + let mut builder = HttpServerBuilder::new(http_opts) .with_sql_handler(frontend_ref.clone()) - .with_logs_handler(frontend_ref.clone()) + .with_logs_handler(frontend_ref.clone()); + if let Some(limiter) = memory_limiter { + builder = builder.with_memory_limiter(limiter); + } + builder .with_prom_handler( frontend_ref.clone(), Some(frontend_ref.clone()), @@ -1254,7 +1279,7 @@ pub async fn setup_pg_server_with_prom_native_histogram( let instance = setup_standalone_instance(name, store_type).await; // Prometheus remote-write HTTP app with native histograms enabled. - let http_server = build_test_prom_server(instance.fe_instance().clone(), false) + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, None) .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) .build(); let app = http_server.build(http_server.make_app()).unwrap(); diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 4b78e5fe276..7cc02cfc456 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -3342,7 +3342,7 @@ async fn test_histogram_ingestion_storage_lifecycle(#[case] metric_engine: bool) let mut context = TestContext::new(builder).await; let make_router = |frontend: Arc| { let builder = if metric_engine { - build_test_prom_server(frontend.clone(), true) + build_test_prom_server(frontend.clone(), true, None) } else { HttpServerBuilder::new(HttpOptions::default()) .with_sql_handler(frontend.clone()) diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index c73f80da78e..e9cfee7e9d4 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -136,6 +136,7 @@ macro_rules! grpc_tests { test_grpc_timezone, test_grpc_tls_config, test_grpc_memory_limit, + test_grpc_compressed_memory_reservation, ); )* }; @@ -358,6 +359,104 @@ pub async fn test_grpc_message_size_limit_send(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +/// Transport-compressed gRPC requests reserve the worst-case decoded message +/// size (`max_recv_message_size`) against the aggregate quota *before* tonic +/// decompresses them. With a quota smaller than that bound, compressed +/// requests must fail fast with `RESOURCE_EXHAUSTED` (before decoding), while +/// uncompressed requests are still admitted through the exact post-decode +/// charge. With the default unlimited limiter, compression keeps working. +pub async fn test_grpc_compressed_memory_reservation(store_type: StorageType) { + use api::v1::query_request::Query; + use api::v1::{GreptimeRequest, QueryRequest}; + use tonic::codec::CompressionEncoding; + + let config = GrpcServerConfig { + max_recv_message_size: 1024 * 1024, + ..Default::default() + }; + // Quota smaller than the decoding bound: compressed requests cannot be + // pre-admitted. + let memory_limiter = ServerMemoryLimiter::new(512 * 1024, OnExhaustedPolicy::Fail); + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation", + None, + Some(config), + Some(memory_limiter), + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + + let insert_request = || { + tonic::Request::new(GreptimeRequest { + header: Some(RequestHeader { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + ..Default::default() + }), + request: Some(RequestBody::Query(QueryRequest { + query: Some(Query::Sql( + "INSERT INTO grpc_quota VALUES (1000, 1.0)".to_string(), + )), + })), + }) + }; + + // Compressed request: the pre-decode reservation (1 MiB) exceeds the + // 512 KiB quota and must be rejected before any decoding happens. + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + let status = compressed.handle(insert_request()).await.unwrap_err(); + assert_eq!( + status.code(), + tonic::Code::ResourceExhausted, + "compressed request must be pre-admitted against the quota, got: {status}" + ); + + // Uncompressed request on the same server: still admitted (its exact + // serialized size fits the quota). + let mut plain = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + plain.handle(insert_request()).await.unwrap(); + + let _ = fe_grpc_server.shutdown().await; + + // Default (unlimited) limiter: compressed requests keep working. + let (_db, fe_grpc_server) = setup_grpc_server_with( + store_type, + "test_grpc_compressed_memory_reservation_unlimited", + None, + Some(GrpcServerConfig::default()), + None, + ) + .await; + let addr = fe_grpc_server.bind_addr().unwrap().to_string(); + let db = Database::new_with_dbname( + format!("{}-{}", DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME), + Client::with_urls(vec![addr.clone()]), + ); + db.sql("CREATE TABLE grpc_quota(ts TIMESTAMP TIME INDEX, v DOUBLE)") + .await + .unwrap(); + let mut compressed = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap() + .send_compressed(CompressionEncoding::Zstd); + compressed.handle(insert_request()).await.unwrap(); + let _ = fe_grpc_server.shutdown().await; +} + pub async fn test_grpc_message_size_limit_recv(store_type: StorageType) { let config = GrpcServerConfig { max_recv_message_size: 10, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 780140e3500..a05c58444e3 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -89,8 +89,10 @@ use table::table_name::TableName; use tests_integration::test_util::{ MockInstanceImpl, StorageType, TestGuard, assert_wal_delta, build_test_prom_server, setup_test_http_app, setup_test_http_app_with_frontend, + setup_test_http_app_with_frontend_and_custom_options, setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, + setup_test_prom_app_with_frontend_and_memory_limiter, setup_test_prom_app_with_frontend_batched, }; use urlencoding::encode; @@ -166,6 +168,9 @@ macro_rules! http_tests { test_promql_over_non_millisecond_physical_tables, test_prometheus_label_replace_response, test_prom_http_api, + test_prom_remote_compression_bomb_rejected, + test_prom_remote_write_decoded_body_charged, + test_influxdb_write_decoded_body_charged, test_config_api, test_dashboard_api, test_prometheus_remote_write, @@ -1725,6 +1730,185 @@ pub async fn test_prom_http_api(store_type: StorageType) { guard.remove_all().await; } +/// A tiny compressed body must not be able to demand an unbounded decoded +/// allocation on the Prometheus remote-storage endpoints. The five-byte raw +/// Snappy varint (`ff ff ff ff 0f`) declares a ~4 GiB decoded length and must +/// be rejected by the decoded-size cap. +pub async fn test_prom_remote_compression_bomb_rejected(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = setup_test_prom_app_with_frontend(store_type, "prom_remote_bomb").await; + let client = TestClient::new(app).await; + + // Remote read with the five-byte varint bomb. + let res = client + .post("/v1/prometheus/read") + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!(res.status(), StatusCode::BAD_REQUEST); + + // Remote write v1 with the same bomb, via both codecs (exercising the + // fallback: the zstd attempt fails on malformed input and the snappy + // fallback must reject the declared size). + for encoding in ["snappy", "zstd"] { + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", encoding) + .body(Vec::from([0xffu8, 0xff, 0xff, 0xff, 0x0f])) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::BAD_REQUEST, + "bomb with Content-Encoding {encoding} must be rejected" + ); + } + + // The server must still serve valid requests afterwards. + let write_request = WriteRequest { + timeseries: prom_store::mock_timeseries(), + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + guard.remove_all().await; +} + +/// The decoded (not encoded) size of a Prometheus remote write body must be +/// charged against the aggregate request-memory quota. +pub async fn test_prom_remote_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let (app, mut guard) = setup_test_prom_app_with_frontend_and_memory_limiter( + store_type, + "prom_remote_write_quota", + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A valid write request whose decoded protobuf clearly exceeds the quota + // while its Snappy wire body stays far below it. + let write_request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: "__name__".to_string(), + value: "quota_metric".to_string(), + }, + Label { + name: "payload".to_string(), + // ~8 KiB decoded; compresses to a few hundred bytes. + value: "p".repeat(8 * 1024), + }, + ], + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + }], + ..Default::default() + }], + ..Default::default() + }; + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(compressed) + .send() + .await; + let status = res.status(); + let body = res.text().await; + assert_eq!( + status, + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected, got body: {body}" + ); + + guard.remove_all().await; +} + +/// Decompressed ingestion bodies (routes using `RequestDecompressionLayer`) +/// must be charged against the aggregate request-memory quota as they are +/// decoded, not only their compressed wire size. +pub async fn test_influxdb_write_decoded_body_charged(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + // 4 KiB aggregate quota with fail-fast policy. + let memory_limiter = ServerMemoryLimiter::new(4096, OnExhaustedPolicy::Fail); + let http_opts = servers::http::HttpOptions { + addr: format!("127.0.0.1:{}", common_test_util::ports::get_port()), + ..Default::default() + }; + let (app, mut guard) = setup_test_http_app_with_frontend_and_custom_options( + store_type, + "test_influxdb_decoded_quota", + None, + Some(http_opts), + Some(memory_limiter), + ) + .await; + let client = TestClient::new(app).await; + + // A gzip member expanding to ~8 KiB of line protocol from a tiny wire body. + let lines: Vec = (0..200) + .map(|i| format!("decoded_quota,host=h{} v=1.0,ts={}", i % 20, 1700000000 + i)) + .collect(); + let decoded = lines.join("\n"); + assert!(decoded.len() > 4096); + let compressed = zstd::stream::encode_all(decoded.as_bytes(), 3).unwrap(); + assert!( + compressed.len() < 2048, + "wire body must stay under the quota" + ); + + // Compressed write whose decoded body exceeds the quota: rejected while + // streaming (the extractor surfaces the aborted body as 400). + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::TOO_MANY_REQUESTS, + "decoded body larger than the aggregate quota must be rejected with 429" + ); + + // The server still serves a small compressed write within the quota. + let small = "decoded_quota,host=ok v=1.0,ts=1700000000"; + let compressed_small = zstd::stream::encode_all(small.as_bytes(), 3).unwrap(); + let res = client + .post("/v1/influxdb/write?db=public") + .header("Content-Type", "text/plain") + .header("Content-Encoding", "zstd") + .body(compressed_small) + .send() + .await; + assert_eq!( + res.status(), + StatusCode::NO_CONTENT, + "small compressed write must succeed, got: {}", + res.text().await + ); + + guard.remove_all().await; +} + /// PromQL must behave identically regardless of the physical metric table's /// time index unit: the same samples are remote-written into physical tables /// pre-created with second/micro/nano time indexes (plus a millisecond @@ -3722,7 +3906,7 @@ async fn check_prometheus_remote_write_batched_skip_wal(distributed: bool, v2: b common_telemetry::init_default_ut_logging(); let mut instance = MockInstanceImpl::new(&format!("prom_bulk_skip_wal_v2_{v2}"), distributed).await; - let server = build_test_prom_server(instance.frontend(), true).build(); + let server = build_test_prom_server(instance.frontend(), true, None).build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; write_prometheus_skip_wal_sample(&client, v2, 1000, None).await; @@ -3902,7 +4086,7 @@ async fn setup_redirecting_batched_prom_app( .with_plugin(plugins) .build() .await; - let server = build_test_prom_server(standalone.fe_instance().clone(), true) + let server = build_test_prom_server(standalone.fe_instance().clone(), true, None) .with_greptime_config_options(standalone.opts.datanode_options().to_toml().unwrap()) .build(); let client = TestClient::new(server.build(server.make_app()).unwrap()).await; From a99cb56d7168c88727b34468db4eee1c0f3b4e2c Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 29 Sep 2026 13:17:42 +0000 Subject: [PATCH 078/133] feat(datanode): wire the object store WAL into the datanode and standalone (#9386) * feat(datanode): wire the object store WAL into the datanode and standalone A datanode configured with `provider = "experimental_object_store"` now builds an `ObjectStoreLogStore` instead of failing with the placeholder "not supported yet" error. The builder resolves `storage_provider` against the object store manager (empty selects the default store, an unknown name is rejected as an invalid `storage_provider`), opens the store under the node id and the standalone generation, stops it if a later build step fails, and hands it to `Datanode`, whose shutdown stops it after the region server. Datanode shutdown is now best-effort: every step runs and the first error is returned. A datanode with a metasrv client still rejects the provider. The standalone bootstrap builds `WalProvider::ObjectStore` with the same derived node prefix before the metasrv WAL conversion, so the WAL options it allocates match the prefix the store runs under. `log-store` exports `ObjectStoreLogStore` and its testing hooks. Mito's test utilities can build an engine on the object store WAL over an in-memory object store, and new engine tests cover create, reopen, replay isolation between regions, batch open, drop and offline cleanup, the latest entry id after replay, and the rejection of object store WAL options on a Raft Engine log store. Signed-off-by: jeremyhi * fix(datanode): stop every region engine before the object store WAL The region server returned at the first engine whose stop failed, so the engines after it in map order, Mito included, were never stopped before the datanode stopped the object store WAL. It now attempts every engine, returns the first error and logs the later ones. `ObjectStoreLogStore` is also exported from the `log-store` crate root, and its callers import it from there. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- config/config.md | 4 +- config/datanode.example.toml | 2 +- config/standalone.example.toml | 2 +- src/cmd/src/standalone.rs | 86 +++- src/datanode/src/datanode.rs | 354 +++++++++++++++- src/datanode/src/error.rs | 9 +- src/datanode/src/region_server.rs | 13 +- src/datanode/src/tests.rs | 10 + src/log-store/src/lib.rs | 2 + src/log-store/src/object_store_wal.rs | 9 +- src/log-store/src/object_store_wal/format.rs | 2 + src/log-store/src/object_store_wal/store.rs | 16 +- src/mito2/Cargo.toml | 2 +- src/mito2/src/engine.rs | 2 + src/mito2/src/engine/object_store_wal_test.rs | 379 ++++++++++++++++++ src/mito2/src/test_util.rs | 47 ++- 16 files changed, 875 insertions(+), 64 deletions(-) create mode 100644 src/mito2/src/engine/object_store_wal_test.rs diff --git a/config/config.md b/config/config.md index 116cf8a4f82..06d3b89e075 100644 --- a/config/config.md +++ b/config/config.md @@ -100,7 +100,7 @@ | `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. | | `wal` | -- | -- | The WAL options. | -| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** | +| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** | | `wal.dir` | String | Unset | The directory to store the WAL files.
**It's only used when the provider is `raft_engine`**. | | `wal.file_size` | String | `128MB` | The size of the WAL segment file.
**It's only used when the provider is `raft_engine`**. | | `wal.purge_threshold` | String | `1GB` | The threshold of the WAL size to trigger a purge.
**It's only used when the provider is `raft_engine`**. | @@ -579,7 +579,7 @@ | `meta_client.metadata_cache_ttl` | String | `10m` | TTL of the metadata cache. | | `meta_client.metadata_cache_tti` | String | `5m` | -- | | `wal` | -- | -- | The WAL options. | -| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.**
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** | +| `wal.provider` | String | `raft_engine` | The provider of the WAL.
- `raft_engine`: the wal is stored in the local file system by raft-engine.
- `kafka`: it's remote wal that data is stored in Kafka.
- `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.**
- `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** | | `wal.dir` | String | Unset | The directory to store the WAL files.
**It's only used when the provider is `raft_engine`**. | | `wal.file_size` | String | `128MB` | The size of the WAL segment file.
**It's only used when the provider is `raft_engine`**. | | `wal.purge_threshold` | String | `1GB` | The threshold of the WAL size to trigger a purge.
**It's only used when the provider is `raft_engine`**. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 82b93fd6ccc..3b8f226abb2 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -136,7 +136,7 @@ metadata_cache_tti = "5m" ## - `raft_engine`: the wal is stored in the local file system by raft-engine. ## - `kafka`: it's remote wal that data is stored in Kafka. ## - `noop`: it's a no-op WAL provider that does not store any WAL data.
**Notes: any unflushed data will be lost when the datanode is shutdown.** -## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** +## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** provider = "raft_engine" ## The directory to store the WAL files. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index e3e28563582..c09c815aed0 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -299,7 +299,7 @@ prom_validation_mode = "strict" ## The provider of the WAL. ## - `raft_engine`: the wal is stored in the local file system by raft-engine. ## - `kafka`: it's remote wal that data is stored in Kafka. -## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental and not supported yet.** +## - `experimental_object_store`: the wal is stored as objects in an object store.
**Notes: experimental.** provider = "raft_engine" ## The directory to store the WAL files. diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 60556428117..bc6a62cd4d5 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -44,7 +44,7 @@ use common_meta::procedure_executor::{LocalProcedureExecutor, ProcedureExecutorR use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; use common_meta::sequence::{Sequence, SequenceBuilder}; -use common_meta::wal_provider::{WalProviderRef, build_wal_provider}; +use common_meta::wal_provider::{WalProvider, WalProviderRef, build_wal_provider}; use common_options::plugin_options::StandaloneFlag; use common_procedure::ProcedureManagerRef; use common_query::prelude::set_default_prefix; @@ -52,6 +52,8 @@ use common_telemetry::info; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, TracingOptions}; use common_time::timezone::set_default_timezone; use common_version::{short_version, verbose_version}; +use common_wal::config::DatanodeWalConfig; +use common_wal::config::object_store::STANDALONE_GENERATION; use datanode::config::{DatanodeOptions, StorageConfig}; use datanode::datanode::{Datanode, DatanodeBuilder}; use datanode::region_server::RegionServer; @@ -83,6 +85,34 @@ use crate::{App, create_resource_limit_metrics, error, log_versions, maybe_activ pub const APP_NAME: &str = "greptime-standalone"; +/// Builds the WAL provider that allocates region WAL options in standalone mode. +/// +/// The object store WAL allocates the node prefix of `node_id`, the datanode id +/// of the standalone instance; an absent id is treated as 0, the id +/// `StandaloneOptions::datanode_options` assigns. +pub async fn build_standalone_wal_provider( + wal: &DatanodeWalConfig, + node_id: Option, + kv_backend: KvBackendRef, +) -> Result { + match wal { + DatanodeWalConfig::ObjectStore(config) => Ok(WalProvider::ObjectStore { + prefix: config.node_prefix(node_id.unwrap_or(0), STANDALONE_GENERATION), + }), + DatanodeWalConfig::RaftEngine(_) + | DatanodeWalConfig::Kafka(_) + | DatanodeWalConfig::Noop => { + let metasrv_wal_config = wal + .clone() + .try_into() + .context(error::InvalidWalProviderSnafu)?; + build_wal_provider(&metasrv_wal_config, kv_backend) + .await + .context(error::BuildWalProviderSnafu) + } + } +} + fn standalone_local_file_access( storage: &StorageConfig, ) -> common_datasource::error::Result { @@ -610,15 +640,8 @@ impl StartCommand { .step(10) .build(), ); - let kafka_options = opts - .wal - .clone() - .try_into() - .context(error::InvalidWalProviderSnafu)?; - let wal_provider = build_wal_provider(&kafka_options, kv_backend.clone()) - .await - .context(error::BuildWalProviderSnafu)?; - let wal_provider = Arc::new(wal_provider); + let wal_provider = + Arc::new(build_standalone_wal_provider(&opts.wal, node_id, kv_backend.clone()).await?); let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( table_id_allocator.clone(), wal_provider.clone(), @@ -1035,9 +1058,12 @@ mod tests { use clap::{CommandFactory, Parser}; use common_base::readable_size::ReadableSize; use common_config::ENV_VAR_SEP; + use common_meta::ddl::allocator::wal_options::WalOptionsAllocator; + use common_meta::kv_backend::memory::MemoryKvBackend; use common_options::plugin_options::StandaloneFlag; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; - use common_wal::config::DatanodeWalConfig; + use common_wal::config::object_store::ObjectStoreWalConfig; + use common_wal::options::{ObjectStoreWalOptions, WalOptions}; use frontend::frontend::FrontendOptions; use object_store::config::{FileConfig, GcsConfig}; use servers::grpc::GrpcOptions; @@ -1045,6 +1071,44 @@ mod tests { use super::*; use crate::options::GlobalOptions; + #[tokio::test] + async fn test_build_standalone_wal_provider() { + let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + + let config = ObjectStoreWalConfig { + prefix: "cluster-a/wal".to_string(), + ..Default::default() + }; + // The regions persist the node prefix the datanode runs its store under. + let node_prefix = config.node_prefix(0, STANDALONE_GENERATION); + assert_eq!(node_prefix, "cluster-a/wal/datanodes/0/epochs/0"); + let provider = build_standalone_wal_provider( + &DatanodeWalConfig::ObjectStore(config), + Some(0), + kv_backend.clone(), + ) + .await + .unwrap(); + assert!(matches!( + &provider, + WalProvider::ObjectStore { prefix } if *prefix == node_prefix + )); + let regions = vec![0, 1]; + let wal_options = provider.allocate(®ions, false).await.unwrap(); + for region in regions { + assert_eq!( + wal_options[®ion], + WalOptions::ObjectStore(ObjectStoreWalOptions::new(node_prefix.clone())) + ); + } + + let provider = + build_standalone_wal_provider(&DatanodeWalConfig::default(), Some(0), kv_backend) + .await + .unwrap(); + assert!(matches!(provider, WalProvider::RaftEngine)); + } + #[test] fn test_standalone_local_file_access_config() { let data_home = create_temp_dir("standalone_copy_root"); diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index b3d2c981f1e..8a80083aef3 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -36,9 +36,10 @@ use common_stat::ResourceStatImpl; use common_telemetry::{error, info, warn}; use common_wal::config::DatanodeWalConfig; use common_wal::config::kafka::DatanodeKafkaConfig; -use common_wal::config::object_store::ObjectStoreWalConfig; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use common_wal::config::raft_engine::RaftEngineConfig; use file_engine::engine::FileRegionEngine; +use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; use log_store::kafka::{GlobalIndexCollector, default_index_file}; use log_store::noop::log_store::NoopLogStore; @@ -49,12 +50,14 @@ use mito2::config::MitoConfig; use mito2::engine::{MitoEngine, MitoEngineBuilder}; use mito2::region::opener::PartitionExprFetcherRef; use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; +use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; use query::QueryEngineFactory; use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef}; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; +use store_api::logstore::LogStore; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{ RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState, @@ -67,7 +70,7 @@ use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu, InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu, - ObjectStoreWalNotSupportedSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, + ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ @@ -90,6 +93,7 @@ pub struct Datanode { greptimedb_telemetry_task: Arc, leases_notifier: Option>, plugins: Plugins, + object_store_log_store: Option>, } impl Datanode { @@ -135,21 +139,40 @@ impl Datanode { self.services = services; } + /// Shuts down every component in order. A failed step does not skip the + /// later ones; the first error is returned and the rest are logged. pub async fn shutdown(&mut self) -> Result<()> { - self.services - .shutdown_all() - .await - .context(ShutdownServerSnafu)?; + let mut first_error = None; + record_shutdown_error( + &mut first_error, + self.services + .shutdown_all() + .await + .context(ShutdownServerSnafu), + ); let _ = self.greptimedb_telemetry_task.stop().await; if let Some(heartbeat_task) = &self.heartbeat_task { - heartbeat_task - .close() - .map_err(BoxedError::new) - .context(ShutdownInstanceSnafu)?; + record_shutdown_error( + &mut first_error, + heartbeat_task + .close() + .map_err(BoxedError::new) + .context(ShutdownInstanceSnafu), + ); } - self.region_server.stop().await?; - Ok(()) + record_shutdown_error(&mut first_error, self.region_server.stop().await); + if let Some(log_store) = &self.object_store_log_store { + record_shutdown_error( + &mut first_error, + log_store + .stop() + .await + .map_err(BoxedError::new) + .context(ShutdownInstanceSnafu), + ); + } + first_error.map_or(Ok(()), Err) } pub fn region_server(&self) -> RegionServer { @@ -173,6 +196,8 @@ pub struct DatanodeBuilder { local_file_access: LocalFileAccess, #[cfg(feature = "enterprise")] extension_range_provider_factory: Option, + /// Set as soon as the object store log store exists so that a failed build can stop it. + object_store_log_store: Option>, } impl DatanodeBuilder { @@ -189,6 +214,7 @@ impl DatanodeBuilder { #[cfg(feature = "enterprise")] extension_range_provider_factory: None, topic_stats_reporter: None, + object_store_log_store: None, } } @@ -256,6 +282,22 @@ impl DatanodeBuilder { } pub async fn build(mut self) -> Result { + self.try_build().await + } + + /// Builds the datanode and stops the object store log store if a later step fails. + async fn try_build(&mut self) -> Result { + let result = self.build_inner().await; + if result.is_err() + && let Some(log_store) = &self.object_store_log_store + && let Err(err) = log_store.stop().await + { + warn!(err; "Failed to stop the object store log store after a failed build"); + } + result + } + + async fn build_inner(&mut self) -> Result { validate_region_engine_config(&self.opts.region_engine)?; let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?; set_default_prefix(self.opts.default_column_prefix.as_deref()) @@ -269,6 +311,7 @@ impl DatanodeBuilder { // writable upon open. let controlled_by_metasrv = meta_client.is_some(); if let DatanodeWalConfig::ObjectStore(config) = &self.opts.wal { + ensure!(!controlled_by_metasrv, ObjectStoreWalNotStandaloneSnafu); validate_object_store_wal_config(config)?; } @@ -372,6 +415,7 @@ impl DatanodeBuilder { region_event_receiver, leases_notifier, plugins: self.plugins.clone(), + object_store_log_store: self.object_store_log_store.take(), }) } @@ -660,11 +704,57 @@ impl DatanodeBuilder { builder.try_build().await.context(BuildMitoEngineSnafu)? } - DatanodeWalConfig::ObjectStore(_) => return ObjectStoreWalNotSupportedSnafu.fail(), + DatanodeWalConfig::ObjectStore(object_store_config) => { + let object_store = wal_object_store( + &object_store_manager, + &opts.storage, + &object_store_config.storage_provider, + )?; + let node_id = opts.node_id.context(MissingNodeIdSnafu)?; + let log_store = + Self::build_object_store_log_store(object_store, object_store_config, node_id) + .await?; + self.object_store_log_store = Some(log_store.clone()); + + let builder = MitoEngineBuilder::new( + &opts.storage.data_home, + config, + log_store, + object_store_manager, + schema_metadata_manager, + file_ref_manager, + partition_expr_fetcher, + plugins, + ); + + #[cfg(feature = "enterprise")] + let builder = builder.with_extension_range_provider_factory( + self.extension_range_provider_factory.take(), + ); + + builder.try_build().await.context(BuildMitoEngineSnafu)? + } }; Ok(mito_engine) } + /// Builds [ObjectStoreLogStore] under the node prefix of `node_id` in the + /// standalone generation, the prefix the regions persist in their WAL options. + async fn build_object_store_log_store( + object_store: ObjectStore, + config: &ObjectStoreWalConfig, + node_id: u64, + ) -> Result> { + info!( + "Creating object store logstore with config: {:?}, node id: {}", + config, node_id + ); + ObjectStoreLogStore::try_new(object_store, config, node_id, STANDALONE_GENERATION) + .await + .map_err(Box::new) + .context(OpenLogStoreSnafu) + } + /// Builds [RaftEngineLogStore]. async fn build_raft_engine_log_store( data_home: &str, @@ -714,6 +804,43 @@ impl DatanodeBuilder { } } +/// Keeps the first shutdown error and logs the later ones. +pub(crate) fn record_shutdown_error(first_error: &mut Option, result: Result<()>) { + if let Err(err) = result { + if first_error.is_none() { + *first_error = Some(err); + } else { + warn!(err; "Ignored a later shutdown error"); + } + } +} + +/// Returns the object store that holds the WAL objects: the default one for an +/// empty `storage_provider`, otherwise the configured provider with that name. +fn wal_object_store( + object_store_manager: &ObjectStoreManager, + storage: &StorageConfig, + storage_provider: &str, +) -> Result { + if storage_provider.is_empty() { + return Ok(object_store_manager.default_object_store().clone()); + } + object_store_manager + .find(storage_provider) + .cloned() + .with_context(|| { + let configured = std::iter::once(&storage.store) + .chain(storage.providers.iter()) + .map(|store| store.config_name()) + .collect::>(); + InvalidObjectStoreWalConfigSnafu { + field: "storage_provider", + value: storage_provider, + reason: format!("must name a configured storage provider, one of {configured:?}"), + } + }) +} + /// Rejects repeated engine types before their configs can silently overwrite each other. fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> { let mut engine_indices = HashMap::new(); @@ -923,6 +1050,7 @@ mod tests { use std::io::Write; use std::path::Path; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use cache::build_datanode_cache_registry; @@ -939,15 +1067,22 @@ mod tests { use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use log_store::ObjectStoreLogStore; + use log_store::error::Error as LogStoreError; + use meta_client::client::MetaClientBuilder; use mito2::engine::MITO_ENGINE_NAME; + use object_store::ObjectStore; + use object_store::manager::ObjectStoreManager; + use store_api::logstore::LogStore; use store_api::region_request::RegionRequest; use store_api::storage::RegionId; - use crate::config::{DatanodeOptions, RegionEngineConfig}; + use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::datanode::{ DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config, + wal_object_store, }; - use crate::error::Error; + use crate::error::{self, Error}; use crate::tests::{MockRegionEngine, mock_region_server}; async fn setup_table_datanode(kv: &KvBackendRef) { @@ -1059,6 +1194,14 @@ mod tests { std::fs::read_dir(dir).unwrap().next().is_none() } + /// Returns true if the store rejects even an empty append, which it does only once stopped. + async fn is_stopped(log_store: &ObjectStoreLogStore) -> bool { + matches!( + log_store.append_batch(vec![]).await, + Err(LogStoreError::ObjectStoreWalStopped { .. }) + ) + } + #[test] fn test_validate_unique_region_engine_config() { let mut opts = DatanodeOptions::default(); @@ -1219,18 +1362,185 @@ mod tests { } #[tokio::test] - async fn test_build_rejects_object_store_wal_as_not_supported() { - let data_home = create_temp_dir("object-store-wal-not-supported"); + async fn test_build_rejects_object_store_wal_with_meta_client() { + let data_home = create_temp_dir("object-store-wal-meta-client"); + let meta_client = Arc::new(MetaClientBuilder::datanode_default_options(0).build()); + + let mut builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig::default(), + ); + builder.with_meta_client(meta_client); + let err = build_err(builder).await; + + assert_matches!(err, Error::ObjectStoreWalNotStandalone { .. }); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + // The builder stops before it creates any storage or log store. + assert!(is_empty_dir(data_home.path())); + } + + #[tokio::test] + async fn test_build_rejects_unknown_wal_storage_provider() { + let data_home = create_temp_dir("object-store-wal-unknown-provider"); + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig { + storage_provider: "s3".to_string(), + ..Default::default() + }, + ); + let err = build_err(builder).await; + + let Error::InvalidObjectStoreWalConfig { + field, + value, + reason, + .. + } = &err + else { + panic!("unexpected error: {err:?}"); + }; + assert_eq!(*field, "storage_provider"); + assert_eq!(value, "s3"); + assert!(reason.contains(r#"["File"]"#), "{reason}"); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + // No log store is created under the prefix. + assert!(!data_home.path().join("wal").exists()); + } + + #[test] + fn test_wal_object_store_selects_provider() { + fn fs_object_store(dir: &Path) -> ObjectStore { + let builder = object_store::services::Fs::default().root(dir.to_str().unwrap()); + ObjectStore::new(builder).unwrap() + } + + let default_dir = create_temp_dir("wal-object-store-default"); + let named_dir = create_temp_dir("wal-object-store-named"); + let mut manager = ObjectStoreManager::new("File", fs_object_store(default_dir.path())); + manager.add("Wal-Store", fs_object_store(named_dir.path())); + let storage = StorageConfig::default(); + + let root = |store: ObjectStore| store.info().root().clone(); + assert_eq!( + root(wal_object_store(&manager, &storage, "").unwrap()), + root(manager.default_object_store().clone()) + ); + assert_eq!( + root(wal_object_store(&manager, &storage, "wal-store").unwrap()), + root(manager.find("Wal-Store").unwrap().clone()) + ); + assert_matches!( + wal_object_store(&manager, &storage, "missing").unwrap_err(), + Error::InvalidObjectStoreWalConfig { + field: "storage_provider", + .. + } + ); + } + + #[tokio::test] + async fn test_build_and_shutdown_with_raft_engine_wal() { + let data_home = create_temp_dir("raft-engine-wal"); + let builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::default(), + Arc::new(MemoryKvBackend::new()), + ); + + let mut datanode = builder.build().await.unwrap(); + assert!(datanode.object_store_log_store.is_none()); + assert!(data_home.path().join("wal").exists()); + datanode.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn test_build_and_shutdown_with_object_store_wal() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal"); let builder = object_store_wal_builder( data_home.path().to_str().unwrap(), ObjectStoreWalConfig::default(), ); - let err = build_err(builder).await; - assert_matches!(err, Error::ObjectStoreWalNotSupported { .. }); - assert_eq!(StatusCode::Unsupported, err.status_code()); - // No log store is created under the prefix. - assert!(!data_home.path().join("wal").exists()); + let mut datanode = builder.build().await.unwrap(); + let log_store = datanode.object_store_log_store.clone().unwrap(); + assert!(!is_stopped(&log_store).await); + // The store runs under the node prefix derived from the configured root. + assert!(format!("{log_store:?}").contains(r#"prefix: "wal/datanodes/0/epochs/0""#)); + + datanode.shutdown().await.unwrap(); + assert!(is_stopped(&log_store).await); + } + + #[tokio::test] + async fn test_build_with_named_wal_storage_provider() { + let data_home = create_temp_dir("object-store-wal-named-provider"); + // The default file store is registered under its provider name. + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig { + storage_provider: "file".to_string(), + ..Default::default() + }, + ); + + let mut datanode = builder.build().await.unwrap(); + assert!(datanode.object_store_log_store.is_some()); + datanode.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn test_shutdown_stops_every_engine_and_the_log_store_after_engine_failures() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-shutdown-failure"); + let builder = object_store_wal_builder( + data_home.path().to_str().unwrap(), + ObjectStoreWalConfig::default(), + ); + let mut datanode = builder.build().await.unwrap(); + let log_store = datanode.object_store_log_store.clone().unwrap(); + // Both engines fail to stop, so whichever the region server stops first fails. + let stop_calls = Arc::new(AtomicUsize::new(0)); + for name in ["failing-a", "failing-b"] { + let stop_calls = stop_calls.clone(); + let (engine, _) = MockRegionEngine::with_custom_apply_fn(name, move |engine| { + engine.handle_stop_mock_fn = Some(Box::new(move || { + stop_calls.fetch_add(1, Ordering::Relaxed); + error::UnexpectedSnafu { + violated: "stop failed", + } + .fail() + })); + }); + datanode.region_server().register_engine(engine); + } + + let err = datanode.shutdown().await.unwrap_err(); + + assert_matches!(err, Error::StopRegionEngine { .. }); + assert_eq!(2, stop_calls.load(Ordering::Relaxed)); + assert!(is_stopped(&log_store).await); + } + + #[tokio::test] + async fn test_failed_build_stops_object_store_log_store() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("object-store-wal-failed-build"); + let kv_backend = Arc::new(MemoryKvBackend::new()); + // Regions registered in the metadata but missing on disk fail the region open + // that runs after the log store and the engines are built. + setup_table_datanode(&(kv_backend.clone() as _)).await; + let mut builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()), + kv_backend, + ); + + assert!(builder.try_build().await.is_err()); + + let log_store = builder.object_store_log_store.as_ref().unwrap(); + assert!(is_stopped(log_store).await); } #[test] diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index b62708aa343..539712711b7 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -319,8 +319,8 @@ pub enum Error { location: Location, }, - #[snafu(display("Object store WAL is not supported yet"))] - ObjectStoreWalNotSupported { + #[snafu(display("Object store WAL is only supported in standalone mode"))] + ObjectStoreWalNotStandalone { #[snafu(implicit)] location: Location, }, @@ -499,6 +499,7 @@ impl ErrorExt for Error { | TomlFormat { .. } | DuplicateRegionEngineConfig { .. } | InvalidObjectStoreWalConfig { .. } + | ObjectStoreWalNotStandalone { .. } | BuildDatanode { .. } => StatusCode::InvalidArguments, PayloadNotExist { .. } @@ -523,9 +524,7 @@ impl ErrorExt for Error { OpenLogStore { source, .. } => source.status_code(), MetaClientInit { source, .. } => source.status_code(), - UnsupportedOutput { .. } - | NotYetImplemented { .. } - | ObjectStoreWalNotSupported { .. } => StatusCode::Unsupported, + UnsupportedOutput { .. } | NotYetImplemented { .. } => StatusCode::Unsupported, HandleRegionRequest { source, .. } | GetRegionMetadata { source, .. } | HandleBatchOpenRequest { source, .. } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index aba09de8797..7d5218f30b7 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -92,6 +92,7 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; +use crate::datanode::record_shutdown_error; use crate::error::{ self, BuildRegionRequestsSnafu, ConcurrentQueryLimiterClosedSnafu, ConcurrentQueryLimiterTimeoutSnafu, DataFusionSnafu, DecodeLogicalPlanSnafu, @@ -1948,15 +1949,19 @@ impl RegionServerInner { drop(self.mito_engine.write().unwrap().take()); let engines = self.engines.write().unwrap().drain().collect::>(); + let mut first_error = None; for (engine_name, engine) in engines { - engine + let result = engine .stop() .await - .context(StopRegionEngineSnafu { name: &engine_name })?; - info!("Region engine {engine_name} is stopped"); + .context(StopRegionEngineSnafu { name: &engine_name }); + if result.is_ok() { + info!("Region engine {engine_name} is stopped"); + } + record_shutdown_error(&mut first_error, result); } - Ok(()) + first_error.map_or(Ok(()), Err) } } diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 843ac275220..8884c646b48 100644 --- a/src/datanode/src/tests.rs +++ b/src/datanode/src/tests.rs @@ -121,6 +121,8 @@ pub type MockSyncRegionHandler = Box< dyn Fn(RegionId, SyncRegionFromRequest) -> Result + Send + Sync, >; +pub type MockStopHandler = Box Result<(), Error> + Send + Sync>; + pub struct MockRegionEngine { sender: Sender<(RegionId, RegionRequest)>, pub(crate) handle_request_delay: Option, @@ -128,6 +130,7 @@ pub struct MockRegionEngine { pub(crate) handle_set_readonly_gracefully_mock_fn: Option, pub(crate) handle_get_metadata_mock_fn: Option, pub(crate) handle_sync_region_mock_fn: Option, + pub(crate) handle_stop_mock_fn: Option, pub(crate) mock_role: Option>, engine: String, } @@ -144,6 +147,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -165,6 +169,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -186,6 +191,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: Some(mock_fn), handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }), @@ -208,6 +214,7 @@ impl MockRegionEngine { handle_set_readonly_gracefully_mock_fn: None, handle_get_metadata_mock_fn: None, handle_sync_region_mock_fn: None, + handle_stop_mock_fn: None, mock_role: None, engine: engine.to_string(), }; @@ -269,6 +276,9 @@ impl RegionEngine for MockRegionEngine { } async fn stop(&self) -> Result<(), BoxedError> { + if let Some(mock_fn) = &self.handle_stop_mock_fn { + return mock_fn().map_err(BoxedError::new); + } Ok(()) } diff --git a/src/log-store/src/lib.rs b/src/log-store/src/lib.rs index 5773efa7aaa..0cc133d81e7 100644 --- a/src/log-store/src/lib.rs +++ b/src/log-store/src/lib.rs @@ -19,3 +19,5 @@ pub mod noop; pub mod object_store_wal; pub mod raft_engine; pub mod test_util; + +pub use object_store_wal::ObjectStoreLogStore; diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 3c1958b5cdb..f36108f3bd5 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -56,19 +56,12 @@ //! id name the object that holds the entry, the low bits its position among //! the entries of its region in that object. -#[allow(dead_code)] mod batch; -#[allow(dead_code)] mod catalog; -#[allow(dead_code)] mod format; -#[allow(dead_code)] mod io; - -#[allow(dead_code)] mod store; #[allow(unused_imports)] pub(crate) use batch::entry_id; -#[allow(unused_imports)] -pub(crate) use store::ObjectStoreLogStore; +pub use store::ObjectStoreLogStore; diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index b3d46049523..6dd46c87242 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -42,6 +42,7 @@ const SEGMENT_HEADER_LEN: usize = 8 + 4; /// offset, segment length and segment CRC32. pub(crate) const FOOTER_ENTRY_LEN: usize = 8 + 8 + 8 + 4 + 8 + 8 + 4; /// Length of the entry count the footer starts with. +#[cfg(test)] const FOOTER_COUNT_LEN: usize = 4; /// Header of a WAL object. @@ -185,6 +186,7 @@ pub(crate) fn encode_object(header: Header, records: &[Record]) -> Result u64 { let segments = footer.iter().map(|entry| entry.segment_len).sum::(); let framing = HEADER_LEN + FOOTER_COUNT_LEN + footer.len() * FOOTER_ENTRY_LEN + TRAILER_LEN; diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 13d4220c6ab..2a0c75f1048 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -88,7 +88,7 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// `enqueued` mode it returns on admission and the object is created in the /// background. While [`MAX_SEALED_BATCHES`] batches /// wait to become durable no append is admitted. -pub(crate) struct ObjectStoreLogStore { +pub struct ObjectStoreLogStore { prefix: String, ack_mode: AckMode, io: Arc, @@ -124,7 +124,7 @@ impl ObjectStoreLogStore { /// recovering the catalog from the objects that already exist and writing /// the object that starts the epoch of this instance. Recovery fails on the /// first corrupted or conflicting object. - pub(crate) async fn try_new( + pub async fn try_new( object_store: ObjectStore, config: &ObjectStoreWalConfig, node_id: u64, @@ -336,7 +336,7 @@ fn positive_bytes(bytes: u64, name: &str) -> Result { impl ObjectStoreLogStore { /// Waits until the actor has admitted at least `expected` append calls /// since the store was built. - pub(crate) async fn wait_for_admitted_appends(&self, expected: usize) -> Result<()> { + pub async fn wait_for_admitted_appends(&self, expected: usize) -> Result<()> { self.admitted_appends .clone() .wait_for(|count| *count >= expected) @@ -348,7 +348,7 @@ impl ObjectStoreLogStore { /// Seals the open batch regardless of its size and age and returns once /// its object is durable and indexed, or with the error that failed it. - pub(crate) async fn seal_open_batch(&self) -> Result<()> { + pub async fn seal_open_batch(&self) -> Result<()> { ensure!( !self.stopped.load(Ordering::Acquire), ObjectStoreWalStoppedSnafu @@ -368,24 +368,24 @@ impl ObjectStoreLogStore { /// [`release_creates`](Self::release_creates), so a test can observe /// entries that are admitted but not durable. A create that is parked when /// the store is dropped never runs. - pub(crate) fn hold_creates(&self) { + pub fn hold_creates(&self) { self.creates_held.send_replace(true); } /// Lets the creates parked by [`hold_creates`](Self::hold_creates) run. - pub(crate) fn release_creates(&self) { + pub fn release_creates(&self) { self.creates_held.send_replace(false); } /// Makes every create that runs from now on fail with a transient object /// store error instead of writing. - pub(crate) fn fail_creates(&self) { + pub fn fail_creates(&self) { self.creates_fail.store(true, Ordering::Release); } /// Sets the stopped flag without sending the stop command, which is the /// state a store is in between the two steps of [`stop`](LogStore::stop). - pub(crate) fn begin_stop(&self) { + pub fn begin_stop(&self) { self.stopped.store(true, Ordering::Release); } } diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 46032deef38..39dd96f687b 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -97,7 +97,7 @@ common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true criterion = { workspace = true, features = ["async", "async_tokio"] } dotenv.workspace = true -log-store.workspace = true +log-store = { workspace = true, features = ["testing"] } mito-codec = { workspace = true, features = ["testing"] } object-store = { workspace = true, features = ["services-memory"] } rskafka.workspace = true diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index fbfd339cff7..abfb470c1be 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -51,6 +51,8 @@ pub mod listener; #[cfg(test)] mod merge_mode_test; #[cfg(test)] +mod object_store_wal_test; +#[cfg(test)] mod open_test; #[cfg(test)] mod parallel_test; diff --git a/src/mito2/src/engine/object_store_wal_test.rs b/src/mito2/src/engine/object_store_wal_test.rs new file mode 100644 index 00000000000..d4225b6fcf4 --- /dev/null +++ b/src/mito2/src/engine/object_store_wal_test.rs @@ -0,0 +1,379 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Tests for regions whose WAL options select the object store WAL. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use api::v1::Rows; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; +use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use futures::TryStreamExt; +use log_store::ObjectStoreLogStore; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{ + PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionOpenRequest, + RegionRequest, +}; +use store_api::storage::{RegionId, ScanRequest}; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema, +}; +use crate::wal::EntryId; + +fn wal_config() -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + prefix: "cluster-a/wal".to_string(), + ..Default::default() + } +} + +/// The node prefix a standalone datanode derives from the configured root, +/// which the regions persist in their WAL options. +fn node_prefix() -> String { + wal_config().node_prefix(0, STANDALONE_GENERATION) +} + +fn wal_options() -> HashMap { + let options = WalOptions::ObjectStore(ObjectStoreWalOptions::new(node_prefix())); + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&options).unwrap(), + )]) +} + +fn object_store_provider(region_id: RegionId) -> Provider { + Provider::object_store_provider(region_id, node_prefix()) +} + +fn open_request(table_dir: String) -> RegionOpenRequest { + RegionOpenRequest { + engine: String::new(), + table_dir, + options: wal_options(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + } +} + +async fn new_engine(env: &mut TestEnv) -> (MitoEngine, Arc) { + env.create_engine_with_object_store_wal(MitoConfig::default(), &wal_config()) + .await +} + +/// Creates a region with object store WAL options and returns its table dir and row schema. +async fn create_region( + engine: &MitoEngine, + region_id: RegionId, +) -> (String, Vec) { + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + (table_dir, schema) +} + +/// Creates a region and writes rows `[start, end)` to it in one WAL entry. +async fn create_region_with_rows( + engine: &MitoEngine, + region_id: RegionId, + start: usize, + end: usize, +) -> String { + let (table_dir, schema) = create_region(engine, region_id).await; + put_rows( + engine, + region_id, + Rows { + schema, + rows: build_rows(start, end), + }, + ) + .await; + table_dir +} + +async fn close_region(engine: &MitoEngine, region_id: RegionId) { + engine + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); +} + +async fn count_rows(engine: &MitoEngine, region_id: RegionId) -> usize { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + RecordBatches::try_collect(stream) + .await + .unwrap() + .iter() + .map(|batch| batch.num_rows()) + .sum() +} + +/// Returns the ids of the entries the store still hands out for `region_id`. +async fn stored_entry_ids(log_store: &ObjectStoreLogStore, region_id: RegionId) -> Vec { + log_store + .read(&object_store_provider(region_id), 0, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .map(|entry| entry.entry_id()) + .collect() +} + +#[tokio::test] +async fn test_create_and_reopen_region_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-reopen").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let provider = object_store_provider(region_id); + + let table_dir = create_region_with_rows(&engine, region_id, 0, 3).await; + assert_eq!(provider, engine.get_region(region_id).unwrap().provider); + let last_entry_id = log_store.latest_entry_id(&provider).unwrap(); + assert_eq!( + vec![last_entry_id], + stored_entry_ids(&log_store, region_id).await + ); + + reopen_region(&engine, region_id, table_dir, true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!(provider, region.provider); + assert_eq!( + last_entry_id, + region.version_control.current().last_entry_id + ); + assert_eq!(3, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_topic_latest_entry_id_follows_replay() { + let mut env = TestEnv::with_prefix("object-store-wal-latest-entry-id").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_id = RegionId::new(1, 1); + let provider = object_store_provider(region_id); + + let (table_dir, schema) = create_region(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!(0, region.version_control.current().version.flushed_entry_id); + assert_eq!(0, region.topic_latest_entry_id.load(Ordering::Relaxed)); + + put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: build_rows(0, 3), + }, + ) + .await; + flush_region(&engine, region_id, None).await; + let flushed_entry_id = region.version_control.current().version.flushed_entry_id; + assert_eq!( + flushed_entry_id, + log_store.latest_entry_id(&provider).unwrap() + ); + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows(3, 5), + }, + ) + .await; + let latest_entry_id = log_store.latest_entry_id(&provider).unwrap(); + assert!(latest_entry_id > flushed_entry_id); + + // Replay refills the memtable with the unflushed entry, so the topic latest + // entry id stays at the replay start. + reopen_region(&engine, region_id, table_dir.clone(), true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!( + flushed_entry_id, + region.version_control.current().version.flushed_entry_id + ); + assert_eq!( + latest_entry_id, + region.version_control.current().last_entry_id + ); + assert_eq!( + flushed_entry_id, + region.topic_latest_entry_id.load(Ordering::Relaxed) + ); + assert_eq!(5, count_rows(&engine, region_id).await); + + // Nothing is replayed after a flush, so the topic latest entry id comes from the store. + flush_region(&engine, region_id, None).await; + reopen_region(&engine, region_id, table_dir, true, wal_options()).await; + let region = engine.get_region(region_id).unwrap(); + assert_eq!( + latest_entry_id, + region.version_control.current().version.flushed_entry_id + ); + assert_eq!( + latest_entry_id, + region.topic_latest_entry_id.load(Ordering::Relaxed) + ); + assert_eq!(5, count_rows(&engine, region_id).await); +} + +#[tokio::test] +async fn test_open_region_replays_only_its_own_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-isolation").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + + let table_dir_a = create_region_with_rows(&engine, region_a, 0, 3).await; + let table_dir_b = create_region_with_rows(&engine, region_b, 0, 5).await; + close_region(&engine, region_a).await; + close_region(&engine, region_b).await; + // The store holds the entries of both regions under the same prefix. + assert_eq!(1, stored_entry_ids(&log_store, region_a).await.len()); + assert_eq!(1, stored_entry_ids(&log_store, region_b).await.len()); + + engine + .handle_request(region_a, RegionRequest::Open(open_request(table_dir_a))) + .await + .unwrap(); + assert_eq!(3, count_rows(&engine, region_a).await); + assert!(!engine.is_region_exists(region_b)); + + engine + .handle_request(region_b, RegionRequest::Open(open_request(table_dir_b))) + .await + .unwrap(); + assert_eq!(5, count_rows(&engine, region_b).await); + assert_eq!(3, count_rows(&engine, region_a).await); +} + +#[tokio::test] +async fn test_batch_open_regions_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-batch-open").await; + let (engine, _) = new_engine(&mut env).await; + + let mut regions = Vec::new(); + for i in 1..=3 { + let region_id = RegionId::new(1, i); + let num_rows = i as usize + 1; + let table_dir = create_region_with_rows(&engine, region_id, 0, num_rows).await; + close_region(&engine, region_id).await; + regions.push((region_id, table_dir, num_rows)); + } + + let requests = regions + .iter() + .map(|(region_id, table_dir, _)| (*region_id, open_request(table_dir.clone()))) + .collect(); + let responses = engine + .handle_batch_open_requests(4, requests) + .await + .unwrap(); + assert_eq!(regions.len(), responses.len()); + for (_, response) in responses { + response.unwrap(); + } + + for (region_id, _, num_rows) in regions { + assert_eq!( + object_store_provider(region_id), + engine.get_region(region_id).unwrap().provider + ); + assert_eq!(num_rows, count_rows(&engine, region_id).await); + } +} + +#[tokio::test] +async fn test_drop_and_offline_cleanup_with_object_store_wal() { + let mut env = TestEnv::with_prefix("object-store-wal-cleanup").await; + let (engine, log_store) = new_engine(&mut env).await; + let region_a = RegionId::new(1, 1); + let region_b = RegionId::new(1, 2); + + let _ = create_region_with_rows(&engine, region_a, 0, 3).await; + let table_dir_b = create_region_with_rows(&engine, region_b, 0, 3).await; + + engine + .handle_request( + region_a, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ) + .await + .unwrap(); + // The drop obsoletes the entries of region A only. + assert!(stored_entry_ids(&log_store, region_a).await.is_empty()); + assert_eq!(1, stored_entry_ids(&log_store, region_b).await.len()); + + close_region(&engine, region_b).await; + engine + .handle_request( + region_b, + RegionRequest::CleanUp(RegionCleanUpRequest { + engine: String::new(), + table_dir: table_dir_b, + path_type: PathType::Bare, + options: wal_options(), + }), + ) + .await + .unwrap(); + assert!(stored_entry_ids(&log_store, region_b).await.is_empty()); +} + +#[tokio::test] +async fn test_object_store_wal_options_reject_raft_engine_log_store() { + let mut env = TestEnv::with_prefix("object-store-wal-raft-engine").await; + let engine = env.create_engine(MitoConfig::default()).await; + + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]) + .build(); + let err = engine + .handle_request(RegionId::new(1, 1), RegionRequest::Create(request)) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); +} diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 165f2930f41..d5104bb16b7 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -42,11 +42,13 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_telemetry::{debug, warn}; use common_test_util::temp_dir::{TempDir, create_temp_dir}; +use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; use datatypes::arrow::array::{TimestampMillisecondArray, UInt8Array, UInt64Array}; use datatypes::extension::json::{Json2ExtensionType, JsonExtensionType}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; +use log_store::ObjectStoreLogStore; use log_store::kafka::log_store::KafkaLogStore; use log_store::raft_engine::log_store::RaftEngineLogStore; use log_store::test_util::log_store_util; @@ -56,7 +58,7 @@ use object_store::layers::mock::{ Buffer, Deleter, Metadata, MockLayer, MockLayerBuilder, OpDelete, Result as MockResult, Writer, }; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; -use object_store::services::Fs; +use object_store::services::{Fs, Memory}; use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::client::{Client, ClientBuilder}; use rskafka::record::Record; @@ -492,6 +494,49 @@ impl TestEnv { self.new_mito_engine_with_plugins(config, plugins).await } + /// Creates a new engine with specific config on `log_store` under this env. + pub(crate) async fn create_engine_with_log_store( + &mut self, + config: MitoConfig, + log_store: Arc, + ) -> MitoEngine { + let object_store_manager = Arc::new(self.create_object_store_manager()); + self.object_store_manager = Some(object_store_manager.clone()); + + let data_home = self.data_home().display().to_string(); + MitoEngine::new( + &data_home, + config, + log_store, + object_store_manager, + self.schema_metadata_manager.clone(), + self.file_ref_manager.clone(), + self.partition_expr_fetcher.clone(), + Plugins::new(), + ) + .await + .unwrap() + } + + /// Creates a new engine with specific config on an [ObjectStoreLogStore] that + /// keeps its objects in memory, opened as datanode 0 in the standalone + /// generation. Returns the store so that a test can drive its testing hooks. + pub(crate) async fn create_engine_with_object_store_wal( + &mut self, + config: MitoConfig, + wal_config: &ObjectStoreWalConfig, + ) -> (MitoEngine, Arc) { + let wal_object_store = ObjectStore::new(Memory::default()).unwrap(); + let log_store = + ObjectStoreLogStore::try_new(wal_object_store, wal_config, 0, STANDALONE_GENERATION) + .await + .unwrap(); + let engine = self + .create_engine_with_log_store(config, log_store.clone()) + .await; + (engine, log_store) + } + /// Creates a new engine with specific config and existing logstore and object store manager. pub async fn create_follower_engine(&mut self, config: MitoConfig) -> MitoEngine { self.new_mito_engine(config).await From 78a7b9329271d09b0717bfb53f573077143e243a Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 30 Sep 2026 01:42:11 +0000 Subject: [PATCH 079/133] fix(promql): align timestamp(), label_join and label_replace with Prometheus semantics (#9385) * fix(promql): align timestamp() and label_join with Prometheus semantics - timestamp() over a selector reports the selected sample's timestamp, without adding the offset. - timestamp() over any other expression reports the evaluation time instead of the input value. - label_join that overwrites an existing label rejects duplicate label sets at runtime. Signed-off-by: Dennis Zhuang * fix(promql): align label_replace and label_join edge cases with Prometheus - label_replace leaves non-matching series unchanged instead of copying the source value into the destination label. - label_replace may overwrite an existing label; duplicate label sets are rejected at runtime like label_join. - label_join with no source labels removes the destination label, and an empty source label name is rejected. - Empty label values produced by these functions are NULL, the same as an absent label. Signed-off-by: Dennis Zhuang * fix(promql): join absent labels as empty strings in label_join concat_ws skips NULL arguments together with their separator, so a label that is absent (NULL) dropped its separator, e.g. label_join(label_join(vector(1), "a", "", "missing"), "b", "-", "a", "a") lost b="-". Read every source label as coalesce(label, '') like label_replace does. Signed-off-by: Dennis Zhuang * fix(promql): skip rows without a sample in timestamp() timestamp() replaced the input fields with the timestamp before the empty-value filter ran, so a row whose fields were all NULL got a timestamp, e.g. timestamp(-m) for a NULL sample. Filter such rows before the projection, for both the selector and the evaluation-time paths. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/promql/src/functions/vector_matching.rs | 5 + src/query/src/promql/error.rs | 7 - src/query/src/promql/planner.rs | 305 ++++++++++-------- src/query/src/promql/planner/test.rs | 23 +- src/query/src/promql/planner/test/delta.rs | 4 +- tests-integration/tests/http.rs | 25 ++ .../common/promql/at_modifier.result | 117 +++---- .../standalone/common/promql/at_modifier.sql | 21 +- .../standalone/common/promql/label.result | 199 ++++++++---- .../cases/standalone/common/promql/label.sql | 21 +- .../promql/native_time_selection.result | 9 + .../common/promql/native_time_selection.sql | 2 + .../common/promql/timestamp_fn.result | 89 ++++- .../standalone/common/promql/timestamp_fn.sql | 28 ++ 14 files changed, 553 insertions(+), 302 deletions(-) diff --git a/src/promql/src/functions/vector_matching.rs b/src/promql/src/functions/vector_matching.rs index fdef83f6f05..ef96e8aa5ea 100644 --- a/src/promql/src/functions/vector_matching.rs +++ b/src/promql/src/functions/vector_matching.rs @@ -34,6 +34,8 @@ pub enum MatchGroupViolation { ImplicitManyToOne, /// A group modifier left several matches with the same result label set. AmbiguousGroupLabels, + /// A function rewrote labels so that several series share the same label set. + DuplicateLabelSet, } impl MatchGroupViolation { @@ -54,6 +56,9 @@ impl MatchGroupViolation { Self::AmbiguousGroupLabels => format!( "multiple matches for labels {group}: grouping labels must ensure unique matches" ), + Self::DuplicateLabelSet => { + "vector cannot contain metrics with the same labelset".to_string() + } } } } diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index 6cd45d9e57c..3b24d33612d 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -217,12 +217,6 @@ pub enum Error { location: Location, }, - #[snafu(display("vector cannot contain metrics with the same labelset"))] - SameLabelSet { - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Invalid regular expression in label_replace(): {}", regex))] InvalidRegularExpression { regex: String, @@ -256,7 +250,6 @@ impl ErrorExt for Error { | CombineTableColumnMismatch { .. } | UnexpectedPlanExpr { .. } | UnsupportedMatcherOp { .. } - | SameLabelSet { .. } | TimestampOutOfRange { .. } | SystemTimeOutOfRange { .. } | AtModifierTimestampOutOfRange { .. } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index dc023762b00..88464bc9ee8 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -107,7 +107,7 @@ use crate::promql::error::{ CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, - MultipleVectorSnafu, NoMetricMatcherSnafu, Result, SameLabelSetSnafu, TableNameNotFoundSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, Result, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, }; @@ -1752,47 +1752,9 @@ impl PromPlanner { DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) }) .collect::>(); - // `timestamp()` preserves the shifted selector timeline even though - // SeriesNormalize now retains raw native timestamp storage. Decimal - // arithmetic shifts before truncating to milliseconds. - let unit_factor = match ident(&time_index_column) - .get_type(normalize.schema()) - .context(DataFusionPlanningSnafu)? - { - ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0), - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0), - ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3), - ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6), - _ => unreachable!("time index is a timestamp"), - }; - let sample_time = ident(&time_index_column) - .cast_to(&ArrowDataType::Int64, normalize.schema()) - .context(DataFusionPlanningSnafu)? - .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema()) - .context(DataFusionPlanningSnafu)?; - let sample_time = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(sample_time), - op: Operator::Multiply, - right: Box::new(lit(ScalarValue::Decimal128( - Some(unit_factor.0), - unit_factor.1, - unit_factor.2, - ))), - }); - let sample_time = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(sample_time), - op: Operator::Plus, - right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))), - }) - .cast_to(&ArrowDataType::Int64, normalize.schema()) - .context(DataFusionPlanningSnafu)? - .cast_to(&ArrowDataType::Float64, normalize.schema()) - .context(DataFusionPlanningSnafu)?; - let sample_time = DfExpr::BinaryExpr(BinaryExpr { - left: Box::new(sample_time), - op: Operator::Divide, - right: Box::new(lit(1000.0)), - }); + // The time index still holds the raw sample timestamp here, which is what + // `timestamp()` reports regardless of `offset` and `@`. + let sample_time = Self::timestamp_seconds_expr(&time_index_column, normalize.schema())?; project_exprs.push(sample_time.alias(×tamp_value_column)); let normalize = LogicalPlanBuilder::from(normalize) .project(project_exprs) @@ -1846,40 +1808,51 @@ impl PromPlanner { }), }; if let Some(timestamp_value_column) = timestamp_value_column { - self.create_timestamp_func_plan(manipulate, ×tamp_value_column) + self.create_timestamp_func_plan(manipulate, ident(timestamp_value_column)) } else { Ok(manipulate) } } - /// Builds a projection plan for the PromQL `timestamp()` function. - /// Projects the time index column as the value column for each row. + /// Converts the timestamp column `column` into PromQL seconds, truncated to milliseconds. + fn timestamp_seconds_expr(column: &str, schema: &DFSchema) -> Result { + let column = DfExpr::Column(Column::from_name(column)); + let ArrowDataType::Timestamp(_, timezone) = + column.get_type(schema).context(DataFusionPlanningSnafu)? + else { + unreachable!("time index is a timestamp") + }; + let millis = column + .cast_to( + &ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone), + schema, + ) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Int64, schema) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Float64, schema) + .context(DataFusionPlanningSnafu)?; + Ok(DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(millis), + op: Operator::Divide, + right: Box::new(lit(1000.0)), + })) + } + + /// Builds a projection plan for the PromQL `timestamp()` function, which reports + /// `timestamp_value` as the value of each row, along with the original tag and time index + /// columns. /// - /// # Arguments - /// * `input` - Input [`LogicalPlan`] after instant-vector selection. - /// * `timestamp_value_column` - Private column containing each selected sample's timestamp. - /// - /// # Returns - /// Returns a [`Result`] where the resulting logical plan projects the timestamp - /// column as the value column, along with the original tag and time index columns. - /// - /// # Timestamp vs. Time Function - /// - /// - **Timestamp Function (`timestamp()`)**: In PromQL, the `timestamp()` function returns the - /// timestamp (time index) of each sample as the value column. - /// - /// - **Time Function (`time()`)**: The `time()` function returns the evaluation time of the query - /// as a scalar value. - /// - /// # Side Effects /// Updates the planner context's field columns to the timestamp column name. - /// fn create_timestamp_func_plan( &mut self, input: LogicalPlan, - timestamp_value_column: &str, + timestamp_value: DfExpr, ) -> Result { - let time_expr = ident(timestamp_value_column).alias(DEFAULT_FIELD_COLUMN); + // A row whose fields are all NULL holds no sample, so it must not get a timestamp. The + // check reads the input fields, which the projection below replaces. + let has_sample = self.create_empty_values_filter_expr(true)?; + let time_expr = timestamp_value.alias(DEFAULT_FIELD_COLUMN); self.ctx.field_columns = vec![time_expr.schema_name().to_string()]; let mut project_exprs = Vec::with_capacity(self.ctx.tag_columns.len() + 2); project_exprs.push(self.create_time_index_column_expr()?); @@ -1887,6 +1860,8 @@ impl PromPlanner { project_exprs.extend(self.create_tag_column_exprs()?); LogicalPlanBuilder::from(input) + .filter(has_sample) + .context(DataFusionPlanningSnafu)? .project(project_exprs) .context(DataFusionPlanningSnafu)? .build() @@ -2029,9 +2004,31 @@ impl PromPlanner { // transform function arguments let args = self.create_function_args(&args.args)?; + // Only a vector selector keeps the timestamps of its samples. Any other expression + // produces samples at the evaluation time, which is what `timestamp()` reports for it. + let mut timestamp_arg = args.input.as_ref(); + while let Some(PromExpr::Paren(ParenExpr { expr })) = timestamp_arg { + timestamp_arg = Some(expr); + } + let timestamp_of_selector = + func.name == "timestamp" && matches!(timestamp_arg, Some(PromExpr::VectorSelector(_))); let input = if let Some(prom_expr) = &args.input { - self.prom_expr_to_plan_inner(prom_expr, func.name == "timestamp", query_engine_state) - .await? + let input = self + .prom_expr_to_plan_inner(prom_expr, timestamp_of_selector, query_engine_state) + .await?; + if func.name == "timestamp" && !timestamp_of_selector { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let eval_time = Self::timestamp_seconds_expr(&time_index_column, input.schema())?; + self.create_timestamp_func_plan(input, eval_time)? + } else { + input + } } else { self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); self.ctx.reset_table_name_and_schema(); @@ -2057,6 +2054,7 @@ impl PromPlanner { // `create_function_expr`: the function that reads it (`predict_linear`) then depends on an // argument instead of on planner state written by the selector below it. let range_fold_offset = self.ctx.range_fold_offset.take(); + let input_tag_columns = self.ctx.tag_columns.clone(); let (mut func_exprs, new_tags) = self.create_function_expr( func, args.literals.clone(), @@ -2108,13 +2106,31 @@ impl PromPlanner { _ => builder, }; + // Rewriting a label the input series already have can map several of them onto the same + // label set, which PromQL rejects. A new label keeps the series distinct. + let may_duplicate_label_sets = matches!(func.name, "label_join" | "label_replace") + && new_tags.iter().any(|tag| input_tag_columns.contains(tag)); + // Update context tags after building plan // We can't push them before planning, because they won't exist until projection. for tag in new_tags { self.ctx.tag_columns.push(tag); } - let plan = builder.build().context(DataFusionPlanningSnafu)?; + let mut plan = builder.build().context(DataFusionPlanningSnafu)?; + if may_duplicate_label_sets { + let labels = self.ctx.tag_columns.clone(); + plan = Self::assert_unique_match_group( + plan, + labels + .iter() + .map(|label| DfExpr::Column(Column::from_name(label))) + .collect(), + labels, + self.create_time_index_column_expr()?, + MatchGroupViolation::DuplicateLabelSet, + )?; + } common_telemetry::debug!("Created PromQL function plan: {plan:?} for {call_expr:?}"); Ok(plan) @@ -3620,6 +3636,7 @@ impl PromPlanner { let (concat_expr, dst_label) = Self::build_concat_labels_expr( &mut other_input_exprs, &self.ctx, + input_schema, query_engine_state, )?; @@ -3641,8 +3658,8 @@ impl PromPlanner { } "label_replace" => { self.ctx.use_tsid = false; - if let Some((replace_expr, dst_label)) = self - .build_regexp_replace_label_expr(&mut other_input_exprs, query_engine_state)? + if let Some((replace_expr, dst_label)) = + self.build_regexp_replace_label_expr(&mut other_input_exprs, input_schema)? { // Reserve the current field columns except the `dst_label`. for value in &self.ctx.field_columns { @@ -3652,10 +3669,8 @@ impl PromPlanner { } } - ensure!( - !self.ctx.tag_columns.contains(&dst_label), - SameLabelSetSnafu - ); + // Remove it from tag columns if exists to avoid duplicated column names + self.ctx.tag_columns.retain(|tag| *tag != dst_label); new_tags.push(dst_label); // Add the new label expr to evaluate exprs.push(replace_expr); @@ -3964,7 +3979,7 @@ impl PromPlanner { fn build_regexp_replace_label_expr( &self, other_input_exprs: &mut VecDeque, - query_engine_state: &QueryEngineState, + input_schema: &DFSchemaRef, ) -> Result> { // label_replace(vector, dst_label, replacement, src_label, regex) let dst_label = match other_input_exprs.pop_front() { @@ -4000,72 +4015,64 @@ impl PromPlanner { .fail()?, }; - // Validate the regex before using it + // Like Prometheus, match the whole source value. A series whose source value matches gets + // `dst_label` set to the expanded replacement; any other series is left unchanged. // doc: https://prometheus.io/docs/prometheus/latest/querying/functions/#label_replace - regex::Regex::new(®ex).map_err(|_| { - InvalidRegularExpressionSnafu { - regex: regex.clone(), - } - .build() - })?; + let anchored = format!("^(?s:{regex})$"); + let compiled = regex::Regex::new(&anchored) + .map_err(|_| InvalidRegularExpressionSnafu { regex }.build())?; + let dst_exists = self.ctx.tag_columns.contains(&dst_label); - // If the src_label exists and regex is empty, keep everything unchanged. - if self.ctx.tag_columns.contains(&src_label) && regex.is_empty() { - return Ok(None); - } - - // If the src_label doesn't exists, and + // A missing source label reads as the empty string for every series, so the result is + // the same for all of them and is decided here. if !self.ctx.tag_columns.contains(&src_label) { - if replacement.is_empty() { - // the replacement is empty, keep everything unchanged. + let Some(captures) = compiled.captures("") else { return Ok(None); - } else { - // the replacement is not empty, always adds dst_label with replacement value. + }; + let mut value = String::new(); + captures.expand(&replacement, &mut value); + if value.is_empty() { + // Setting a label to the empty string removes it, which is a no-op for a new + // label. + if !dst_exists { + return Ok(None); + } return Ok(Some(( - // alias literal `replacement` as dst_label - lit(replacement).alias(&dst_label), + lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label, ))); } + return Ok(Some((lit(value).alias(&dst_label), dst_label))); } - // Preprocess the regex: - // https://github.com/prometheus/prometheus/blob/d902abc50d6652ba8fe9a81ff8e5cce936114eba/promql/functions.go#L1575C32-L1575C37 - let regex = format!("^(?s:{regex})$"); + let src = Self::label_value_expr(&src_label, input_schema)?; + let matched = DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::regex::regexp_like(), + args: vec![src.clone(), lit(anchored.clone())], + }); + let replaced = Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::regex::regexp_replace(), + args: vec![src, lit(anchored), lit(replacement)], + })); + let unchanged = if dst_exists { + DfExpr::Column(Column::from_name(&dst_label)) + .cast_to(&ArrowDataType::Utf8, input_schema) + .context(DataFusionPlanningSnafu)? + } else { + lit(ScalarValue::Utf8(None)) + }; + let replace_expr = when(matched, replaced) + .otherwise(unchanged) + .context(DataFusionPlanningSnafu)?; - let session_state = query_engine_state.session_state(); - let func = session_state - .scalar_functions() - .get("regexp_replace") - .context(UnsupportedExprSnafu { - name: "regexp_replace", - })?; - - // regexp_replace(src_label, regex, replacement) - let args = vec![ - if src_label.is_empty() { - DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) - } else { - DfExpr::Column(Column::from_name(src_label)) - }, - DfExpr::Literal(ScalarValue::Utf8(Some(regex)), None), - DfExpr::Literal(ScalarValue::Utf8(Some(replacement)), None), - ]; - - Ok(Some(( - DfExpr::ScalarFunction(ScalarFunction { - func: func.clone(), - args, - }) - .alias(&dst_label), - dst_label, - ))) + Ok(Some((replace_expr.alias(&dst_label), dst_label))) } /// Build expr for `label_join` function fn build_concat_labels_expr( other_input_exprs: &mut VecDeque, ctx: &PromPlannerContext, + input_schema: &DFSchemaRef, query_engine_state: &QueryEngineState, ) -> Result<(DfExpr, String)> { // label_join(vector, dst_label, separator, src_label_1, src_label_2, ...) @@ -4097,17 +4104,19 @@ impl PromPlanner { let src_labels = other_input_exprs .iter() .map(|expr| { - // Cast source label into column or null literal + // `concat_ws` skips NULL arguments together with their separator, while an + // absent label joins as the empty string. match expr { DfExpr::Literal(ScalarValue::Utf8(Some(label)), None) => { if label.is_empty() { - Ok(DfExpr::Literal(ScalarValue::Null, None)) + FunctionInvalidArgumentSnafu { + fn_name: "label_join", + } + .fail() } else if available_columns.contains(label.as_str()) { - // Label exists in the table schema - Ok(DfExpr::Column(Column::from_name(label))) + Self::label_value_expr(label, input_schema) } else { - // Label doesn't exist, treat as empty string (null) - Ok(DfExpr::Literal(ScalarValue::Null, None)) + Ok(lit("")) } } other => UnexpectedPlanExprSnafu { @@ -4120,12 +4129,10 @@ impl PromPlanner { } }) .collect::>>()?; - ensure!( - !src_labels.is_empty(), - FunctionInvalidArgumentSnafu { - fn_name: "label_join" - } - ); + // Joining no labels yields the empty string, i.e. removes `dst_label`. + if src_labels.is_empty() { + return Ok((lit(ScalarValue::Utf8(None)).alias(&dst_label), dst_label)); + } let session_state = query_engine_state.session_state(); let func = session_state @@ -4139,15 +4146,37 @@ impl PromPlanner { args.extend(src_labels); Ok(( - DfExpr::ScalarFunction(ScalarFunction { + Self::empty_label_to_null(DfExpr::ScalarFunction(ScalarFunction { func: func.clone(), args, - }) + })) .alias(&dst_label), dst_label, )) } + /// The value of `label` as a string, where NULL (the series has no such label) reads as the + /// empty string, as in PromQL. + fn label_value_expr(label: &str, input_schema: &DFSchemaRef) -> Result { + let value = DfExpr::Column(Column::from_name(label)) + .cast_to(&ArrowDataType::Utf8, input_schema) + .context(DataFusionPlanningSnafu)?; + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![value, lit("")], + })) + } + + /// An empty label value means the label is absent in PromQL. Label functions represent it + /// as NULL, like a series that never had the label, so that both compare equal when label + /// sets are matched and neither is reported as a label. + fn empty_label_to_null(value: DfExpr) -> DfExpr { + DfExpr::ScalarFunction(ScalarFunction { + func: datafusion_functions::core::nullif(), + args: vec![value, lit("")], + }) + } + fn create_time_index_column_expr(&self) -> Result { Ok(DfExpr::Column(Column::from_name( self.ctx diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 382f4791e28..215acb85cf9 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1699,12 +1699,13 @@ async fn single_timestamp_plan_preserves_source_value() { "Filter: value IS NOT NULL [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(1,1,0) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ - \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", + \n Filter: some_metric.field_0 IS NOT NULL [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ + \n TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]", ); assert_eq!(plan.display_indent_schema().to_string(), expected); @@ -2413,9 +2414,9 @@ async fn at_modifier_keeps_multi_series_roots_out_of_promoted_subtree() { async fn at_modifier_does_not_promote_label_join() { for query in [ // Directly above the anchored instant selector... - "label_join(some_metric @ 300, \"tag_0\", \"-\", \"\")", + "label_join(some_metric @ 300, \"tag_0\", \"-\", \"tag_0\", \"tag_0\")", // ... and below another call, which is planned as usual over the join. - "abs(label_join(some_metric @ 300, \"tag_0\", \"-\", \"\"))", + "abs(label_join(some_metric @ 300, \"tag_0\", \"-\", \"tag_0\", \"tag_0\"))", ] { let plan = build_at_modifier_plan(query, 0, 1000).await; let plan_str = plan.display_indent_schema().to_string(); @@ -2448,7 +2449,7 @@ async fn at_modifier_does_not_promote_label_join() { // A range call below the join is still promoted on its own: the anchored window is folded // once per series, and the join above it is evaluated at every step over that replay. let plan = build_at_modifier_plan( - "label_join(rate(some_metric[5m] @ 300), \"tag_0\", \"-\", \"\")", + "label_join(rate(some_metric[5m] @ 300), \"tag_0\", \"-\", \"tag_0\", \"tag_0\")", 0, 1000, ) @@ -5902,7 +5903,7 @@ async fn test_label_join() { let expected = r#" Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] - Projection: up.timestamp, up.field_0, concat_ws(Utf8(","), up.tag_1, up.tag_2, up.tag_3) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] + Projection: up.timestamp, up.field_0, nullif(concat_ws(Utf8(","), coalesce(up.tag_1, Utf8("")), coalesce(up.tag_2, Utf8("")), coalesce(up.tag_3, Utf8(""))), Utf8("")) AS foo, up.tag_0, up.tag_1, up.tag_2, up.tag_3 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8] PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] PromSeriesDivide: tags=["tag_0", "tag_1", "tag_2", "tag_3"] [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] Sort: up.tag_0 ASC NULLS FIRST, up.tag_1 ASC NULLS FIRST, up.tag_2 ASC NULLS FIRST, up.tag_3 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, tag_1:Utf8, tag_2:Utf8, tag_3:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] @@ -5937,7 +5938,7 @@ async fn test_label_replace() { let expected = r#" Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] - Projection: up.timestamp, up.field_0, regexp_replace(up.tag_0, Utf8("^(?s:(.*):.*)$"), Utf8("$1")) AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] + Projection: up.timestamp, up.field_0, CASE WHEN regexp_like(coalesce(up.tag_0, Utf8("")), Utf8("^(?s:(.*):.*)$")) THEN nullif(regexp_replace(coalesce(up.tag_0, Utf8("")), Utf8("^(?s:(.*):.*)$"), Utf8("$1")), Utf8("")) ELSE Utf8(NULL) END AS foo, up.tag_0 [timestamp:Timestamp(ms), field_0:Float64;N, foo:Utf8;N, tag_0:Utf8] PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] Sort: up.tag_0 ASC NULLS FIRST, up.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N] diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index a647821e0b3..f0c918351d1 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -1102,12 +1102,12 @@ async fn delta_offsets_survive_optimized_plan_serialization() { ( "timestamp positive offset", r#"timestamp(delta_metric{series="cumulative"} offset 60s)"#, - 120.0, + 60.0, ), ( "timestamp negative offset", r#"timestamp(delta_metric{series="cumulative"} offset -60s)"#, - 120.0, + 180.0, ), ( "range positive offset", diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index a05c58444e3..9b6ef9d8ae2 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -1022,6 +1022,31 @@ pub async fn test_prometheus_label_replace_response(store_type: StorageType) { .unwrap() ); + // A series whose source value does not match keeps its labels, without `host_copy`. + let query = encode(r#"label_replace(demo, "host_copy", "$1", "host", "other(.*)")"#); + let res = client + .get(&format!("/v1/prometheus/api/v1/query?query={query}&time=0")) + .send() + .await; + + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!(body.status, "success"); + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "vector", + "result": [{ + "metric": { + "__name__": "demo", + "host": "host1" + }, + "value": [0.0, "1.1"] + }] + })) + .unwrap() + ); + guard.remove_all().await; } diff --git a/tests/cases/standalone/common/promql/at_modifier.result b/tests/cases/standalone/common/promql/at_modifier.result index 0a19eac692c..27570231833 100644 --- a/tests/cases/standalone/common/promql/at_modifier.result +++ b/tests/cases/standalone/common/promql/at_modifier.result @@ -621,9 +621,6 @@ TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset -- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every -- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects -- the timestamp above it. --- --- Only an `@` without an `offset` is asserted here: the value of `timestamp()` for an anchored --- selector that also carries an `offset` is a pre-existing question of its own, out of scope here. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); @@ -642,95 +639,105 @@ TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); | 1970-01-01T00:04:00 | 300.0 | b | +---------------------+-------+------+ +-- With an `offset`, the anchor selects the sample at 240s, and `timestamp()` reports 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 120, '60s') timestamp(at_modifier_gauge @ 300 offset 1m); + ++---------------------+-------+------+ +| ts | value | host | ++---------------------+-------+------+ +| 1970-01-01T00:00:00 | 240.0 | a | +| 1970-01-01T00:00:00 | 240.0 | b | +| 1970-01-01T00:01:00 | 240.0 | a | +| 1970-01-01T00:01:00 | 240.0 | b | +| 1970-01-01T00:02:00 | 240.0 | a | +| 1970-01-01T00:02:00 | 240.0 | b | ++---------------------+-------+------+ + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED -- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | Projection: at_modifier_gauge.ts, value AS value, at_modifier_gauge.host | -| | Projection: at_modifier_gauge.ts, __promql_timestamp_value_ AS value, at_modifier_gauge.host | -| | PromInstantManipulate: range=[0..240000], lookback=[240001], interval=[60000], time index=[ts] | -| | PromInstantManipulate: range=[0..0], lookback=[300000], interval=[60000], time index=[ts] | -| | Projection: at_modifier_gauge.ts, at_modifier_gauge.val, at_modifier_gauge.host, CAST(CAST(CAST(CAST(CAST(at_modifier_gauge.ts AS Int64) AS Decimal128(19, 0)) AS Decimal128(21, 0)) + Decimal128(0,19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ | -| | PromSeriesNormalize: offset=[-300000], time index=[ts], filter NaN: [false] | -| | PromSeriesDivide: tags=["host"] | -| | Sort: at_modifier_gauge.host ASC NULLS FIRST, at_modifier_gauge.ts ASC NULLS FIRST | -| | Filter: at_modifier_gauge.ts >= TimestampMillisecond(1, None) AND at_modifier_gauge.ts <= TimestampMillisecond(300000, None) | -| | TableScan: at_modifier_gauge, partial_filters=[at_modifier_gauge.ts >= TimestampMillisecond(1, None), at_modifier_gauge.ts <= TimestampMillisecond(300000, None)] | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Projection: at_modifier_gauge.ts, value AS value, at_modifier_gauge.host | +| | Projection: at_modifier_gauge.ts, __promql_timestamp_value_ AS value, at_modifier_gauge.host | +| | Filter: at_modifier_gauge.val IS NOT NULL | +| | PromInstantManipulate: range=[0..240000], lookback=[240001], interval=[60000], time index=[ts] | +| | PromInstantManipulate: range=[0..0], lookback=[300000], interval=[60000], time index=[ts] | +| | Projection: at_modifier_gauge.ts, at_modifier_gauge.val, at_modifier_gauge.host, CAST(CAST(at_modifier_gauge.ts AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ | +| | PromSeriesNormalize: offset=[-300000], time index=[ts], filter NaN: [false] | +| | PromSeriesDivide: tags=["host"] | +| | Sort: at_modifier_gauge.host ASC NULLS FIRST, at_modifier_gauge.ts ASC NULLS FIRST | +| | Filter: at_modifier_gauge.ts >= TimestampMillisecond(1, None) AND at_modifier_gauge.ts <= TimestampMillisecond(300000, None) | +| | TableScan: at_modifier_gauge, partial_filters=[at_modifier_gauge.ts >= TimestampMillisecond(1, None), at_modifier_gauge.ts <= TimestampMillisecond(300000, None)] | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- 11. `label_join` above an anchored selector. The call rewrites the label the input series are --- told apart by (`host` becomes the empty string), so it is never promoted: it must stay above the +-- told apart by (`host` becomes `a-a` and `b-b`), so it is never promoted: it must stay above the -- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, -- with their own value, at every step. A promoted `label_join` would instead be replayed through -- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row -- per step. --- --- The invariant asserted here is scoped to the replay: it must not drop or mix the input rows. --- It is not a claim about the final PromQL semantics of this query: joining `host` to one value --- leaves two samples with the same label set at the same timestamp, which Prometheus rejects, --- while `label_join` does not validate that yet (`label_replace` errors on such a rewrite --- instead). That duplicate-labelset validation gap is pre-existing and out of scope here. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "", ""); +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "-", "host", "host"); +---------------------+------+------+ | ts | val | host | +---------------------+------+------+ -| 1970-01-01T00:05:00 | 12.0 | | -| 1970-01-01T00:05:00 | 5.0 | | -| 1970-01-01T00:06:00 | 12.0 | | -| 1970-01-01T00:06:00 | 5.0 | | -| 1970-01-01T00:07:00 | 12.0 | | -| 1970-01-01T00:07:00 | 5.0 | | -| 1970-01-01T00:08:00 | 12.0 | | -| 1970-01-01T00:08:00 | 5.0 | | +| 1970-01-01T00:05:00 | 12.0 | b-b | +| 1970-01-01T00:05:00 | 5.0 | a-a | +| 1970-01-01T00:06:00 | 12.0 | b-b | +| 1970-01-01T00:06:00 | 5.0 | a-a | +| 1970-01-01T00:07:00 | 12.0 | b-b | +| 1970-01-01T00:07:00 | 5.0 | a-a | +| 1970-01-01T00:08:00 | 12.0 | b-b | +| 1970-01-01T00:08:00 | 5.0 | a-a | +---------------------+------+------+ -- The same join below another call: neither is promoted, and both are evaluated at every step over -- the per-series replay of the anchored selector. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "", "")); +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "-", "host", "host")); +---------------------+----------+------+ | ts | abs(val) | host | +---------------------+----------+------+ -| 1970-01-01T00:05:00 | 12.0 | | -| 1970-01-01T00:05:00 | 5.0 | | -| 1970-01-01T00:06:00 | 12.0 | | -| 1970-01-01T00:06:00 | 5.0 | | -| 1970-01-01T00:07:00 | 12.0 | | -| 1970-01-01T00:07:00 | 5.0 | | -| 1970-01-01T00:08:00 | 12.0 | | -| 1970-01-01T00:08:00 | 5.0 | | +| 1970-01-01T00:05:00 | 12.0 | b-b | +| 1970-01-01T00:05:00 | 5.0 | a-a | +| 1970-01-01T00:06:00 | 12.0 | b-b | +| 1970-01-01T00:06:00 | 5.0 | a-a | +| 1970-01-01T00:07:00 | 12.0 | b-b | +| 1970-01-01T00:07:00 | 5.0 | a-a | +| 1970-01-01T00:08:00 | 12.0 | b-b | +| 1970-01-01T00:08:00 | 5.0 | a-a | +---------------------+----------+------+ -- A range call below the join is still promoted on its own (it is the direct call over the -- anchored range selector): the anchored window is folded once per series, and the join above that -- replay reports the rate of both hosts at every step. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "", ""); +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "-", "host", "host"); +---------------------+------------------------------------------+------+ | ts | prom_rate(ts_range,val,ts,Int64(300000)) | host | +---------------------+------------------------------------------+------+ -| 1970-01-01T00:05:00 | 0.5 | | -| 1970-01-01T00:05:00 | 1.0 | | -| 1970-01-01T00:06:00 | 0.5 | | -| 1970-01-01T00:06:00 | 1.0 | | -| 1970-01-01T00:07:00 | 0.5 | | -| 1970-01-01T00:07:00 | 1.0 | | -| 1970-01-01T00:08:00 | 0.5 | | -| 1970-01-01T00:08:00 | 1.0 | | +| 1970-01-01T00:05:00 | 0.5 | b-b | +| 1970-01-01T00:05:00 | 1.0 | a-a | +| 1970-01-01T00:06:00 | 0.5 | b-b | +| 1970-01-01T00:06:00 | 1.0 | a-a | +| 1970-01-01T00:07:00 | 0.5 | b-b | +| 1970-01-01T00:07:00 | 1.0 | a-a | +| 1970-01-01T00:08:00 | 0.5 | b-b | +| 1970-01-01T00:08:00 | 1.0 | a-a | +---------------------+------------------------------------------+------+ DROP TABLE at_modifier_gauge; diff --git a/tests/cases/standalone/common/promql/at_modifier.sql b/tests/cases/standalone/common/promql/at_modifier.sql index 9e9bc21a9c1..b4fd9d1a552 100644 --- a/tests/cases/standalone/common/promql/at_modifier.sql +++ b/tests/cases/standalone/common/promql/at_modifier.sql @@ -258,12 +258,13 @@ TQL EVAL (330, 390, '60s') predict_linear(at_modifier_gauge{host="a"}[5m] offset -- on its own and `timestamp()` reports the timestamp of the sample the anchor selected at every -- step (300s here). Its plan keeps the anchored selection below the replay of the grid and projects -- the timestamp above it. --- --- Only an `@` without an `offset` is asserted here: the value of `timestamp()` for an anchored --- selector that also carries an `offset` is a pre-existing question of its own, out of scope here. -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); +-- With an `offset`, the anchor selects the sample at 240s, and `timestamp()` reports 240s. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 120, '60s') timestamp(at_modifier_gauge @ 300 offset 1m); + -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -271,30 +272,24 @@ TQL EVAL (0, 240, '60s') timestamp(at_modifier_gauge @ 300); TQL EXPLAIN (0, 240, '60s') timestamp(at_modifier_gauge @ 300); -- 11. `label_join` above an anchored selector. The call rewrites the label the input series are --- told apart by (`host` becomes the empty string), so it is never promoted: it must stay above the +-- told apart by (`host` becomes `a-a` and `b-b`), so it is never promoted: it must stay above the -- per-series replay of the selector and be evaluated at every step. Both hosts are still reported, -- with their own value, at every step. A promoted `label_join` would instead be replayed through -- the labels it just rewrote and report the two series as a single timeline, i.e. one mixed row -- per step. --- --- The invariant asserted here is scoped to the replay: it must not drop or mix the input rows. --- It is not a claim about the final PromQL semantics of this query: joining `host` to one value --- leaves two samples with the same label set at the same timestamp, which Prometheus rejects, --- while `label_join` does not validate that yet (`label_replace` errors on such a rewrite --- instead). That duplicate-labelset validation gap is pre-existing and out of scope here. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "", ""); +TQL EVAL (300, 480, '60s') label_join(at_modifier_gauge @ 300, "host", "-", "host", "host"); -- The same join below another call: neither is promoted, and both are evaluated at every step over -- the per-series replay of the anchored selector. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "", "")); +TQL EVAL (300, 480, '60s') abs(label_join(at_modifier_gauge @ 300, "host", "-", "host", "host")); -- A range call below the join is still promoted on its own (it is the direct call over the -- anchored range selector): the anchored window is folded once per series, and the join above that -- replay reports the rate of both hosts at every step. -- SQLNESS SORT_RESULT 3 1 -TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "", ""); +TQL EVAL (300, 480, '60s') label_join(rate(at_modifier_counter_total[5m] @ 300), "host", "-", "host", "host"); DROP TABLE at_modifier_gauge; diff --git a/tests/cases/standalone/common/promql/label.result b/tests/cases/standalone/common/promql/label.result index 112b6261715..3fe8d559221 100644 --- a/tests/cases/standalone/common/promql/label.result +++ b/tests/cases/standalone/common/promql/label.result @@ -21,9 +21,23 @@ INSERT INTO TABLE test VALUES Affected Rows: 8 -- Missing source labels -- +-- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-"); -Error: 1004(InvalidArguments), Invalid function argument for label_join ++---------------------+-----+----------+-------+------------+ +| ts | val | new_host | host | idc | ++---------------------+-----+----------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+----------+-------+------------+ -- dst_label is equal to source label -- -- SQLNESS SORT_RESULT 3 1 @@ -63,24 +77,27 @@ TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "idc", "host" | 1970-01-01T00:00:15 | 7 | idc4:zone3-host1 | idc4:zone3 | +---------------------+-----+------------------+------------+ --- test the empty source label -- --- SQLNESS SORT_RESULT 3 1 +-- Absent labels join as empty strings and keep their separators -- +TQL EVAL (0, 15, '5s') label_join(label_join(vector(1), "a", "", "missing"), "b", "-", "a", "a"); + ++---------------------+----------------+---+---+ +| time | greptime_value | b | a | ++---------------------+----------------+---+---+ +| 1970-01-01T00:00:00 | 1.0 | - | | +| 1970-01-01T00:00:05 | 1.0 | - | | +| 1970-01-01T00:00:10 | 1.0 | - | | +| 1970-01-01T00:00:15 | 1.0 | - | | ++---------------------+----------------+---+---+ + +-- An empty source label name is invalid -- TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", ""); -+---------------------+-----+------+------------+ -| ts | val | host | idc | -+---------------------+-----+------+------------+ -| 1970-01-01T00:00:00 | 1 | | idc1 | -| 1970-01-01T00:00:05 | 1 | | idc1 | -| 1970-01-01T00:00:05 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | | idc1 | -| 1970-01-01T00:00:10 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | | idc1 | -| 1970-01-01T00:00:15 | 3 | | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | | idc4:zone3 | -+---------------------+-----+------+------------+ +Error: 1004(InvalidArguments), Invalid function argument for label_join + +-- Both hosts are joined to the same `host` at 0s, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_join(test, "host", "-", "idc"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "host"); @@ -114,13 +131,13 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "$2", "idc", +---------------------+-----+---------+-------+------------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | | 1970-01-01T00:00:05 | 3 | zone1 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | | 1970-01-01T00:00:10 | 3 | zone1 | host1 | idc2:zone1 | | 1970-01-01T00:00:10 | 5 | zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | | 1970-01-01T00:00:15 | 3 | zone1 | host1 | idc2:zone1 | | 1970-01-01T00:00:15 | 5 | zone2 | host1 | idc3:zone2 | | 1970-01-01T00:00:15 | 7 | zone3 | host1 | idc4:zone3 | @@ -129,20 +146,20 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "$2", "idc", -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "idc99", "idc", "idc2.*"); -+---------------------+-----+------------+-------+------------+ -| ts | val | new_idc | host | idc | -+---------------------+-----+------------+-------+------------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | idc3:zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 3 | idc99 | host1 | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | idc3:zone2 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | idc4:zone3 | host1 | idc4:zone3 | -+---------------------+-----+------------+-------+------------+ ++---------------------+-----+---------+-------+------------+ +| ts | val | new_idc | host | idc | ++---------------------+-----+---------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | idc99 | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+---------+-------+------------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", "(.*):(.*)"); @@ -150,26 +167,61 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", +---------------------+-----+---------+-------+------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------+ -| 1970-01-01T00:00:00 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:05 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:05 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:10 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:10 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:10 | 6 | idc3 | host2 | idc3 | -| 1970-01-01T00:00:15 | 2 | idc1 | host2 | idc1 | -| 1970-01-01T00:00:15 | 4 | idc2 | host2 | idc2 | -| 1970-01-01T00:00:15 | 6 | idc3 | host2 | idc3 | -| 1970-01-01T00:00:15 | 8 | idc4 | host2 | idc4 | +| 1970-01-01T00:00:00 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:05 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:05 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:10 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:10 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:10 | 6 | | host2 | idc3 | +| 1970-01-01T00:00:15 | 2 | | host2 | idc1 | +| 1970-01-01T00:00:15 | 4 | | host2 | idc2 | +| 1970-01-01T00:00:15 | 6 | | host2 | idc3 | +| 1970-01-01T00:00:15 | 8 | | host2 | idc4 | +---------------------+-----+---------+-------+------+ -- dst_label is equal to source label -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc", "$2", "idc", "(.*):(.*)"); -Error: 1004(InvalidArguments), vector cannot contain metrics with the same labelset ++---------------------+-----+------+-------+ +| ts | val | idc | host | ++---------------------+-----+------+-------+ +| 1970-01-01T00:00:00 | 2 | idc1 | host2 | +| 1970-01-01T00:00:05 | 2 | idc1 | host2 | +| 1970-01-01T00:00:05 | 4 | idc2 | host2 | +| 1970-01-01T00:00:10 | 2 | idc1 | host2 | +| 1970-01-01T00:00:10 | 4 | idc2 | host2 | +| 1970-01-01T00:00:10 | 6 | idc3 | host2 | +| 1970-01-01T00:00:15 | 2 | idc1 | host2 | +| 1970-01-01T00:00:15 | 4 | idc2 | host2 | +| 1970-01-01T00:00:15 | 6 | idc3 | host2 | +| 1970-01-01T00:00:15 | 8 | idc4 | host2 | ++---------------------+-----+------+-------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "idc", "$2", "idc", "(.*):(.*)"); + ++---------------------+-----+-------+-------+ +| ts | val | idc | host | ++---------------------+-----+-------+-------+ +| 1970-01-01T00:00:00 | 1 | idc1 | host1 | +| 1970-01-01T00:00:05 | 1 | idc1 | host1 | +| 1970-01-01T00:00:05 | 3 | zone1 | host1 | +| 1970-01-01T00:00:10 | 1 | idc1 | host1 | +| 1970-01-01T00:00:10 | 3 | zone1 | host1 | +| 1970-01-01T00:00:10 | 5 | zone2 | host1 | +| 1970-01-01T00:00:15 | 1 | idc1 | host1 | +| 1970-01-01T00:00:15 | 3 | zone1 | host1 | +| 1970-01-01T00:00:15 | 5 | zone2 | host1 | +| 1970-01-01T00:00:15 | 7 | zone3 | host1 | ++---------------------+-----+-------+-------+ + +-- Both hosts are rewritten to the same `host`, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_replace(test, "host", "x", "host", ".*"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset -- test the empty source label -- --- TODO(dennis): we can't remove the label currently -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc2", "", "", ""); @@ -221,20 +273,20 @@ TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} + label_replace(vector(1), -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "host2", "host", ""); -+---------------------+-----+-------+------------+ -| ts | val | host | idc | -+---------------------+-----+-------+------------+ -| 1970-01-01T00:00:00 | 1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:10 | 5 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 3 | host1 | idc2:zone1 | -| 1970-01-01T00:00:15 | 5 | host1 | idc3:zone2 | -| 1970-01-01T00:00:15 | 7 | host1 | idc4:zone3 | -+---------------------+-----+-------+------------+ ++---------------------+-----+-------+-------+------------+ +| ts | val | host2 | host | idc | ++---------------------+-----+-------+-------+------------+ +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:10 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 3 | | host1 | idc2:zone1 | +| 1970-01-01T00:00:15 | 5 | | host1 | idc3:zone2 | +| 1970-01-01T00:00:15 | 7 | | host1 | idc4:zone3 | ++---------------------+-----+-------+-------+------------+ -- Empty regex with not existing source label -- SQLNESS SORT_RESULT 3 1 @@ -274,6 +326,23 @@ TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "", "instance", | 1970-01-01T00:00:15 | 7 | host1 | idc4:zone3 | +---------------------+-----+-------+------------+ +-- An empty replacement removes an existing label -- +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{idc="idc1"}, "idc", "", "instance", ""); + ++---------------------+-----+-----+-------+ +| ts | val | idc | host | ++---------------------+-----+-----+-------+ +| 1970-01-01T00:00:00 | 1 | | host1 | +| 1970-01-01T00:00:00 | 2 | | host2 | +| 1970-01-01T00:00:05 | 1 | | host1 | +| 1970-01-01T00:00:05 | 2 | | host2 | +| 1970-01-01T00:00:10 | 1 | | host1 | +| 1970-01-01T00:00:10 | 2 | | host2 | +| 1970-01-01T00:00:15 | 1 | | host1 | +| 1970-01-01T00:00:15 | 2 | | host2 | ++---------------------+-----+-----+-------+ + -- Empty regex and different label value -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} * label_replace(vector(1), "host", "host2", "host", ""); @@ -307,10 +376,10 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "idc99", "id +---------------------+-----+---------+-------+------+ | ts | val | new_idc | host | idc | +---------------------+-----+---------+-------+------+ -| 1970-01-01T00:00:00 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:05 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:10 | 1 | idc1 | host1 | idc1 | -| 1970-01-01T00:00:15 | 1 | idc1 | host1 | idc1 | +| 1970-01-01T00:00:00 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:05 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:10 | 1 | | host1 | idc1 | +| 1970-01-01T00:00:15 | 1 | | host1 | idc1 | +---------------------+-----+---------+-------+------+ -- SQLNESS SORT_RESULT 3 1 diff --git a/tests/cases/standalone/common/promql/label.sql b/tests/cases/standalone/common/promql/label.sql index 6a77e9b5214..6c2d64d42ac 100644 --- a/tests/cases/standalone/common/promql/label.sql +++ b/tests/cases/standalone/common/promql/label.sql @@ -17,6 +17,7 @@ INSERT INTO TABLE test VALUES (15000, 'host2', 'idc4',8); -- Missing source labels -- +-- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-"); -- dst_label is equal to source label -- @@ -27,10 +28,15 @@ TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "host"); -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", "idc", "host"); --- test the empty source label -- --- SQLNESS SORT_RESULT 3 1 +-- Absent labels join as empty strings and keep their separators -- +TQL EVAL (0, 15, '5s') label_join(label_join(vector(1), "a", "", "missing"), "b", "-", "a", "a"); + +-- An empty source label name is invalid -- TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "host", "-", ""); +-- Both hosts are joined to the same `host` at 0s, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_join(test, "host", "-", "idc"); + -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "host"); @@ -50,8 +56,13 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "new_idc", "$2", "idc", -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc", "$2", "idc", "(.*):(.*)"); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "idc", "$2", "idc", "(.*):(.*)"); + +-- Both hosts are rewritten to the same `host`, leaving two series with the same label set -- +TQL EVAL (0, 15, '5s') label_replace(test, "host", "x", "host", ".*"); + -- test the empty source label -- --- TODO(dennis): we can't remove the label currently -- -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 15, '5s') label_replace(test{host="host2"}, "idc2", "", "", ""); @@ -77,6 +88,10 @@ TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "host2", "insta -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') label_replace(test{host="host1"}, "host2", "", "instance", ""); +-- An empty replacement removes an existing label -- +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '5s') label_replace(test{idc="idc1"}, "idc", "", "instance", ""); + -- Empty regex and different label value -- SQLNESS SORT_RESULT 3 1 TQL EVAL(0, 15, '5s') {__name__="test",host="host1"} * label_replace(vector(1), "host", "host2", "host", ""); diff --git a/tests/cases/standalone/common/promql/native_time_selection.result b/tests/cases/standalone/common/promql/native_time_selection.result index 2c0e20847cb..945189e68f6 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.result +++ b/tests/cases/standalone/common/promql/native_time_selection.result @@ -613,6 +613,15 @@ TQL EVAL (1, 1, '1s', '1s') native_time_sec offset -500ms; | 11.0 | 1970-01-01T00:00:01 | +------+---------------------+ +-- `timestamp()` reports the timestamp of the selected sample (1s), without the offset applied. +TQL EVAL (1, 1, '1s', '1s') timestamp(native_time_sec offset -500ms); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:00:01 | 1.0 | ++---------------------+-------+ + DROP TABLE native_time_sec; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/native_time_selection.sql b/tests/cases/standalone/common/promql/native_time_selection.sql index ca795b0b3ad..204b047ef23 100644 --- a/tests/cases/standalone/common/promql/native_time_selection.sql +++ b/tests/cases/standalone/common/promql/native_time_selection.sql @@ -219,4 +219,6 @@ CREATE TABLE native_time_sec (ts TIMESTAMP(0) TIME INDEX, val DOUBLE); INSERT INTO native_time_sec VALUES (0, 10), (1, 11), (2, 12); TQL EVAL (1, 1, '1s', '1s') native_time_sec offset 500ms; TQL EVAL (1, 1, '1s', '1s') native_time_sec offset -500ms; +-- `timestamp()` reports the timestamp of the selected sample (1s), without the offset applied. +TQL EVAL (1, 1, '1s', '1s') timestamp(native_time_sec offset -500ms); DROP TABLE native_time_sec; diff --git a/tests/cases/standalone/common/promql/timestamp_fn.result b/tests/cases/standalone/common/promql/timestamp_fn.result index 1f5682649a4..fc4e9c46fa6 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.result +++ b/tests/cases/standalone/common/promql/timestamp_fn.result @@ -113,17 +113,17 @@ tql eval (0, 60, '30s') abs(timestamp(timestamp_test) - avg(timestamp(timestamp_ -- Test Issue 6707 tql eval timestamp(demo_memory_usage_bytes * 1); -+------+--------------------+ -| time | value * Float64(1) | -+------+--------------------+ -+------+--------------------+ ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ tql eval timestamp(-demo_memory_usage_bytes); -+------+-----------+ -| time | (- value) | -+------+-----------+ -+------+-----------+ ++------+-------+ +| time | value | ++------+-------+ ++------+-------+ tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; @@ -133,6 +133,79 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; | 1970-01-01T00:01:00 | 60.0 | +---------------------+-------+ +-- timestamp() reports the timestamp of the selected sample, not the evaluation time +tql eval (60, 60, '1s') timestamp(timestamp_test offset 30s); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:00 | 1.0 | ++---------------------+-------+ + +tql eval (60, 60, '1s') timestamp((timestamp_test offset 30s)); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:00 | 1.0 | ++---------------------+-------+ + +-- The result of a function is sampled at the evaluation time +tql eval (90, 90, '1s') timestamp(abs(timestamp_test)); + ++---------------------+-------+ +| ts | value | ++---------------------+-------+ +| 1970-01-01T00:01:30 | 90.0 | ++---------------------+-------+ + +-- A row whose fields are all NULL holds no sample and gets no timestamp; a row with any +-- non-NULL field does. +create table timestamp_null (ts timestamp time index, host string primary key, a double, b double); + +Affected Rows: 0 + +insert into timestamp_null values + (0, 'none', NULL, NULL), + (0, 'only_a', 1.0, NULL), + (0, 'only_b', NULL, 2.0); + +Affected Rows: 3 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(timestamp_null); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 0.0 | only_a | +| 1970-01-01T00:01:00 | 0.0 | only_b | ++---------------------+-------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(-timestamp_null); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 60.0 | only_a | +| 1970-01-01T00:01:00 | 60.0 | only_b | ++---------------------+-------+--------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(abs(timestamp_null)); + ++---------------------+-------+--------+ +| ts | value | host | ++---------------------+-------+--------+ +| 1970-01-01T00:01:00 | 60.0 | only_a | +| 1970-01-01T00:01:00 | 60.0 | only_b | ++---------------------+-------+--------+ + +drop table timestamp_null; + +Affected Rows: 0 + -- Test timestamp() with multiple metrics create table timestamp_test2 (ts timestamp time index, val double); diff --git a/tests/cases/standalone/common/promql/timestamp_fn.sql b/tests/cases/standalone/common/promql/timestamp_fn.sql index 123bdeb5881..cc307007e4b 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.sql +++ b/tests/cases/standalone/common/promql/timestamp_fn.sql @@ -43,6 +43,34 @@ tql eval timestamp(-demo_memory_usage_bytes); tql eval (0, 60, '30s') timestamp(timestamp_test) == 60; +-- timestamp() reports the timestamp of the selected sample, not the evaluation time +tql eval (60, 60, '1s') timestamp(timestamp_test offset 30s); + +tql eval (60, 60, '1s') timestamp((timestamp_test offset 30s)); + +-- The result of a function is sampled at the evaluation time +tql eval (90, 90, '1s') timestamp(abs(timestamp_test)); + +-- A row whose fields are all NULL holds no sample and gets no timestamp; a row with any +-- non-NULL field does. +create table timestamp_null (ts timestamp time index, host string primary key, a double, b double); + +insert into timestamp_null values + (0, 'none', NULL, NULL), + (0, 'only_a', 1.0, NULL), + (0, 'only_b', NULL, 2.0); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(timestamp_null); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(-timestamp_null); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') timestamp(abs(timestamp_null)); + +drop table timestamp_null; + -- Test timestamp() with multiple metrics create table timestamp_test2 (ts timestamp time index, val double); From 8553db3f276f0bf0fc526214b64c6382d478b079 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 30 Sep 2026 05:20:59 +0000 Subject: [PATCH 080/133] feat(mito): wait for WAL durability before publishing a manifest watermark (#9410) * feat(mito): wait for WAL durability before publishing a manifest watermark A log store that acknowledges appends before their entries are durable can hand out entry ids that a crash loses. Before a flush, a full truncate or a discard of unflushed data records an entry id in the manifest, Mito now waits on `LogStore::wait_durable` for that id. The flush waits after its SSTs are written and observes its cancellation, so a drop or truncate queued behind it is not held up by the upload. Add engine tests of the object store WAL across restarts and of the barrier, and a log store testing hook that stores an object and then reports its create as failed. Signed-off-by: jeremyhi * test(mito): synchronize the WAL barrier tests on store signals The barrier tests waited with fixed sleeps, which do not establish that a flush or truncate has reached its durability wait, and the lost backlog test could stop the store before the seal was handled. The object store WAL gets testing hooks that observe parked creates and durability waits, and one that ends the actor as a crash would; the tests wait on them, crash the store and the engine before a reopen until nothing holds the old store, and cover a truncate or discard of unflushed data that completes once its entry is durable. Signed-off-by: jeremyhi * test(mito): rename the recovery state and the crash test in the WAL recovery tests Signed-off-by: jeremyhi * test(mito): drop a comment that restates the empty recovery state Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/object_store_wal.rs | 3 +- src/log-store/src/object_store_wal/batch.rs | 2 +- src/log-store/src/object_store_wal/store.rs | 203 ++- src/mito2/src/engine.rs | 2 + .../engine/object_store_wal_recovery_test.rs | 1150 +++++++++++++++++ src/mito2/src/error.rs | 16 +- src/mito2/src/flush.rs | 18 + src/mito2/src/wal.rs | 56 +- src/mito2/src/worker/handle_flush.rs | 3 + src/mito2/src/worker/handle_manifest.rs | 39 +- 10 files changed, 1465 insertions(+), 27 deletions(-) create mode 100644 src/mito2/src/engine/object_store_wal_recovery_test.rs diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index f36108f3bd5..2ec36b30463 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -62,6 +62,5 @@ mod format; mod io; mod store; -#[allow(unused_imports)] -pub(crate) use batch::entry_id; +pub use batch::entry_id; pub use store::ObjectStoreLogStore; diff --git a/src/log-store/src/object_store_wal/batch.rs b/src/log-store/src/object_store_wal/batch.rs index 078af921def..b1873c8a219 100644 --- a/src/log-store/src/object_store_wal/batch.rs +++ b/src/log-store/src/object_store_wal/batch.rs @@ -45,7 +45,7 @@ pub(crate) const OBJECT_SEQ_LIMIT: u64 = 1 << (u64::BITS - POSITION_BITS); /// at one so that id zero, the watermark of a region without entries, is never /// assigned. A region's ids increase with the object sequence and have gaps /// wherever other regions or other positions took the sequence. -pub(crate) fn entry_id(object_seq: u64, position: u64) -> EntryId { +pub fn entry_id(object_seq: u64, position: u64) -> EntryId { debug_assert!(object_seq < OBJECT_SEQ_LIMIT && (1..POSITION_LIMIT).contains(&position)); (object_seq << POSITION_BITS) | position } diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 2a0c75f1048..5a6454123dd 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -104,6 +104,12 @@ pub struct ObjectStoreLogStore { creates_held: watch::Sender, #[cfg(any(test, feature = "testing"))] creates_fail: Arc, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: Arc, + #[cfg(any(test, feature = "testing"))] + parked_creates: watch::Receiver, + #[cfg(any(test, feature = "testing"))] + durability_waits: watch::Receiver, } type ObsoleteEntryIds = Arc>>; @@ -203,6 +209,12 @@ impl ObjectStoreLogStore { let (creates_held_tx, creates_held_rx) = watch::channel(false); #[cfg(any(test, feature = "testing"))] let creates_fail = Arc::new(AtomicBool::new(false)); + #[cfg(any(test, feature = "testing"))] + let next_create_fails_after_write = Arc::new(AtomicBool::new(false)); + #[cfg(any(test, feature = "testing"))] + let (parked_creates_tx, parked_creates_rx) = watch::channel(0); + #[cfg(any(test, feature = "testing"))] + let (durability_waits_tx, durability_waits_rx) = watch::channel(0); let actor = Actor { io: io.clone(), @@ -237,6 +249,12 @@ impl ObjectStoreLogStore { creates_held: creates_held_rx, #[cfg(any(test, feature = "testing"))] creates_fail: creates_fail.clone(), + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: next_create_fails_after_write.clone(), + #[cfg(any(test, feature = "testing"))] + parked_creates: Arc::new(parked_creates_tx), + #[cfg(any(test, feature = "testing"))] + durability_waits: durability_waits_tx, }; common_runtime::spawn_global(actor.run()); Ok(Arc::new(Self { @@ -255,6 +273,12 @@ impl ObjectStoreLogStore { creates_held: creates_held_tx, #[cfg(any(test, feature = "testing"))] creates_fail, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write, + #[cfg(any(test, feature = "testing"))] + parked_creates: parked_creates_rx, + #[cfg(any(test, feature = "testing"))] + durability_waits: durability_waits_rx, })) } @@ -313,6 +337,20 @@ impl ObjectStoreLogStore { } } +/// The transient object store error of a create that a testing hook fails. +#[cfg(any(test, feature = "testing"))] +fn injected_create_failure(io: &dyn WalObjectIo, object_seq: u64) -> Result { + let error = object_store::Error::new( + object_store::ErrorKind::Unexpected, + "injected create failure", + ) + .set_temporary(); + Err(error).context(crate::error::WalObjectStoreSnafu { + operation: "write", + path: io.object_path(object_seq), + }) +} + /// Records the obsolete watermark of `region_id`, which never moves down. fn record_obsolete(obsolete_entry_ids: &ObsoleteEntryIds, region_id: RegionId, entry_id: EntryId) { obsolete_entry_ids @@ -383,6 +421,57 @@ impl ObjectStoreLogStore { self.creates_fail.store(true, Ordering::Release); } + /// Makes the next create that writes its object report a transient object + /// store error afterwards, so the object exists although its create + /// failed. + pub fn fail_next_create_after_write(&self) { + self.next_create_fails_after_write + .store(true, Ordering::Release); + } + + /// Waits until at least `expected` creates have been parked by + /// [`hold_creates`](Self::hold_creates) since the store was built. + pub async fn wait_for_parked_creates(&self, expected: usize) -> Result<()> { + self.parked_creates + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Waits until at least `expected` calls of + /// [`wait_durable`](LogStore::wait_durable) have had to wait for an entry + /// that is not durable since the store was built. + pub async fn wait_for_durability_waits(&self, expected: usize) -> Result<()> { + self.durability_waits + .clone() + .wait_for(|count| *count >= expected) + .await + .ok() + .map(|_| ()) + .context(ObjectStoreWalStoppedSnafu) + } + + /// Ends the actor the way a crash of the process would: nothing more is + /// written, creates that have not completed are dropped, and every caller + /// still waiting for the store fails. Returns once the actor has exited. + pub async fn crash(&self) { + self.stopped.store(true, Ordering::Release); + let (response_tx, response_rx) = oneshot::channel(); + if self + .command_tx + .send(Command::Crash { + response: response_tx, + }) + .await + .is_ok() + { + let _ = response_rx.await; + } + } + /// Sets the stopped flag without sending the stop command, which is the /// state a store is in between the two steps of [`stop`](LogStore::stop). pub fn begin_stop(&self) { @@ -662,6 +751,8 @@ enum Command { Seal { response: oneshot::Sender>, }, + #[cfg(any(test, feature = "testing"))] + Crash { response: oneshot::Sender<()> }, } /// An append waiting for the object that holds its entries. @@ -797,6 +888,12 @@ struct Actor { creates_held: watch::Receiver, #[cfg(any(test, feature = "testing"))] creates_fail: Arc, + #[cfg(any(test, feature = "testing"))] + next_create_fails_after_write: Arc, + #[cfg(any(test, feature = "testing"))] + parked_creates: Arc>, + #[cfg(any(test, feature = "testing"))] + durability_waits: watch::Sender, } impl Actor { @@ -835,6 +932,12 @@ impl Actor { Some(Command::Seal { response }) => { self.handle_seal(response); } + #[cfg(any(test, feature = "testing"))] + Some(Command::Crash { response }) => { + self.handle_crash(); + let _ = response.send(()); + return; + } // Every sender is gone: the store was dropped without // `stop`. The creates in flight are dropped with the actor. None => return, @@ -1076,6 +1179,10 @@ impl Actor { let mut creates_held = self.creates_held.clone(); #[cfg(any(test, feature = "testing"))] let creates_fail = self.creates_fail.clone(); + #[cfg(any(test, feature = "testing"))] + let next_create_fails_after_write = self.next_create_fails_after_write.clone(); + #[cfg(any(test, feature = "testing"))] + let parked_creates = self.parked_creates.clone(); self.creates.push(Box::pin(async move { if !delay.is_zero() { tokio::time::sleep(delay).await; @@ -1083,21 +1190,16 @@ impl Actor { // The store was dropped while the create was parked: it never // runs. #[cfg(any(test, feature = "testing"))] + if *creates_held.borrow() { + parked_creates.send_modify(|count| *count += 1); + } + #[cfg(any(test, feature = "testing"))] if creates_held.wait_for(|held| !*held).await.is_err() { return (object_seq, Err(ObjectStoreWalStoppedSnafu.build())); } #[cfg(any(test, feature = "testing"))] if creates_fail.load(Ordering::Acquire) { - let error = object_store::Error::new( - object_store::ErrorKind::Unexpected, - "injected create failure", - ) - .set_temporary(); - let result = Err(error).context(crate::error::WalObjectStoreSnafu { - operation: "write", - path: io.object_path(object_seq), - }); - return (object_seq, result); + return (object_seq, injected_create_failure(io.as_ref(), object_seq)); } // Any conflict poisons an `enqueued` store, so only the // `durable` mode reads the epoch of the existing object. @@ -1109,6 +1211,10 @@ impl Actor { } result => result, }; + #[cfg(any(test, feature = "testing"))] + if result.is_ok() && next_create_fails_after_write.swap(false, Ordering::AcqRel) { + return (object_seq, injected_create_failure(io.as_ref(), object_seq)); + } (object_seq, result) })); } @@ -1354,6 +1460,8 @@ impl Actor { entry_id, response, }); + #[cfg(any(test, feature = "testing"))] + self.durability_waits.send_modify(|count| *count += 1); } /// Makes sure no id of the region at or below `entry_id` is assigned @@ -1489,6 +1597,21 @@ impl Actor { let _ = response.send(result); } + /// Fails every caller that waits for the store; the creates are + /// dropped with the actor. + #[cfg(any(test, feature = "testing"))] + fn handle_crash(&mut self) { + let stopped = || ObjectStoreWalStoppedSnafu.build(); + for batch in self.sealed.drain(..) { + batch.fail(stopped); + } + self.reset_open_batch(); + self.fail_unacknowledged(stopped); + for response in self.stop.drain(..) { + let _ = response.send(Err(stopped())); + } + } + fn is_stopped(&self) -> bool { self.stopped.load(Ordering::Acquire) } @@ -2759,6 +2882,9 @@ mod tests { admitted_appends: watch::channel(0).1, creates_held: watch::channel(false).0, creates_fail: Arc::default(), + next_create_fails_after_write: Arc::default(), + parked_creates: watch::channel(0).1, + durability_waits: watch::channel(0).1, }; (store, command_rx) } @@ -3956,6 +4082,63 @@ mod tests { store.stop().await.unwrap(); } + #[tokio::test] + async fn test_store_hook_fails_the_next_create_after_it_writes() { + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + + // Object 1 is stored, but its create reports a failure, so the + // append fails and the region has no durable entry. + store.fail_next_create_after_write(); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(0, latest(&store, region_id)); + + // Only that create failed. + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!( + HashMap::from([(region_id, id(2, 1))]), + response.last_entry_ids + ); + assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_hooks_observe_parked_creates_waits_and_crash() { + let store = open(memory_store(), &enqueued(eager())).await; + let region_id = region(1); + + // The append is acknowledged and its create parks; a wait for its + // entry has to wait. + store.hold_creates(); + append(&store, region_id, "a1").await.unwrap(); + timeout(WAIT, store.wait_for_parked_creates(1)) + .await + .unwrap() + .unwrap(); + let wait = { + let store = store.clone(); + tokio::spawn(async move { store.wait_durable(&provider(region_id), id(1, 1)).await }) + }; + timeout(WAIT, store.wait_for_durability_waits(1)) + .await + .unwrap() + .unwrap(); + assert!(!wait.is_finished()); + + // The crash fails the waiter and the store, and the parked create + // never writes its object. + store.crash().await; + assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); + assert_stopped(&append(&store, region_id, "a2").await.unwrap_err()); + assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); + } + #[tokio::test] async fn test_store_create_held_when_the_store_is_dropped_never_runs() { let (io, _) = RecordingIo::over(memory_store()); diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index abfb470c1be..47752adbe75 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -51,6 +51,8 @@ pub mod listener; #[cfg(test)] mod merge_mode_test; #[cfg(test)] +mod object_store_wal_recovery_test; +#[cfg(test)] mod object_store_wal_test; #[cfg(test)] mod open_test; diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs new file mode 100644 index 00000000000..6017f745d2e --- /dev/null +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -0,0 +1,1150 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Recovery tests for regions on a real [`ObjectStoreLogStore`] over an +//! in-memory object store. The store seals objects only through its testing +//! hooks, so every test controls which entries are durable. +//! +//! Every open of the store writes a start object above every present object, +//! so the first data object of a store opened on an empty prefix has +//! sequence 1. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use api::v1::Rows; +use common_base::readable_size::ReadableSize; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; +use common_recordbatch::RecordBatches; +use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig, STANDALONE_GENERATION}; +use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use log_store::ObjectStoreLogStore; +use log_store::object_store_wal::entry_id; +use object_store::ObjectStore; +use object_store::services::Memory; +use rstest::rstest; +use store_api::logstore::LogStore; +use store_api::logstore::provider::Provider; +use store_api::mito_engine_options::SKIP_WAL_KEY; +use store_api::region_engine::{RegionEngine, RegionRole}; +use store_api::region_request::{ + PathType, RegionDropRequest, RegionFlushRequest, RegionOpenRequest, RegionPutRequest, + RegionRequest, RegionTruncateRequest, +}; +use store_api::storage::{RegionId, ScanRequest}; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::region::MitoRegionRef; +use crate::test_util::{ + CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, flush_region, put_rows, + rows_schema, +}; + +/// The root the store derives its node prefix from. +const ROOT: &str = "cluster-a/wal"; +/// The node prefix of datanode 0 in the standalone generation under [`ROOT`]. +const PREFIX: &str = "cluster-a/wal/datanodes/0/epochs/0"; +/// The engine runs two workers and these regions map to different ones, so +/// their writes can be admitted into the same open batch. +const REGION_A: RegionId = RegionId::new(1, 1); +const REGION_B: RegionId = RegionId::new(1, 2); +const WAIT: Duration = Duration::from_secs(30); + +fn memory_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap() +} + +/// The configuration of a store with `ack_mode` that never seals a batch on +/// its own. +fn store_config(ack_mode: AckMode) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + storage_provider: String::new(), + prefix: ROOT.to_string(), + flush_interval: Duration::from_secs(3600), + max_batch_bytes: ReadableSize(u64::MAX), + ack_mode, + ..Default::default() + } +} + +/// Opens a store under [`PREFIX`] with `ack_mode` that never seals a batch +/// on its own. +async fn open_store(object_store: &ObjectStore, ack_mode: AckMode) -> Arc { + ObjectStoreLogStore::try_new( + object_store.clone(), + &store_config(ack_mode), + 0, + STANDALONE_GENERATION, + ) + .await + .unwrap() +} + +async fn new_engine(env: &mut TestEnv, store: Arc) -> MitoEngine { + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + env.create_engine_with_log_store(config, store).await +} + +fn wal_options() -> HashMap { + let options = WalOptions::ObjectStore(ObjectStoreWalOptions::new(PREFIX.to_string())); + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&options).unwrap(), + )]) +} + +fn provider(region_id: RegionId) -> Provider { + Provider::object_store_provider(region_id, PREFIX.to_string()) +} + +/// Creates a region on the object store WAL with `extra_options` and returns +/// its table dir and row schema. +async fn create_region( + engine: &MitoEngine, + region_id: RegionId, + extra_options: &[(&str, &str)], +) -> (String, Vec) { + let mut builder = + CreateRequestBuilder::new().insert_option(WAL_OPTIONS_KEY, &wal_options()[WAL_OPTIONS_KEY]); + for (key, value) in extra_options { + builder = builder.insert_option(key, value); + } + let request = builder.build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + (table_dir, schema) +} + +/// Opens a region whose options select the object store WAL under +/// [`PREFIX`] plus `extra_options`, and makes it writable. +async fn open_region( + engine: &MitoEngine, + region_id: RegionId, + table_dir: &str, + extra_options: &[(&str, &str)], +) -> std::result::Result<(), BoxedError> { + let mut options = wal_options(); + options.extend( + extra_options + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())), + ); + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir: table_dir.to_string(), + options, + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await?; + engine.set_region_role(region_id, RegionRole::Leader) +} + +fn rows(schema: &[api::v1::ColumnSchema], start: usize, end: usize) -> Rows { + Rows { + schema: schema.to_vec(), + rows: build_rows(start, end), + } +} + +/// Returns every row of the region as a table, so two scans can be compared +/// row by row. +async fn scan_rows(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + RecordBatches::try_collect(stream) + .await + .unwrap() + .pretty_print() + .unwrap() +} + +fn region(engine: &MitoEngine, region_id: RegionId) -> MitoRegionRef { + engine.get_region(region_id).unwrap() +} + +/// The entry ids a region tracks in memory and in its manifest, and the +/// rows its memtables hold. +#[derive(Debug, PartialEq, Eq)] +struct RegionRecoveryState { + flushed_entry_id: u64, + last_entry_id: u64, + topic_latest_entry_id: u64, + manifest_flushed_entry_id: u64, + memtable_rows: u64, +} + +async fn region_recovery_state(engine: &MitoEngine, region_id: RegionId) -> RegionRecoveryState { + let region = region(engine, region_id); + let current = region.version_control.current(); + let manifest = region.manifest_ctx.manifest().await; + RegionRecoveryState { + flushed_entry_id: current.version.flushed_entry_id, + last_entry_id: current.last_entry_id, + topic_latest_entry_id: region.topic_latest_entry_id.load(Ordering::Relaxed), + manifest_flushed_entry_id: manifest.flushed_entry_id, + memtable_rows: current.version.memtables.num_rows(), + } +} + +const EMPTY_RECOVERY_STATE: RegionRecoveryState = RegionRecoveryState { + flushed_entry_id: 0, + last_entry_id: 0, + topic_latest_entry_id: 0, + manifest_flushed_entry_id: 0, + memtable_rows: 0, +}; + +/// Writes rows through the engine and seals them into one WAL object. +/// +/// A put in the `durable` mode blocks until its entries are durable, so the +/// writes run in the background until the store has admitted all of them and +/// the batch is sealed by hand. +struct SealedWriter { + store: Arc, + admitted: usize, +} + +impl SealedWriter { + fn new(store: &Arc) -> Self { + Self { + store: store.clone(), + admitted: 0, + } + } + + async fn put_and_seal(&mut self, engine: &MitoEngine, writes: Vec<(RegionId, Rows)>) { + let handles = writes + .into_iter() + .map(|(region_id, rows)| { + let engine = engine.clone(); + tokio::spawn(async move { put_rows(&engine, region_id, rows).await }) + }) + .collect::>(); + self.admitted += handles.len(); + tokio::time::timeout(WAIT, self.store.wait_for_admitted_appends(self.admitted)) + .await + .expect("writes must reach the store") + .unwrap(); + self.store.seal_open_batch().await.unwrap(); + for handle in handles { + handle.await.unwrap(); + } + } +} + +fn latest(store: &ObjectStoreLogStore, region_id: RegionId) -> u64 { + store.latest_entry_id(&provider(region_id)).unwrap() +} + +/// Returns the sequences of the WAL objects under the prefix, in order. +async fn wal_object_seqs(object_store: &ObjectStore) -> Vec { + let mut seqs = object_store + .list(&format!("{PREFIX}/objects/")) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.metadata().is_file()) + .map(|entry| { + let path = entry.path(); + path.rsplit('/') + .next() + .and_then(|name| name.strip_suffix(".wal")) + .and_then(|seq| seq.parse().ok()) + .unwrap_or_else(|| panic!("unexpected WAL object key {path}")) + }) + .collect::>(); + seqs.sort_unstable(); + seqs +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_reopen_after_partial_flush_replays_only_unflushed_regions(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-partial-flush").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Objects 1 and 2 each hold one entry of both regions: the ids of a + // region are its position under the sequence of its object. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal( + &engine, + vec![ + (REGION_B, rows(&schema_b, 3, 5)), + (REGION_A, rows(&schema_a, 2, 4)), + ], + ) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(entry_id(2, 1), latest(&store, REGION_B)); + + // Flushing empties the memtables, so the topic latest entry id follows + // the store. + flush_region(&engine, REGION_A, None).await; + let flushed_a = RegionRecoveryState { + flushed_entry_id: entry_id(2, 1), + last_entry_id: entry_id(2, 1), + topic_latest_entry_id: entry_id(2, 1), + manifest_flushed_entry_id: entry_id(2, 1), + memtable_rows: 0, + }; + assert_eq!(flushed_a, region_recovery_state(&engine, REGION_A).await); + let rows_a = scan_rows(&engine, REGION_A).await; + let rows_b = scan_rows(&engine, REGION_B).await; + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + + // Region A replays nothing: it flushed both entries, so the topic latest + // entry id comes from the store. Region B replays both of its entries. + assert_eq!(flushed_a, region_recovery_state(&engine, REGION_A).await); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 5, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_B).await + ); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(entry_id(2, 1), latest(&store, REGION_B)); + assert_eq!(rows_a, scan_rows(&engine, REGION_A).await); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(4, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_reopen_after_crash_replays_durable_entries_once(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-crash").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is durable as an object but the manifest never learns of it. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 3))]) + .await; + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + + drop(writer); + crash(engine, store).await; + + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(1, 1), + memtable_rows: 3, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(3, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_open_region_rejects_mismatched_wal_prefix() { + let mut env = TestEnv::with_prefix("object-store-wal-prefix-mismatch").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 3))]) + .await; + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The process now runs its store under the next generation while the + // region still persists the prefix it was created with. + let config = store_config(AckMode::Durable); + let other_prefix = config.node_prefix(0, STANDALONE_GENERATION + 1); + let store = + ObjectStoreLogStore::try_new(object_store.clone(), &config, 0, STANDALONE_GENERATION + 1) + .await + .unwrap(); + let engine = new_engine(&mut env, store).await; + let err = open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + let message = err.output_msg(); + assert!( + message.contains(PREFIX) && message.contains(&other_prefix), + "unexpected error: {message}" + ); + assert!(!engine.is_region_exists(REGION_A)); +} + +#[tokio::test] +async fn test_reopen_on_empty_prefix_without_durable_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-empty-prefix").await; + let store = open_store(&memory_store(), AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + // Region A never writes; region B skips the WAL and flushes its rows. + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[(SKIP_WAL_KEY, "true")]).await; + put_rows(&engine, REGION_B, rows(&schema_b, 0, 3)).await; + flush_region(&engine, REGION_B, None).await; + let rows_b = scan_rows(&engine, REGION_B).await; + + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // A fresh object store holds no object under the prefix. + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[(SKIP_WAL_KEY, "true")]) + .await + .unwrap(); + + assert_eq!(0, latest(&store, REGION_A)); + assert_eq!(0, latest(&store, REGION_B)); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + // Region B keeps the object store provider but wrote nothing to it; its + // rows come back from the SST alone. + assert_eq!(provider(REGION_B), region(&engine, REGION_B).provider); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_B).await.memtable_rows + ); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + + // The region is usable and its first entry lands in the first object + // after the start object. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 0, 2))]) + .await; + assert_eq!(vec![0, 1], wal_object_seqs(&object_store).await); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_entry_ids_continue_across_two_restarts(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-two-restarts").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry of object 1 is flushed, the entry of object 2 is not. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 0, 2))]) + .await; + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 2, 4))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The first restart writes start object 3 and replays the entry of + // object 2 only. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(2, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 2, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(4, engine.get_region_statistic(REGION_A).unwrap().num_rows); + + // The sequence continues at object 4, whose entry is flushed; the entry + // of object 5 is not. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 4, 6))]) + .await; + assert_eq!(entry_id(4, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(4, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 6, 8))]) + .await; + assert_eq!(entry_id(5, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // The second restart replays the entry of object 5 only. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(4, 1), + last_entry_id: entry_id(5, 1), + topic_latest_entry_id: entry_id(4, 1), + manifest_flushed_entry_id: entry_id(4, 1), + memtable_rows: 2, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(5, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(8, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!( + 8, + region(&engine, REGION_A) + .version_control + .committed_sequence() + ); +} + +/// Sends a put without asserting its result. +async fn try_put( + engine: &MitoEngine, + region_id: RegionId, + rows: Rows, +) -> std::result::Result<(), BoxedError> { + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows, + hint: None, + partition_expr_version: None, + }), + ) + .await + .map(|_| ()) +} + +#[tokio::test] +async fn test_reopen_after_a_failed_create_shows_only_the_acknowledged_write() { + let mut env = TestEnv::with_prefix("object-store-wal-failed-create").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let key_rows = |value| Rows { + schema: schema.clone(), + rows: build_rows_for_key("k", 0, 1, value), + }; + + // Object 1 is stored, but its create is reported as failed, so the write + // fails. + let failed = { + let engine = engine.clone(); + let rows = key_rows(1); + tokio::spawn(async move { try_put(&engine, REGION_A, rows).await }) + }; + tokio::time::timeout(WAIT, store.wait_for_admitted_appends(1)) + .await + .unwrap() + .unwrap(); + store.fail_next_create_after_write(); + store.seal_open_batch().await.unwrap_err(); + assert!(failed.await.unwrap().is_err()); + assert_eq!(vec![0, 1], wal_object_seqs(&object_store).await); + assert_eq!(0, latest(&store, REGION_A)); + + // The same key and timestamp is written again under the same row + // sequence, and acknowledged from object 2. + let mut writer = SealedWriter { + store: store.clone(), + admitted: 1, + }; + writer + .put_and_seal(&engine, vec![(REGION_A, key_rows(2))]) + .await; + assert_eq!(vec![0, 1, 2], wal_object_seqs(&object_store).await); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + let region_a = region(&engine, REGION_A); + assert_eq!(1, region_a.version_control.committed_sequence()); + let expected = "\ ++-------+---------+---------------------+ +| tag_0 | field_0 | ts | ++-------+---------+---------------------+ +| k | 2.0 | 1970-01-01T00:00:00 | ++-------+---------+---------------------+"; + assert_eq!(expected, scan_rows(&engine, REGION_A).await); + + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // Object 1 is off the chain, so only the acknowledged write replays. + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 1, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(expected, scan_rows(&engine, REGION_A).await); +} + +/// Requests a flush of the region in the background and returns its handle; +/// the result is up to the caller. +fn spawn_flush( + engine: &MitoEngine, + region_id: RegionId, +) -> tokio::task::JoinHandle> { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest { + row_group_size: None, + reason: None, + }), + ) + .await + .map(|_| ()) + }) +} + +/// Seals the open batch in the background; with creates held, the seal +/// returns only once they are released. +fn spawn_seal( + store: &Arc, +) -> tokio::task::JoinHandle> { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) +} + +/// Waits until a testing hook of the store observes the state it waits for. +async fn wait_for(hook: impl Future>) { + tokio::time::timeout(WAIT, hook) + .await + .expect("the store must reach the state") + .unwrap(); +} + +/// Tears the engine and its store down as a crash of the process would: the +/// store writes nothing more and fails every caller still waiting for it, +/// then the engine stops. Returns once nothing holds the store any more, so a +/// store reopened on the prefix never runs beside the old one. +async fn crash(engine: MitoEngine, store: Arc) { + store.crash().await; + engine.stop().await.unwrap(); + drop(engine); + let released = Arc::downgrade(&store); + drop(store); + tokio::time::timeout(WAIT, async { + while released.strong_count() > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("nothing may hold the crashed store"); +} + +#[tokio::test] +async fn test_flush_waits_until_the_wal_is_durable() { + let mut env = TestEnv::with_prefix("object-store-wal-flush-barrier").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The write is acknowledged on admission; its object is not created + // while creates are held. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(1, 1), + memtable_rows: 3, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, latest(&store, REGION_A)); + + // The flush writes its SST, then waits for the entry to be durable + // before it records the entry as flushed in the manifest. + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + assert!(!seal.is_finished()); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(1, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 0, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(3, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_enqueued_crash_before_the_object_exists_replays_durable_entries() { + let mut env = TestEnv::with_prefix("object-store-wal-enqueued-crash").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged and a flush is attempted, but its object is + // never created: the flush waits and the manifest keeps watermark 0. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // The process dies with the create still held: the waiting flush and + // the seal fail, and the object is never written. + store.crash().await; + assert!(flush.await.unwrap().is_err()); + assert!(seal.await.unwrap().is_err()); + crash(engine, store).await; + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // Nothing is durable, so nothing replays. The restart writes start + // object 1 and the region's next entry lands in object 2; the lost entry + // was inside the unpersisted backlog. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, latest(&store, REGION_A)); + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 3, 5))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!( + entry_id(2, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + assert_eq!(vec![0, 1, 2], wal_object_seqs(&object_store).await); + let rows_before = scan_rows(&engine, REGION_A).await; + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); + drop(writer); + crash(engine, store).await; + + // The second restart replays the entry that became durable and skips + // nothing: the manifest never named an entry that was not durable. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_drop_cancels_a_flush_waiting_for_wal_durability() { + let mut env = TestEnv::with_prefix("object-store-wal-flush-barrier-drop").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The flush reaches the durability barrier and waits for a create that + // is held. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let flush = spawn_flush(&engine, REGION_A); + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!flush.is_finished()); + + // The drop cancels the flush instead of waiting behind the upload. + tokio::time::timeout( + WAIT, + engine.handle_request( + REGION_A, + RegionRequest::Drop(RegionDropRequest { + fast_path: false, + force: false, + partial_drop: false, + }), + ), + ) + .await + .unwrap() + .unwrap(); + assert!(!engine.is_region_exists(REGION_A)); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + assert!( + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .is_err() + ); + + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); +} + +#[rstest] +#[case(RegionTruncateRequest::All)] +#[case(RegionTruncateRequest::Unflushed)] +#[tokio::test] +async fn test_enqueued_truncate_waits_until_the_wal_is_durable( + #[case] request: RegionTruncateRequest, +) { + let mut env = TestEnv::with_prefix("object-store-wal-truncate-barrier").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged and a truncate is attempted, but its object + // is never created: the truncate waits and the manifest keeps frontier 0. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + assert_eq!( + entry_id(1, 1), + region_recovery_state(&engine, REGION_A).await.last_entry_id + ); + let truncate = { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request(REGION_A, RegionRequest::Truncate(request)) + .await + .map(|_| ()) + }) + }; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!truncate.is_finished()); + let manifest = region(&engine, REGION_A).manifest_ctx.manifest().await; + assert_eq!(0, manifest.flushed_entry_id); + assert_eq!(None, manifest.truncated_entry_id); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // The process dies with the create still held: the waiting truncate and + // the seal fail. + store.crash().await; + assert!(truncate.await.unwrap().is_err()); + assert!(seal.await.unwrap().is_err()); + crash(engine, store).await; + + // Nothing is durable and the manifest names no entry, so the region + // starts over after start object 1, and its next entry becomes durable + // in object 2. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + EMPTY_RECOVERY_STATE, + region_recovery_state(&engine, REGION_A).await + ); + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema, 3, 5))]) + .await; + assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); + let rows_before = scan_rows(&engine, REGION_A).await; + drop(writer); + crash(engine, store).await; + + // The second restart replays the durable entry instead of skipping it. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(rows_before, scan_rows(&engine, REGION_A).await); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[rstest] +#[case(RegionTruncateRequest::All)] +#[case(RegionTruncateRequest::Unflushed)] +#[tokio::test] +async fn test_enqueued_truncate_completes_once_the_wal_is_durable( + #[case] request: RegionTruncateRequest, +) { + let mut env = TestEnv::with_prefix("object-store-wal-truncate-durable").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir, schema) = create_region(&engine, REGION_A, &[]).await; + let truncated_entry_id = + matches!(request, RegionTruncateRequest::All).then_some(entry_id(1, 1)); + + // The truncate waits for the acknowledged entry, then records it once + // its object is created. + store.hold_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let truncate = { + let engine = engine.clone(); + tokio::spawn(async move { + engine + .handle_request(REGION_A, RegionRequest::Truncate(request)) + .await + .map(|_| ()) + }) + }; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + wait_for(store.wait_for_durability_waits(1)).await; + assert!(!truncate.is_finished()); + store.release_creates(); + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(WAIT, truncate) + .await + .unwrap() + .unwrap() + .unwrap(); + let manifest = region(&engine, REGION_A).manifest_ctx.manifest().await; + assert_eq!(entry_id(1, 1), manifest.flushed_entry_id); + assert_eq!(truncated_entry_id, manifest.truncated_entry_id); + assert_eq!(0, engine.get_region_statistic(REGION_A).unwrap().num_rows); + engine.stop().await.unwrap(); + drop(engine); + drop(store); + + // The restart replays nothing: the discarded rows stay discarded. + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir, &[]) + .await + .unwrap(); + assert_eq!( + RegionRecoveryState { + flushed_entry_id: entry_id(1, 1), + last_entry_id: entry_id(1, 1), + topic_latest_entry_id: entry_id(1, 1), + manifest_flushed_entry_id: entry_id(1, 1), + memtable_rows: 0, + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(0, engine.get_region_statistic(REGION_A).unwrap().num_rows); +} + +#[tokio::test] +async fn test_flush_does_not_publish_a_frontier_for_a_lost_enqueued_backlog() { + let mut env = TestEnv::with_prefix("object-store-wal-lost-backlog").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Enqueued).await; + let engine = new_engine(&mut env, store.clone()).await; + let (_, schema) = create_region(&engine, REGION_A, &[]).await; + + // The entry is acknowledged, then its create fails while stop has begun + // but the stop command is not handled yet: the backlog is lost. + store.hold_creates(); + store.fail_creates(); + put_rows(&engine, REGION_A, rows(&schema, 0, 3)).await; + let seal = spawn_seal(&store); + wait_for(store.wait_for_parked_creates(1)).await; + store.begin_stop(); + store.release_creates(); + assert!( + tokio::time::timeout(WAIT, seal) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert_eq!(vec![0], wal_object_seqs(&object_store).await); + + // A flush in that window fails at the barrier instead of recording the + // lost entry as flushed. + let flush = spawn_flush(&engine, REGION_A); + assert!( + tokio::time::timeout(WAIT, flush) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert_eq!( + 0, + region_recovery_state(&engine, REGION_A) + .await + .manifest_flushed_entry_id + ); + assert!(store.stop().await.is_err()); +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 7d43548dd2b..f738044c9ef 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -453,6 +453,14 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Failed to wait for the WAL to be durable, region_id: {}", region_id))] + WaitWalDurable { + region_id: RegionId, + #[snafu(implicit)] + location: Location, + source: BoxedError, + }, + // Shared error for each writer in the write group. #[snafu(display("Failed to write region"))] WriteGroup { source: Arc }, @@ -1439,9 +1447,10 @@ impl ErrorExt for Error { OpenDal { .. } | ManifestDeltaNotFound { .. } | ReadParquet { .. } => { StatusCode::StorageUnavailable } - WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } => { - source.status_code() - } + WriteWal { source, .. } + | ReadWal { source, .. } + | DeleteWal { source, .. } + | WaitWalDurable { source, .. } => source.status_code(), CompressObject { .. } | DecompressObject { .. } | SerdeJson { .. } @@ -1651,6 +1660,7 @@ impl ErrorExt for Error { WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } + | WaitWalDurable { source, .. } | FetchManifests { source, .. } | External { source, .. } => source.retry_hint(), diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index a5d9a7d0a73..b5d8c4e56ea 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -69,6 +69,7 @@ use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, }; use crate::sst::{FlatSchemaOptions, FormatType, to_flat_sst_arrow_schema}; +use crate::wal::DurabilityBarrier; use crate::worker::WorkerListener; /// Global write buffer (memtable) manager. @@ -282,6 +283,8 @@ pub(crate) struct RegionFlushTask { /// /// This is used to generate the file meta. pub(crate) partition_expr: Option, + /// Waits until the WAL is durable through the entry id the flush records. + pub(crate) durability_barrier: DurabilityBarrier, } struct FlushTaskWaiters { @@ -509,6 +512,17 @@ impl RegionFlushTask { .await; } + // The manifest may name only a durable entry as flushed. A log store + // that acknowledges appends before their entries are durable can have + // handed out `last_entry_id` for an entry that is still in its backlog. + // A DDL that cancels the flush must not wait behind that upload. + let durable = self.durability_barrier.wait(version_data.last_entry_id); + tokio::pin!(durable); + match CancellableFuture::new(durable.as_mut(), state.cancel_handle()).await { + Ok(result) => result?, + Err(_) => return FlushCancelledSnafu.fail(), + } + let edit = RegionEdit { files_to_add: file_metas, files_to_remove: Vec::new(), @@ -1703,6 +1717,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), } } @@ -1853,6 +1868,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }; task.push_sender(OptionOutputTx::from(output_tx)); scheduler @@ -2143,6 +2159,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }) .collect(); // Schedule first task. @@ -2447,6 +2464,7 @@ mod tests { flush_semaphore: Arc::new(Semaphore::new(2)), is_staging: false, partition_expr: None, + durability_barrier: DurabilityBarrier::noop(), }) .collect(); // Schedule first task. diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index fa549622898..772aac336cf 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -36,7 +36,7 @@ use store_api::logstore::provider::Provider; use store_api::logstore::{AppendBatchResponse, LogStore, WalIndex}; use store_api::storage::RegionId; -use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WriteWalSnafu}; +use crate::error::{BuildEntrySnafu, DeleteWalSnafu, Result, WaitWalDurableSnafu, WriteWalSnafu}; use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader}; use crate::wal::raw_entry_reader::{LogStoreRawEntryReader, RegionRawEntryReader}; @@ -45,6 +45,36 @@ pub type EntryId = store_api::logstore::entry::Id; /// A stream that yields tuple of WAL entry id and corresponding entry. pub type WalEntryStream<'a> = BoxStream<'a, Result<(EntryId, WalEntry)>>; +/// Waits until the WAL of one region is durable through an entry id. +/// +/// A flush calls it before it records an entry id as flushed in the manifest: +/// a log store that acknowledges an append before its entries are durable +/// hands out entry ids that a crash can lose, and a manifest watermark that +/// names such an id would skip the entries assigned again after a restart. +#[derive(Clone)] +pub(crate) struct DurabilityBarrier( + Arc BoxFuture<'static, Result<()>> + Send + Sync>, +); + +impl DurabilityBarrier { + /// Returns once the region is durable through `entry_id`. + pub(crate) async fn wait(&self, entry_id: EntryId) -> Result<()> { + (self.0)(entry_id).await + } + + /// A barrier that never waits, for tests without a log store. + #[cfg(test)] + pub(crate) fn noop() -> Self { + Self(Arc::new(|_| Box::pin(async { Ok(()) }))) + } +} + +impl std::fmt::Debug for DurabilityBarrier { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("DurabilityBarrier") + } +} + /// Write ahead log. /// /// All regions in the engine shares the same WAL instance. @@ -104,6 +134,30 @@ impl Wal { } } + /// Returns the [DurabilityBarrier] of the region written through `provider`. + pub(crate) fn durability_barrier( + &self, + region_id: RegionId, + provider: &Provider, + ) -> DurabilityBarrier { + let store = self.store.clone(); + let provider = provider.clone(); + DurabilityBarrier(Arc::new(move |entry_id| { + let store = store.clone(); + let provider = provider.clone(); + Box::pin(async move { + if let Provider::Noop = provider { + return Ok(()); + } + store + .wait_durable(&provider, entry_id) + .await + .map_err(BoxedError::new) + .context(WaitWalDurableSnafu { region_id }) + }) + })) + } + /// Returns a [WalEntryReader] pub(crate) fn wal_entry_reader( &self, diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs index bfdfea5d1d9..1d4442d8d03 100644 --- a/src/mito2/src/worker/handle_flush.rs +++ b/src/mito2/src/worker/handle_flush.rs @@ -243,6 +243,9 @@ impl RegionWorkerLoop { flush_semaphore: self.flush_semaphore.clone(), is_staging: region.is_staging(), partition_expr: region.maybe_staging_partition_expr_str(), + durability_barrier: self + .wal + .durability_barrier(region.region_id, ®ion.provider), } } } diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index aa8475a49f1..b574e490ac7 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -32,7 +32,7 @@ use crate::cache::file_cache::{FileType, IndexKey}; use crate::config::IndexBuildMode; use crate::error::{EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, Result}; use crate::manifest::action::{ - RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, + RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, TruncateKind, }; use crate::memtable::MemtableBuilderProvider; use crate::metrics::WRITE_CACHE_INFLIGHT_DOWNLOAD; @@ -484,17 +484,31 @@ impl RegionWorkerLoop { let request_sender = self.sender.clone(); let manifest_ctx = region.manifest_ctx.clone(); let is_staging = region.is_staging(); + let durability_barrier = self + .wal + .durability_barrier(region.region_id, ®ion.provider); // Updates manifest in background. common_runtime::spawn_global(async move { + // The truncated entry id becomes the replay frontier, so the WAL + // must be durable through it before the manifest names it. + let durable = match &truncate.kind { + TruncateKind::All { + truncated_entry_id, .. + } => durability_barrier.wait(*truncated_entry_id).await, + _ => Ok(()), + }; // Write region truncated to manifest. let action_list = RegionMetaActionList::with_action(RegionMetaAction::Truncate(truncate.clone())); - let result = manifest_ctx - .update_manifest(RegionLeaderState::Truncating, action_list, is_staging) - .await - .map(|_| ()); + let result = match durable { + Ok(()) => manifest_ctx + .update_manifest(RegionLeaderState::Truncating, action_list, is_staging) + .await + .map(|_| ()), + Err(e) => Err(e), + }; // Sends the result back to the request sender. let truncate_result = TruncateResult { @@ -531,10 +545,12 @@ impl RegionWorkerLoop { let region_id = region.region_id; let request_sender = self.sender.clone(); let manifest_ctx = region.manifest_ctx.clone(); + let durability_barrier = self.wal.durability_barrier(region_id, ®ion.provider); common_runtime::spawn_global(async move { // The frontier moves to the last written entry and sequence, so replaying the - // WAL after a restart skips everything the memtables held. + // WAL after a restart skips everything the memtables held. The WAL must be + // durable through that entry before the manifest names it. let edit = RegionEdit { files_to_add: Vec::new(), files_to_remove: Vec::new(), @@ -545,10 +561,13 @@ impl RegionWorkerLoop { committed_sequence: None, }; let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)); - let result = manifest_ctx - .update_manifest(RegionLeaderState::Truncating, action_list, false) - .await - .map(|_| ()); + let result = match durability_barrier.wait(discarded_entry_id).await { + Ok(()) => manifest_ctx + .update_manifest(RegionLeaderState::Truncating, action_list, false) + .await + .map(|_| ()), + Err(e) => Err(e), + }; let result = DiscardUnflushedResult { region_id, From 78eefbe4c403807444dff5abe96404e40ee72b9a Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 30 Sep 2026 06:23:33 +0000 Subject: [PATCH 081/133] perf: batch schema export requests (#9408) * perf: batch schema export requests Signed-off-by: jeremyhi * fix: propagate legacy schema export failures Signed-off-by: jeremyhi * fix: keep credentials out of SQL response errors Signed-off-by: jeremyhi * test: use a typo-safe dotted catalog name Signed-off-by: jeremyhi * refactor(cli): address schema export review feedback Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/cli/src/data.rs | 1 + src/cli/src/data/export.rs | 78 ++- src/cli/src/data/export_v2/command.rs | 159 ++---- src/cli/src/data/schema_export.rs | 508 ++++++++++++++++++ src/cli/src/database.rs | 91 +++- .../tests/export_logical_tables.rs | 46 ++ 6 files changed, 689 insertions(+), 194 deletions(-) create mode 100644 src/cli/src/data/schema_export.rs diff --git a/src/cli/src/data.rs b/src/cli/src/data.rs index c880bd81c73..3e1924f12ee 100644 --- a/src/cli/src/data.rs +++ b/src/cli/src/data.rs @@ -18,6 +18,7 @@ mod import; pub mod import_v2; pub(crate) mod path; pub(crate) mod progress; +mod schema_export; pub mod snapshot_storage; pub(crate) mod sql; mod storage_export; diff --git a/src/cli/src/data/export.rs b/src/cli/src/data/export.rs index 148c27316e7..055c5d4ee02 100644 --- a/src/cli/src/data/export.rs +++ b/src/cli/src/data/export.rs @@ -27,6 +27,7 @@ use tokio::sync::Semaphore; use tokio::time::Instant; use crate::common::{ObjectStoreConfig, new_fs_object_store}; +use crate::data::schema_export::append_schema_ddl; use crate::data::storage_export::{ AzblobBackend, FsBackend, GcsBackend, OssBackend, S3Backend, StorageExport, StorageType, }; @@ -335,32 +336,6 @@ impl Export { )) } - async fn show_create( - &self, - show_type: &str, - catalog: &str, - schema: &str, - table: Option<&str>, - ) -> Result { - let sql = match table { - Some(table) => format!( - r#"SHOW CREATE {} "{}"."{}"."{}""#, - show_type, catalog, schema, table - ), - None => format!(r#"SHOW CREATE {} "{}"."{}""#, show_type, catalog, schema), - }; - let records = self - .database_client - .sql_in_public(&sql) - .await? - .context(EmptyResultSnafu)?; - let Value::String(create) = &records[0][1] else { - unreachable!() - }; - - Ok(format!("{};\n", create)) - } - async fn export_create_database(&self) -> Result<()> { let timer = Instant::now(); let db_names = self.get_db_names().await?; @@ -368,9 +343,15 @@ impl Export { let operator = self.build_prefer_fs_operator().await?; for schema in db_names { - let create_database = self - .show_create("DATABASE", &self.catalog, &schema, None) - .await?; + let mut create_database = String::new(); + append_schema_ddl( + &self.database_client, + &self.catalog, + &schema, + std::iter::once(("DATABASE", None)), + &mut create_database, + ) + .await?; let file_path = self.get_file_path(&schema, "create_database.sql"); self.write_to_storage(&operator, &file_path, create_database.into_bytes()) @@ -415,23 +396,28 @@ impl Export { } let file_path = export_self.get_file_path(&schema, "create_tables.sql"); - let mut content = Vec::new(); - - // Add table creation SQL - for (c, s, t) in metric_physical_tables.iter().chain(&remaining_tables) { - let create_table = export_self.show_create("TABLE", c, s, Some(t)).await?; - content.extend_from_slice(create_table.as_bytes()); - } - - // Add view creation SQL - for (c, s, v) in &views { - let create_view = export_self.show_create("VIEW", c, s, Some(v)).await?; - content.extend_from_slice(create_view.as_bytes()); - } + let mut content = String::new(); + let objects = metric_physical_tables + .iter() + .chain(&remaining_tables) + .map(|(_, _, table)| ("TABLE", Some(table.as_str()))) + .chain( + views + .iter() + .map(|(_, _, view)| ("VIEW", Some(view.as_str()))), + ); + append_schema_ddl( + &export_self.database_client, + &export_self.catalog, + &schema, + objects, + &mut content, + ) + .await?; // Write to storage export_self - .write_to_storage(&operator, &file_path, content) + .write_to_storage(&operator, &file_path, content.into_bytes()) .await?; info!( @@ -445,9 +431,11 @@ impl Export { }); } - let success = self.execute_tasks(tasks).await; + for result in futures::future::join_all(tasks).await { + result?; + } let elapsed = timer.elapsed(); - info!("Success {success}/{db_count} jobs, cost: {elapsed:?}"); + info!("Success {db_count}/{db_count} jobs, cost: {elapsed:?}"); Ok(()) } diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index 20e69ef5a99..48779f3c435 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -25,15 +25,15 @@ use common_error::ext::BoxedError; use common_telemetry::info; use serde_json::Value; use servers::http::{ColumnSchema, GreptimeQueryOutput, OutputSchema}; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use crate::Tool; use crate::common::ObjectStoreConfig; use crate::data::export_v2::coordinator::{ExportDataOptions, export_data}; use crate::data::export_v2::error::{ - ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, EmptyResultSnafu, IoSnafu, - ManifestVersionMismatchSnafu, Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, - SchemaOnlyModeMismatchSnafu, SnapshotVerifyFailedSnafu, UnexpectedValueTypeSnafu, + ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, IoSnafu, ManifestVersionMismatchSnafu, + Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, SchemaOnlyModeMismatchSnafu, + SnapshotVerifyFailedSnafu, UnexpectedValueTypeSnafu, }; use crate::data::export_v2::extractor::SchemaExtractor; use crate::data::export_v2::manifest::{ @@ -42,10 +42,11 @@ use crate::data::export_v2::manifest::{ use crate::data::export_v2::schema::{DDL_DIR, SCHEMA_DIR, SCHEMAS_FILE}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; use crate::data::progress::{ProgressMode, build_progress_reporter}; +use crate::data::schema_export::append_schema_ddl; use crate::data::snapshot_storage::{ OpenDalStorage, SnapshotStorage, validate_snapshot_uri, validate_uri, }; -use crate::data::sql::{escape_sql_identifier, escape_sql_literal}; +use crate::data::sql::escape_sql_literal; use crate::database::{DatabaseClient, parse_proxy_opts}; /// Export V2 commands. @@ -620,8 +621,10 @@ impl ExportCreate { info!("Exported {} schemas", schema_snapshot.schemas.len()); // 5. Export DDL files for import recovery. - let ddl_by_schema = self.build_ddl_by_schema(&schema_names).await?; - for (schema, ddl) in ddl_by_schema { + let mut sorted_schemas = schema_names; + sorted_schemas.sort(); + for schema in sorted_schemas { + let ddl = self.build_schema_ddl(&schema).await?; let ddl_path = ddl_path_for_schema(&schema); self.storage.write_text(&ddl_path, &ddl).await?; info!("Exported DDL for schema {} to {}", schema, ddl_path); @@ -659,45 +662,31 @@ impl ExportCreate { Ok(()) } - async fn build_ddl_by_schema(&self, schema_names: &[String]) -> Result> { - let mut schemas = schema_names.to_vec(); - schemas.sort(); - - let mut ddl_by_schema = Vec::with_capacity(schemas.len()); - for schema in schemas { - let create_database = self.show_create("DATABASE", &schema, None).await?; - - let (mut physical_tables, mut tables, mut views) = - self.get_schema_objects(&schema).await?; - physical_tables.sort(); - let mut physical_ddls = Vec::with_capacity(physical_tables.len()); - for table in physical_tables { - physical_ddls.push(self.show_create("TABLE", &schema, Some(&table)).await?); - } - - tables.sort(); - let mut table_ddls = Vec::with_capacity(tables.len()); - for table in tables { - table_ddls.push(self.show_create("TABLE", &schema, Some(&table)).await?); - } - - views.sort(); - let mut view_ddls = Vec::with_capacity(views.len()); - for view in views { - view_ddls.push(self.show_create("VIEW", &schema, Some(&view)).await?); - } - - let ddl = build_schema_ddl( - &schema, - create_database, - physical_ddls, - table_ddls, - view_ddls, - ); - ddl_by_schema.push((schema, ddl)); - } - - Ok(ddl_by_schema) + async fn build_schema_ddl(&self, schema: &str) -> Result { + let (mut physical_tables, mut tables, mut views) = self.get_schema_objects(schema).await?; + physical_tables.sort(); + tables.sort(); + views.sort(); + let objects = std::iter::once(("DATABASE", None)) + .chain( + physical_tables + .iter() + .chain(&tables) + .map(|table| ("TABLE", Some(table.as_str()))), + ) + .chain(views.iter().map(|view| ("VIEW", Some(view.as_str())))); + let mut ddl = format!("-- Schema: {schema}\n"); + append_schema_ddl( + &self.database_client, + &self.config.catalog, + schema, + objects, + &mut ddl, + ) + .await + .context(DatabaseSnafu)?; + ddl.push('\n'); + Ok(ddl) } async fn get_schema_objects( @@ -770,65 +759,6 @@ impl ExportCreate { Ok(tables.into_iter().collect()) } - - async fn show_create( - &self, - show_type: &str, - schema: &str, - table: Option<&str>, - ) -> Result { - let sql = match table { - Some(table) => format!( - r#"SHOW CREATE {} "{}"."{}"."{}""#, - show_type, - escape_sql_identifier(&self.config.catalog), - escape_sql_identifier(schema), - escape_sql_identifier(table) - ), - None => format!( - r#"SHOW CREATE {} "{}"."{}""#, - show_type, - escape_sql_identifier(&self.config.catalog), - escape_sql_identifier(schema) - ), - }; - - let records: Option>> = self - .database_client - .sql_in_public(&sql) - .await - .context(DatabaseSnafu)?; - let rows = records.context(EmptyResultSnafu)?; - let row = rows.first().context(EmptyResultSnafu)?; - let Some(Value::String(create)) = row.get(1) else { - return UnexpectedValueTypeSnafu.fail(); - }; - - Ok(format!("{};\n", create)) - } -} - -fn build_schema_ddl( - schema: &str, - create_database: String, - physical_tables: Vec, - tables: Vec, - views: Vec, -) -> String { - let mut ddl = String::new(); - ddl.push_str(&format!("-- Schema: {}\n", schema)); - ddl.push_str(&create_database); - for stmt in physical_tables { - ddl.push_str(&stmt); - } - for stmt in tables { - ddl.push_str(&stmt); - } - for stmt in views { - ddl.push_str(&stmt); - } - ddl.push('\n'); - ddl } fn validate_resume_config(manifest: &Manifest, config: &ExportConfig) -> Result<()> { @@ -1637,25 +1567,6 @@ mod tests { ); } - #[test] - fn test_build_schema_ddl_order() { - let ddl = build_schema_ddl( - "public", - "CREATE DATABASE public;\n".to_string(), - vec!["PHYSICAL;\n".to_string()], - vec!["TABLE;\n".to_string()], - vec!["VIEW;\n".to_string()], - ); - - let db_pos = ddl.find("CREATE DATABASE").unwrap(); - let physical_pos = ddl.find("PHYSICAL;").unwrap(); - let table_pos = ddl.find("TABLE;").unwrap(); - let view_pos = ddl.find("VIEW;").unwrap(); - assert!(db_pos < physical_pos); - assert!(physical_pos < table_pos); - assert!(table_pos < view_pos); - } - #[tokio::test] async fn test_build_rejects_chunk_window_without_bounds() { let cmd = ExportCreateCommand::parse_from([ diff --git a/src/cli/src/data/schema_export.rs b/src/cli/src/data/schema_export.rs new file mode 100644 index 00000000000..8bb99121d12 --- /dev/null +++ b/src/cli/src/data/schema_export.rs @@ -0,0 +1,508 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Bounded SHOW CREATE requests shared by both export commands. + +use common_catalog::consts::DEFAULT_SCHEMA_NAME; +use servers::http::GreptimeQueryOutput; +use servers::http::result::greptime_result_v1::GreptimedbV1Response; + +use crate::data::sql::escape_sql_identifier; +use crate::database::DatabaseClient; +use crate::error::{InvalidArgumentsSnafu, Result, UnexpectedSnafu}; + +const MAX_STATEMENTS: usize = 128; +const MAX_SQL_BYTES: usize = 128 * 1024; + +/// Appends DDL in dependency order supplied by the caller using bounded requests. +pub(crate) fn append_schema_ddl<'a>( + client: &'a DatabaseClient, + catalog: &'a str, + schema: &'a str, + objects: impl Iterator)> + Send + 'a, + ddl: &'a mut String, +) -> impl std::future::Future> + Send + 'a { + let mut statements = objects + .map(move |(kind, table)| show_create_sql(kind, catalog, schema, table)) + .peekable(); + async move { + loop { + let batch = next_batch(&mut statements)?; + if batch.is_empty() { + return Ok(()); + } + let response = client + .sql_response(&batch.concat(), DEFAULT_SCHEMA_NAME) + .await + .map_err(|error| { + UnexpectedSnafu { + msg: format!("SHOW CREATE batch failed for {}: {error}", batch.concat()), + } + .build() + })?; + append_results(&batch, &response, ddl)?; + } + } +} + +fn show_create_sql(kind: &str, catalog: &str, schema: &str, table: Option<&str>) -> String { + let mut sql = format!( + "SHOW CREATE {kind} \"{}\".\"{}\"", + escape_sql_identifier(catalog), + escape_sql_identifier(schema), + ); + if let Some(table) = table { + sql.push_str(".\""); + sql.push_str(&escape_sql_identifier(table)); + sql.push('"'); + } + sql.push_str(";\n"); + sql +} + +fn next_batch( + statements: &mut std::iter::Peekable>, +) -> Result> { + let mut batch = Vec::new(); + let mut bytes = 0; + while let Some(sql) = statements.peek() { + if sql.len() > MAX_SQL_BYTES { + return InvalidArgumentsSnafu { + msg: format!( + "SHOW CREATE exceeds {MAX_SQL_BYTES} SQL bytes ({}): {sql}", + sql.len() + ), + } + .fail(); + } + if batch.len() == MAX_STATEMENTS || bytes + sql.len() > MAX_SQL_BYTES { + break; + } + bytes += sql.len(); + if let Some(sql) = statements.next() { + batch.push(sql); + } + } + Ok(batch) +} + +fn append_results( + batch: &[String], + response: &GreptimedbV1Response, + ddl: &mut String, +) -> Result<()> { + let outputs = response.output(); + if outputs.len() != batch.len() { + return UnexpectedSnafu { + msg: format!( + "SHOW CREATE expected {} results, received {} for {}", + batch.len(), + outputs.len(), + batch.concat() + ), + } + .fail(); + } + for (sql, output) in batch.iter().zip(outputs) { + let GreptimeQueryOutput::Records(records) = output else { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE records for {sql}"), + } + .fail(); + }; + if records.num_cols() != 2 || records.rows().len() != 1 || records.rows()[0].len() != 2 { + return UnexpectedSnafu { + msg: format!("expected one two-column SHOW CREATE row for {sql}"), + } + .fail(); + } + let row = &records.rows()[0]; + let Some(create) = row[1].as_str().filter(|value| !value.is_empty()) else { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE DDL string for {sql}"), + } + .fail(); + }; + if !row[0].is_string() { + return UnexpectedSnafu { + msg: format!("expected SHOW CREATE object name for {sql}"), + } + .fail(); + } + ddl.push_str(create); + ddl.push_str(";\n"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; + + use super::*; + + fn records(name: &str, ddl: &str) -> Value { + json!({"records": {"schema": {"column_schemas": [ + {"name": "Table", "data_type": "String"}, + {"name": "Create Table", "data_type": "String"} + ]}, "rows": [[name, ddl]], "total_rows": 1}}) + } + + #[test] + fn batch_limits_include_quoted_utf8_and_separators() { + let special = show_create_sql("TABLE", "cat.name", "s\";库", Some("ta.\";表")); + assert_eq!( + special, + "SHOW CREATE TABLE \"cat.name\".\"s\"\";库\".\"ta.\"\";表\";\n" + ); + let mut statements = std::iter::repeat_n(special.clone(), 129).peekable(); + assert_eq!( + next_batch(&mut statements).unwrap(), + vec![special.clone(); 128] + ); + assert_eq!(next_batch(&mut statements).unwrap(), vec![special]); + assert!(next_batch(&mut statements).unwrap().is_empty()); + + let prefix = show_create_sql("TABLE", "c", "s", Some("")); + // Every quote expands to two bytes; this also exercises a multibyte name. + let name = format!("表{}", "\"".repeat((MAX_SQL_BYTES - prefix.len() - 3) / 2)); + let mut exact = show_create_sql("TABLE", "c", "s", Some(&name)); + let padding = MAX_SQL_BYTES - exact.len(); + exact = show_create_sql( + "TABLE", + "c", + "s", + Some(&format!("{name}{}", "a".repeat(padding))), + ); + assert_eq!(exact.len(), MAX_SQL_BYTES); + let small = show_create_sql("TABLE", "c", "s", Some("next")); + let mut statements = [exact.clone(), small.clone()].into_iter().peekable(); + assert_eq!(next_batch(&mut statements).unwrap(), vec![exact.clone()]); + assert_eq!(next_batch(&mut statements).unwrap(), vec![small]); + let mut oversized = std::iter::once(format!("{exact} ")).peekable(); + assert!( + next_batch(&mut oversized) + .unwrap_err() + .to_string() + .contains("exceeds 131072") + ); + + let half = show_create_sql( + "TABLE", + "c", + "s", + Some(&"a".repeat(MAX_SQL_BYTES / 2 - prefix.len())), + ); + let mut statements = [half.clone(), half.clone(), prefix.clone()] + .into_iter() + .peekable(); + assert_eq!( + next_batch(&mut statements).unwrap(), + vec![half.clone(), half] + ); + assert_eq!(next_batch(&mut statements).unwrap(), vec![prefix]); + } + + #[test] + fn show_create_requires_exact_record_shape_and_count() { + let batch = vec![ + "SHOW CREATE TABLE first".into(), + "SHOW CREATE TABLE middle".into(), + ]; + let valid = records("first", "CREATE TABLE first"); + for output in [ + vec![], + vec![valid.clone()], + vec![valid.clone(); 3], + vec![valid.clone(), json!({"affectedrows": 1})], + vec![valid.clone(), records("middle", "")], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([]); + r + }], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([["middle", 4]]); + r + }], + vec![valid.clone(), { + let mut r = valid.clone(); + r["records"]["rows"] = json!([["middle", "ddl", "extra"]]); + r + }], + ] { + let response = + serde_json::from_value(json!({"output": output, "execution_time_ms": 0})).unwrap(); + let error = append_results(&batch, &response, &mut String::new()).unwrap_err(); + assert!(error.to_string().contains("middle"), "{error}"); + } + } + + async fn start_mock_sql_server( + bodies: Vec<(u16, String)>, + ) -> (DatabaseClient, tokio::task::JoinHandle>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + tokio::time::timeout(std::time::Duration::from_secs(5), async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = BufReader::new(socket); + let mut queries = Vec::new(); + for (status, body) in bodies { + let mut headers = String::new(); + loop { + let mut line = String::new(); + if socket.read_line(&mut line).await.unwrap() == 0 { + assert!(headers.is_empty()); + return queries; + } + headers.push_str(&line); + if line == "\r\n" { + break; + } + } + let headers = headers.to_lowercase(); + assert!(headers.starts_with("post /v1/sql ")); + assert!(headers.contains("authorization: basic dxnlcjpwyxnzd29yza==")); + assert!(headers.contains("x-greptime-timeout: 7s")); + let length: usize = headers + .lines() + .find_map(|line| line.strip_prefix("content-length: ")) + .unwrap() + .parse() + .unwrap(); + let mut form = vec![0; length]; + socket.read_exact(&mut form).await.unwrap(); + let params: std::collections::HashMap<_, _> = + url::form_urlencoded::parse(&form).into_owned().collect(); + assert_eq!(params["db"], "catalog-public"); + queries.push(params["sql"].clone()); + socket + .get_mut() + .write_all( + format!( + "HTTP/1.1 {status} Response\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); + } + queries + }) + .await + .expect("mock SQL exchange timed out") + }); + let client = DatabaseClient::new( + address.to_string(), + "catalog".into(), + Some("user:password".into()), + std::time::Duration::from_secs(7), + None, + true, + ); + (client, server) + } + + #[tokio::test] + async fn batches_preserve_order_context_and_reuse_connection() { + let names: Vec<_> = (0..129).map(|i| format!("表.{i}\";x")).collect(); + let outputs: Vec<_> = names + .iter() + .map(|name| records(name, &format!("DDL {name}"))) + .collect(); + let bodies = outputs + .chunks(128) + .map(|chunk| { + ( + 200, + json!({"output": chunk, "execution_time_ms": 0}).to_string(), + ) + }) + .collect(); + let (client, server) = start_mock_sql_server(bodies).await; + let mut ddl = String::new(); + append_schema_ddl( + &client.clone(), + "catalog", + "schema", + names.iter().map(|name| ("TABLE", Some(name.as_str()))), + &mut ddl, + ) + .await + .unwrap(); + assert_eq!( + ddl, + names + .iter() + .map(|name| format!("DDL {name};\n")) + .collect::() + ); + let queries = server.await.unwrap(); + let expected: Vec<_> = (0..129) + .map(|i| format!("SHOW CREATE TABLE \"catalog\".\"schema\".\"表.{i}\"\";x\";\n")) + .collect(); + assert_eq!( + queries, + vec![expected[..128].concat(), expected[128..].concat()] + ); + } + + #[tokio::test] + async fn legacy_command_stops_before_data_after_schema_failure() { + use clap::Parser; + + use crate::data::export::ExportCommand; + + let output = |rows: Value, columns: usize| { + json!({"records": {"schema": {"column_schemas": (0..columns) + .map(|i| json!({"name": i.to_string(), "data_type": "String"})) + .collect::>()}, "rows": rows, "total_rows": rows.as_array().unwrap().len()}}) + }; + let databases = output(json!([["public"]]), 1); + let mut malformed = records("middle", "CREATE TABLE middle"); + malformed["records"]["rows"][0][1] = json!(42); + let bodies = vec![ + vec![databases.clone()], + vec![records("public", "CREATE DATABASE public")], + vec![databases.clone()], + vec![output(json!([]), 3)], + vec![output( + json!([ + ["catalog", "public", "first", "BASE TABLE"], + ["catalog", "public", "middle", "BASE TABLE"], + ["catalog", "public", "last", "BASE TABLE"] + ]), + 4, + )], + vec![ + records("first", "CREATE TABLE first"), + malformed, + records("last", "CREATE TABLE last"), + ], + vec![databases], + vec![json!({"affectedrows": 0})], + ] + .into_iter() + .map(|output| { + ( + 200, + json!({"output": output, "execution_time_ms": 0}).to_string(), + ) + }) + .collect::>(); + for target in ["schema", "all"] { + let (client, server) = start_mock_sql_server(bodies.clone()).await; + let dir = tempfile::tempdir().unwrap(); + let command = ExportCommand::parse_from([ + "export", + "--addr", + client.addr(), + "--database", + "catalog-public", + "--output-dir", + dir.path().to_str().unwrap(), + "--target", + target, + "--auth-basic", + "user:password", + "--timeout", + "7s", + "--no-proxy", + ]); + let tool = command.build().await.unwrap(); + let result = tool.do_work().await; + drop(tool); + let queries = server.await.unwrap(); + assert!(result.unwrap_err().to_string().contains("middle")); + assert!( + dir.path() + .join("catalog/public/create_database.sql") + .is_file() + ); + assert!(!dir.path().join("catalog/public/create_tables.sql").exists()); + assert_eq!(queries.len(), 6); + assert!(queries[5].contains("SHOW CREATE TABLE \"catalog\".\"public\".\"middle\"")); + assert!(!queries.iter().any(|sql| sql.starts_with("COPY"))); + } + } + + #[tokio::test] + async fn sql_response_errors_do_not_expose_copy_credentials() { + let secret = "sentinel-connection-secret"; + let echoed = "sentinel-response-secret"; + let sql = format!( + "COPY DATABASE public TO 's3://bucket/' CONNECTION (SECRET_ACCESS_KEY='{secret}')" + ); + let bodies = vec![ + (200, format!("not json {echoed}")), + (200, json!({"output": [{"records": echoed}], "execution_time_ms": 0}).to_string()), + (400, json!({"code": 3001, "error": echoed}).to_string()), + (200, json!({"error": echoed}).to_string()), + (200, json!({"code": 3001, "output": [records("first", echoed)], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [{"error": echoed}]}).to_string()), + (200, json!({"output": [{"code": 3001, "records": records("first", echoed)["records"]}], "execution_time_ms": 0}).to_string()), + ]; + let count = bodies.len(); + let (client, server) = start_mock_sql_server(bodies).await; + for _ in 0..count { + let error = client + .sql_response(&sql, "public") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("SQL"), "{error}"); + assert!(!error.contains(secret), "{error}"); + assert!(!error.contains(echoed), "{error}"); + assert!(!error.contains("COPY DATABASE"), "{error}"); + } + assert_eq!(server.await.unwrap().len(), count); + } + + #[tokio::test] + async fn rejects_http_and_statement_errors_without_retry() { + let valid = records("first", "CREATE TABLE first"); + let bodies = vec![ + (400, json!({"code": 3001, "error": "middle not found", "execution_time_ms": 0}).to_string()), + (200, json!({"code": 3001, "output": [valid.clone(), valid.clone(), valid.clone()], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [valid.clone(), {"error": "middle denied"}, valid.clone()], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [valid.clone(), {"error": "middle denied", "records": valid["records"]}, valid], "execution_time_ms": 0}).to_string()), + (200, json!({"output": [records("first", "DDL"), {"code": 3001, "records": records("middle", "DDL")["records"]}, records("last", "DDL")], "execution_time_ms": 0}).to_string()), + (200, r#"{"execution_time_ms":0}"#.into()), + (200, "not json".into()), + ]; + let count = bodies.len(); + let (client, server) = start_mock_sql_server(bodies).await; + for _ in 0..count { + let mut ddl = String::new(); + let error = append_schema_ddl( + &client, + "catalog", + "schema", + ["first", "middle", "last"] + .into_iter() + .map(|name| ("TABLE", Some(name))), + &mut ddl, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("middle"), "{error}"); + assert!(ddl.is_empty()); + } + assert_eq!(server.await.unwrap().len(), count); + } +} diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index ac314bc5899..da277e166be 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; use std::time::Duration; use base64::Engine; @@ -37,6 +38,7 @@ pub struct DatabaseClient { timeout: Duration, proxy: Option, no_proxy: bool, + client: Arc>, } pub fn parse_proxy_opts( @@ -86,6 +88,7 @@ impl DatabaseClient { timeout, proxy, no_proxy, + client: Arc::new(tokio::sync::OnceCell::new()), } } @@ -109,15 +112,7 @@ impl DatabaseClient { async fn require_packed_capability(&self, capability: &str) -> Result<()> { let url = format!("http://{}/v1/capabilities", self.addr); - let mut builder = reqwest::Client::builder().timeout(self.timeout); - if let Some(proxy) = self.proxy.clone() { - builder = builder.proxy(proxy); - } - if self.no_proxy { - builder = builder.no_proxy(); - } - let client = builder.build().context(BuildClientSnafu)?; - let mut request = client.get(&url); + let mut request = self.http_client().await?.get(&url).timeout(self.timeout); if let Some(auth) = &self.auth_header { request = request.header("Authorization", auth); } @@ -173,15 +168,9 @@ impl DatabaseClient { ("db", format!("{}-{}", self.catalog, schema)), ("sql", sql.to_string()), ]; - let mut builder = reqwest::Client::builder(); - if let Some(proxy) = self.proxy.clone() { - builder = builder.proxy(proxy); - } - if self.no_proxy { - builder = builder.no_proxy(); - } - let client = builder.build().context(BuildClientSnafu)?; - let mut request = client + let mut request = self + .http_client() + .await? .post(&url) .form(¶ms) .header("Content-Type", "application/x-www-form-urlencoded"); @@ -197,17 +186,69 @@ impl DatabaseClient { let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { reason: format!("bad url: {}", url), })?; - let response = response - .error_for_status() - .with_context(|_| HttpQuerySqlSnafu { - reason: format!("query failed: {}", sql), - })?; - + let status = response.status(); let text = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot get response text".to_string(), })?; + let value: Value = serde_json::from_str(&text).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: format!("invalid SQL response ({status})"), + } + .build() + })?; + if !status.is_success() + || value.get("error").is_some() + || value + .get("code") + .is_some_and(|code| code.as_u64() != Some(0)) + { + return crate::error::UnexpectedSnafu { + msg: format!( + "SQL request failed ({status}, code {:?})", + value.get("code").and_then(Value::as_u64) + ), + } + .fail(); + } + if let Some(outputs) = value.get("output").and_then(Value::as_array) { + for (index, output) in outputs.iter().enumerate() { + if output.get("error").is_some() + || output + .get("code") + .is_some_and(|code| code.as_u64() != Some(0)) + { + return crate::error::UnexpectedSnafu { + msg: format!( + "SQL statement {} failed (code {:?})", + index + 1, + output.get("code").and_then(Value::as_u64) + ), + } + .fail(); + } + } + } + serde_json::from_value(value).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: format!("invalid SQL response ({status})"), + } + .build() + }) + } - serde_json::from_str::(&text).context(SerdeJsonSnafu) + async fn http_client(&self) -> Result<&reqwest::Client> { + self.client + .get_or_try_init(|| async { + let mut builder = reqwest::Client::builder(); + if let Some(proxy) = self.proxy.clone() { + builder = builder.proxy(proxy); + } + if self.no_proxy { + builder = builder.no_proxy(); + } + builder.build().context(BuildClientSnafu) + }) + .await } } diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 80169a391cd..da1232a243b 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1541,6 +1541,20 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { .await; sql(instance, "INSERT INTO audit VALUES ('a',1,1),('z',NULL,3)").await; sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; + let special_name = "audit.dashboard"; + let quoted_special = special_name.replace('"', "\"\""); + sql( + instance, + &format!("CREATE VIEW \"{quoted_special}\" AS SELECT * FROM audit"), + ) + .await; + if packed && !s3 { + for id in 0..128 { + let name = format!("batch_{id:03}"); + sql(instance, &format!("CREATE TABLE {name} (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='v2_a')")).await; + names.push(name); + } + } sql(instance, "CREATE DATABASE z_later").await; if packed { sql(instance, "CREATE DATABASE empty_schema").await; @@ -1573,6 +1587,21 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { names.push("bulk".into()); } let (addr, server) = export_http(instance.clone()).await; + if packed && !s3 { + let client = cli::DatabaseClient::new( + addr.clone(), + "greptime".into(), + None, + std::time::Duration::from_secs(60), + None, + true, + ); + let error = client.sql_in_public( + "SHOW CREATE TABLE audit; SHOW CREATE TABLE missing_middle; SHOW CREATE VIEW dashboard" + ).await.unwrap_err(); + assert!(error.to_string().contains("SQL request failed"), "{error}"); + } + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); for experimental in [false, true] { for layout in ["packed", "invalid"] { @@ -1876,6 +1905,23 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { ) .await ); + for view in ["dashboard", special_name] { + let query = format!( + "SELECT * FROM \"{}\" ORDER BY ts,host", + view.replace('"', "\"\"") + ); + assert_eq!( + values(instance, &query).await, + values(target.fe_instance(), &query).await + ); + assert_eq!( + table(instance, view).await.schema().column_schemas(), + table(target.fe_instance(), view) + .await + .schema() + .column_schemas() + ); + } if packed { let schema_uri = format!("{uri}/schema-only"); let mut schema_args = vec![ From a314ac428144b32decdd85866011e3d29f0025b0 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 30 Sep 2026 08:40:27 +0000 Subject: [PATCH 082/133] fix(mito2): truncate parquet column index min/max for SST writes (#9420) * fix(mito2): truncate parquet column index min/max for SST writes The SST writer and the bulk part encoder disabled column index truncation together with statistics truncation in #6977. Mito never reads the column index, but every page of a large string or binary column still stored its full min and max values, uncompressed. Truncate the column index to the parquet default of 64 bytes and keep chunk statistics untruncated, since primary key pruning decodes them. Signed-off-by: Dennis Zhuang * docs(mito2): note that the column index must not be used to decode primary keys Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/memtable/bulk/part.rs | 6 +- src/mito2/src/sst/parquet.rs | 114 ++++++++++++++++++++++++++++ src/mito2/src/sst/parquet/writer.rs | 5 +- 3 files changed, 121 insertions(+), 4 deletions(-) diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 7e50fbbd4cb..444cc59e317 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -71,7 +71,9 @@ use crate::sst::SeriesEstimator; use crate::sst::index::IndexOutput; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; -use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding}; +use crate::sst::parquet::{ + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding, +}; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1334,7 +1336,7 @@ impl BulkPartEncoder { .set_write_batch_size(row_group_size) .set_max_row_group_row_count(Some(row_group_size)) .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_column_index_truncate_length(None) + .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); props = apply_float_field_encoding(props, &metadata, float_field_encoding); let writer_props = Some(props.build()); diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 301842bde72..9404d4484fc 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -77,6 +77,18 @@ pub(crate) struct Json2TargetLayout { /// batching without changing the row group layout of newly written SSTs. pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; +/// Truncation length of min/max values in the parquet column index. +/// +/// Mito never reads the column index (it is skipped or stripped on load), so truncating +/// it is safe. Chunk statistics stay untruncated because pruning decodes primary keys +/// from them. Without truncation every page of a large string column stores its whole +/// min and max values, uncompressed. +/// +/// Truncated values are only bounds: code that starts reading the column index must not +/// decode primary keys from it. +pub(crate) const COLUMN_INDEX_TRUNCATE_LENGTH: Option = + parquet::file::properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; + /// Applies the configured encoding to direct floating-point field columns. pub(crate) fn apply_float_field_encoding( mut builder: WriterPropertiesBuilder, @@ -901,6 +913,82 @@ mod tests { .await; } + #[tokio::test] + async fn test_column_index_truncates_large_field_values() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let handle = sst_file_handle(0, 1000); + let file_path = FixedPathProvider { + region_file_id: handle.file_id(), + }; + let metadata = build_test_binary_test_region_metadata(); + // A tag long enough that the encoded primary key exceeds the 64-byte truncation length. + let tag = "t".repeat(100); + let values: Vec> = (0..64) + .map(|i| format!("{i:08}").into_bytes().repeat(4 * 1024)) + .collect(); + let batch = new_record_batch_with_binary_values(&tag, &values); + let mut metrics = Metrics::new(WriteType::Flush); + let mut writer = ParquetWriter::new_with_object_store( + object_store.clone(), + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + file_path, + &mut metrics, + ) + .await; + writer + .write_all_flat_as_primary_key( + new_flat_source_from_record_batches(vec![batch]), + None, + &WriteOptions::default(), + ) + .await + .unwrap(); + + let path = handle.file_path(FILE_DIR, PathType::Bare); + let bytes = object_store.read(&path).await.unwrap().to_bytes(); + let options = parquet::arrow::arrow_reader::ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options).unwrap(); + let parquet_meta = builder.metadata().clone(); + let schema = parquet_meta.file_metadata().schema_descr(); + let column = |name: &str| { + (0..schema.num_columns()) + .find(|i| schema.column(*i).name() == name) + .unwrap() + }; + + let field_index = &parquet_meta.column_index().unwrap()[0][column("field_0")]; + let parquet::file::page_index::column_index::ColumnIndexMetaData::BYTE_ARRAY(field_index) = + field_index + else { + panic!("unexpected column index: {field_index:?}"); + }; + let mut pages = 0; + for (min, max) in field_index + .min_values_iter() + .zip(field_index.max_values_iter()) + { + assert!(min.unwrap().len() <= 64); + assert!(max.unwrap().len() <= 64); + pages += 1; + } + assert!(pages > 1, "values should span several pages"); + + // Chunk statistics keep exact primary keys; pruning decodes them. + let pk_stats = parquet_meta + .row_group(0) + .column(column(store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME)) + .statistics() + .unwrap(); + let pk = new_primary_key(&[&tag]); + assert!(pk.len() > 64); + assert_eq!(pk_stats.max_bytes_opt().unwrap(), pk.as_slice()); + } + #[rstest::rstest] #[tokio::test] async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { @@ -1475,6 +1563,32 @@ mod tests { assert!(cached.contains_row_group(3)); } + fn new_record_batch_with_binary_values(tag: &str, values: &[Vec]) -> RecordBatch { + let metadata = build_test_binary_test_region_metadata(); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + let num_rows = values.len(); + let mut tag_0_builder = StringDictionaryBuilder::::new(); + let mut pk_builder = BinaryDictionaryBuilder::::new(); + let pk = new_primary_key(&[tag]); + for _ in 0..num_rows { + tag_0_builder.append_value(tag); + pk_builder.append(&pk).unwrap(); + } + let columns: Vec = vec![ + Arc::new(tag_0_builder.finish()), + Arc::new(datatypes::arrow::array::BinaryArray::from_iter_values( + values.iter().map(|v| v.as_slice()), + )), + Arc::new(TimestampMillisecondArray::from_iter_values( + 0..num_rows as i64, + )), + Arc::new(pk_builder.finish()), + Arc::new(UInt64Array::from_value(1000, num_rows)), + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)), + ]; + RecordBatch::try_new(flat_schema, columns).unwrap() + } + fn new_record_batch_with_binary(tags: &[&str], start: usize, end: usize) -> RecordBatch { assert!(end >= start); let metadata = build_test_binary_test_region_metadata(); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9ef15add34..932654dfdd8 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -62,7 +62,8 @@ use crate::sst::parquet::flat_format::{ }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{ - PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding, + COLUMN_INDEX_TRUNCATE_LENGTH, PARQUET_METADATA_KEY, SstInfo, WriteOptions, + apply_float_field_encoding, }; use crate::sst::{ DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -544,7 +545,7 @@ where .set_compression(Compression::ZSTD(ZstdLevel::default())) .set_encoding(Encoding::PLAIN) .set_max_row_group_row_count(Some(opts.row_group_size)) - .set_column_index_truncate_length(None) + .set_column_index_truncate_length(COLUMN_INDEX_TRUNCATE_LENGTH) .set_statistics_truncate_length(None); let ts_col = ColumnPath::new(vec![ self.metadata.time_index_column().column_schema.name.clone(), From d7f53318766817d2c39da155bfa1fcec46a73cd4 Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Wed, 30 Sep 2026 10:40:17 +0000 Subject: [PATCH 083/133] feat: add CORS support for gRPC-Web on the frontend gRPC server (#9374) * feat: add CORS support for gRPC-Web on the frontend gRPC server A browser-based gRPC-Web client cannot read a cross-origin response. The request carries a non-simple content type, so the browser sends an `OPTIONS` preflight first, and it drops the call when the response has no `Access-Control-Allow-Origin`. Add `enable_cors` and `cors_allowed_origins` to `[grpc]`, threaded from `GrpcOptions` to `GrpcServerConfig` the same way `tls` is. The CORS layer sits outside `tonic_web::GrpcWebLayer`, so it answers the preflight before the request reaches the gRPC routes. It exposes `grpc-status`, `grpc-message` and `grpc-status-details-bin`, which a browser cannot read otherwise. An empty `cors_allowed_origins` allows any origin. `enable_cors` defaults to false, unlike `[http]`. `GrpcOptions` is shared by the public `[grpc]` section and the internal `[internal_grpc]` one, and serde cannot tell them apart, so a true default would also turn CORS on for the internal gRPC listeners: the frontend internal gRPC server, and the datanode and flownode gRPC servers. None of them authenticate callers, and a browser on the host or in the cluster network can reach all of them. Set `enable_cors = true` to turn it on. Tests start a real server on an ephemeral port and cover the preflight, a custom origin list, and the disabled case. Update the example TOMLs and regenerate config/config.md. Signed-off-by: lczllx <2181719471@qq.com> * fix: drop the redundant OPTIONS from the gRPC CORS allow_methods Signed-off-by: lczllx <2181719471@qq.com> * test: assert the allow-headers and allow-methods headers in the gRPC CORS preflight The preflight answers with `access-control-allow-headers: *` from `AllowHeaders::any()`, and with `access-control-allow-methods: post` from the POST-only `allow_methods` list. Pin both, the way the HTTP CORS test pins its own headers: a missing allow-headers header would let the preflight pass the origin check and still have the browser block every gRPC-Web call, which a non-browser client would never notice. Signed-off-by: lczllx <2181719471@qq.com> * docs(config): stop the example configs from setting the new gRPC CORS keys Signed-off-by: lczllx <2181719471@qq.com> * test(servers): cover the exposed gRPC CORS headers and pin the option plumbing Signed-off-by: lczllx <2181719471@qq.com> * docs(config): document the gRPC CORS origin example with `#+` Signed-off-by: lczllx <2181719471@qq.com> * fix(frontend): keep CORS off the internal gRPC server Signed-off-by: lczllx <2181719471@qq.com> * docs(grpc): document that the internal gRPC server never serves CORS Signed-off-by: lczllx <2181719471@qq.com> * Update src/servers/src/grpc.rs Signed-off-by: Dennis Zhuang * chore: trim redundant comments Signed-off-by: Dennis Zhuang * refactor: share CORS origin parsing and simplify gRPC-Web CORS tests Signed-off-by: Dennis Zhuang * refactor: enable gRPC CORS only on the frontend public server Signed-off-by: Dennis Zhuang * docs(config): clarify the gRPC CORS options Signed-off-by: Dennis Zhuang --------- Signed-off-by: lczllx <2181719471@qq.com> Signed-off-by: Dennis Zhuang Co-authored-by: Dennis Zhuang --- config/config.md | 4 + config/frontend.example.toml | 5 + config/standalone.example.toml | 5 + src/cmd/tests/load_config_test.rs | 1 + src/frontend/src/frontend.rs | 3 +- src/frontend/src/server.rs | 71 +++++++++++- src/servers/src/grpc.rs | 173 +++++++++++++++++++++++++++++- src/servers/src/grpc/builder.rs | 9 ++ src/servers/src/http.rs | 27 +++-- tests-integration/tests/http.rs | 2 + 10 files changed, 280 insertions(+), 20 deletions(-) diff --git a/config/config.md b/config/config.md index 06d3b89e075..24351b2b6fd 100644 --- a/config/config.md +++ b/config/config.md @@ -43,6 +43,8 @@ | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | +| `grpc.enable_cors` | Bool | `false` | Enable CORS for gRPC-Web clients in browsers. Disabled by default. | +| `grpc.cors_allowed_origins` | Array | Unset | Origins allowed by gRPC CORS. An empty list allows any origin. | | `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | @@ -303,6 +305,8 @@ | `grpc.server_addr` | String | `127.0.0.1:4001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | | `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `grpc.enable_cors` | Bool | `false` | Enable CORS for gRPC-Web clients in browsers. Disabled by default. | +| `grpc.cors_allowed_origins` | Array | Unset | Origins allowed by gRPC CORS. An empty list allows any origin. | | `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index b528da0304f..2f2c15cfc57 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -96,6 +96,11 @@ runtime_size = 8 ## - `all`: enable all compression. ## Default to `none` flight_compression = "arrow_ipc" +## Enable CORS for gRPC-Web clients in browsers. Disabled by default. +#+ enable_cors = false +## Origins allowed by gRPC CORS. An empty list allows any origin. +## @toml2docs:none-default +#+ cors_allowed_origins = ["https://example.com"] ## The maximum connection age for gRPC connection. ## The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour. ## Refer to https://grpc.io/docs/guides/keepalive/ for more details. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index c09c815aed0..4760e038f49 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -111,6 +111,11 @@ api_server_addr = "127.0.0.1:4006" bind_addr = "127.0.0.1:4001" ## The number of server worker threads. runtime_size = 8 +## Enable CORS for gRPC-Web clients in browsers. Disabled by default. +#+ enable_cors = false +## Origins allowed by gRPC CORS. An empty list allows any origin. +## @toml2docs:none-default +#+ cors_allowed_origins = ["https://example.com"] ## The maximum connection age for gRPC connection. ## The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour. ## Refer to https://grpc.io/docs/guides/keepalive/ for more details. diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index cc095c5c92e..b56299f9916 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -464,6 +464,7 @@ fn test_load_standalone_example_config() { cors_allowed_origins: vec!["https://example.com".to_string()], ..Default::default() }, + grpc: GrpcOptions::default(), query: QueryOptions { memory_pool_size: MemoryLimit::Percentage(50), ..Default::default() diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 71e5888b8fa..90e7abc4517 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -61,7 +61,8 @@ pub struct FrontendOptions { pub http: HttpOptions, pub grpc: GrpcOptions, /// The internal gRPC options for the frontend service. - /// it provide the same service as the public gRPC service, just only for internal use. + /// It serves the same services as the public one plus the internal handler. + /// CORS is always off on it. pub internal_grpc: Option, pub mysql: MysqlOptions, pub postgres: PostgresOptions, diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index a5568b24be5..cf278f1f1d0 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -215,11 +215,15 @@ where external: bool, request_memory_limiter: ServerMemoryLimiter, ) -> Result { - let builder = if let Some(builder) = self.grpc_server_builder.take() { + let mut builder = if let Some(builder) = self.grpc_server_builder.take() { builder } else { self.grpc_server_builder(grpc, request_memory_limiter)? }; + // Browsers never talk to the internal server. + if external && grpc.enable_cors { + builder = builder.with_cors(grpc.cors_allowed_origins.clone()); + } let user_provider = if external { self.plugins.get::() @@ -522,6 +526,7 @@ mod tests { use client::{Client, Database}; use common_grpc::channel_manager::ChannelManager; use meta_client::client::MetaClientBuilder; + use reqwest::header::{ACCESS_CONTROL_ALLOW_ORIGIN, ACCESS_CONTROL_REQUEST_METHOD, ORIGIN}; use servers::grpc::GRPC_SERVER; use servers::grpc::flight::{FlightCraft, FlightCraftRef, TonicStream}; use tonic::{Code, Request, Response, Status, Streaming}; @@ -1029,4 +1034,68 @@ mod tests { // Assert assert!(health_check.is_ok()); } + + #[tokio::test] + async fn test_internal_grpc_server_never_serves_cors() { + let options = FrontendOptions { + http: HttpOptions { + addr: "127.0.0.1:0".to_string(), + ..Default::default() + }, + grpc: GrpcOptions { + enable_cors: true, + ..GrpcOptions::default().with_bind_addr("127.0.0.1:0") + }, + internal_grpc: Some(GrpcOptions { + enable_cors: true, + ..GrpcOptions::default().with_bind_addr("127.0.0.1:0") + }), + mysql: crate::service_config::MysqlOptions { + enable: false, + ..Default::default() + }, + postgres: crate::service_config::PostgresOptions { + enable: false, + ..Default::default() + }, + ..Default::default() + }; + let meta_client = Arc::new( + MetaClientBuilder::new(0, Role::Frontend) + .enable_procedure() + .build(), + ); + let instance = Arc::new( + FrontendBuilder::new_test(&options, meta_client) + .try_build() + .await + .unwrap(), + ); + let mut services = Services::new(options, instance, Default::default()) + .build() + .unwrap(); + + services.start_all().await.unwrap(); + let public_addr = services.addr(GRPC_SERVER).unwrap(); + let internal_addr = services.addr("INTERNAL_GRPC_SERVER").unwrap(); + let public = send_cors_preflight(public_addr).await; + let internal = send_cors_preflight(internal_addr).await; + services.shutdown_all().await.unwrap(); + + assert!(public.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN)); + assert!(!internal.headers().contains_key(ACCESS_CONTROL_ALLOW_ORIGIN)); + } + + async fn send_cors_preflight(addr: std::net::SocketAddr) -> reqwest::Response { + reqwest::Client::new() + .request( + reqwest::Method::OPTIONS, + format!("http://{addr}/greptime.v1.HealthCheck/Check"), + ) + .header(ORIGIN, "https://example.com") + .header(ACCESS_CONTROL_REQUEST_METHOD, "POST") + .send() + .await + .unwrap() + } } diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs index 1be1f8215b5..9e9335b422d 100644 --- a/src/servers/src/grpc.rs +++ b/src/servers/src/grpc.rs @@ -36,6 +36,7 @@ use common_grpc::channel_manager::{ }; use common_telemetry::{error, info, warn}; use futures::FutureExt; +use http::{HeaderName, Method}; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsServiceServer; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; @@ -48,9 +49,11 @@ use tonic::transport::ServerTlsConfig; use tonic::transport::server::TcpIncoming; use tonic::{Request, Response, Status}; use tonic_reflection::server::v1::{ServerReflection, ServerReflectionServer}; +use tower_http::cors::{AllowHeaders, CorsLayer}; use crate::error::{AlreadyStartedSnafu, InternalSnafu, Result, StartGrpcSnafu, TcpBindSnafu}; use crate::grpc::memory_limit::MemoryLimiterExtensionService; +use crate::http::cors_allow_origin; use crate::install_default_crypto_provider; use crate::metrics::MetricsMiddlewareLayer; use crate::otel_arrow::{HeaderInterceptor, OtelArrowServiceHandler}; @@ -87,6 +90,11 @@ pub struct GrpcOptions { /// The HTTP/2 keep-alive timeout. #[serde(with = "humantime_serde")] pub http2_keep_alive_timeout: Duration, + /// Whether to enable CORS, required by gRPC-Web clients in browsers. + /// Only the frontend's public gRPC server honors it. + pub enable_cors: bool, + /// Origins allowed by CORS. Empty allows any origin. + pub cors_allowed_origins: Vec, } impl GrpcOptions { @@ -155,6 +163,8 @@ impl Default for GrpcOptions { max_connection_age: None, http2_keep_alive_interval: Duration::from_secs(10), http2_keep_alive_timeout: Duration::from_secs(3), + enable_cors: false, + cors_allowed_origins: Vec::new(), } } } @@ -176,6 +186,8 @@ impl GrpcOptions { max_connection_age: None, http2_keep_alive_interval: Duration::from_secs(10), http2_keep_alive_timeout: Duration::from_secs(3), + enable_cors: false, + cors_allowed_origins: Vec::new(), } } @@ -237,6 +249,8 @@ pub struct GrpcServer { bind_addr: Option, name: Option, config: GrpcServerConfig, + /// `None` disables CORS. + cors_allowed_origins: Option>, } /// Grpc Server configuration @@ -265,6 +279,24 @@ impl Default for GrpcServerConfig { } impl GrpcServer { + fn cors_layer(&self) -> Result> { + let Some(allowed_origins) = &self.cors_allowed_origins else { + return Ok(None); + }; + let layer = CorsLayer::new() + .allow_methods([Method::POST]) + .allow_origin(cors_allow_origin(allowed_origins)?) + .allow_headers(AllowHeaders::any()) + // Trailers-only responses (most errors) carry `grpc-status` in HTTP + // headers, which browsers hide unless exposed. + .expose_headers([ + HeaderName::from_static("grpc-status"), + HeaderName::from_static("grpc-message"), + HeaderName::from_static("grpc-status-details-bin"), + ]); + Ok(Some(layer)) + } + pub fn create_healthcheck_service(&self) -> HealthCheckServer { HealthCheckServer::new(HealthCheckHandler) } @@ -359,13 +391,18 @@ impl Server for GrpcServer { (incoming, addr) }; - let metrics_layer = tower::ServiceBuilder::new() + let cors_layer = self.cors_layer()?; + + // `Server::layer` prepends, so the first layer added is outermost. CORS must + // sit outside `GrpcWebLayer`, which rejects the `OPTIONS` preflight with a 400. + let middleware_layer = tower::ServiceBuilder::new() .layer(MetricsMiddlewareLayer) + .option_layer(cors_layer) .into_inner(); let mut builder = tonic::transport::Server::builder() .accept_http1(true) - .layer(metrics_layer) + .layer(middleware_layer) .layer(tonic_web::GrpcWebLayer::new()); if let Some(tls_config) = self.tls_config.clone() { @@ -425,9 +462,23 @@ impl Server for GrpcServer { #[cfg(test)] mod tests { - use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; - use super::{DEFAULT_GRPC_ADDR_PORT, format_server_addr, port_from_bind_addr}; + use common_runtime::Runtime; + use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; + use http::header::{ + ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN, + ACCESS_CONTROL_EXPOSE_HEADERS, ACCESS_CONTROL_REQUEST_HEADERS, + ACCESS_CONTROL_REQUEST_METHOD, CONTENT_TYPE, ORIGIN, + }; + use http::{HeaderName, Method, StatusCode}; + + use super::{ + DEFAULT_GRPC_ADDR_PORT, GrpcServer, GrpcServerConfig, format_server_addr, + port_from_bind_addr, + }; + use crate::grpc::builder::GrpcServerBuilder; + use crate::server::Server; #[test] fn test_port_from_bind_addr() { @@ -454,4 +505,118 @@ mod tests { format_server_addr(IpAddr::V6(Ipv6Addr::LOCALHOST), 3002) ); } + + /// The returned server owns the shutdown sender and must outlive the requests. + async fn start_test_grpc_server( + cors_allowed_origins: Option>, + ) -> (GrpcServer, SocketAddr) { + let runtime = Runtime::builder().build().unwrap(); + let mut builder = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime); + if let Some(origins) = cors_allowed_origins { + builder = builder.with_cors(origins); + } + let mut server = builder.build(); + server + .start(SocketAddr::from((Ipv4Addr::LOCALHOST, 0))) + .await + .unwrap(); + let addr = server.bind_addr().unwrap(); + (server, addr) + } + + const TEST_PATH: &str = "greptime.v1.HealthCheck/Check"; + + async fn send_preflight(addr: SocketAddr, origin: &str) -> reqwest::Response { + reqwest::Client::new() + .request(Method::OPTIONS, format!("http://{addr}/{TEST_PATH}")) + .header(ORIGIN, origin) + .header(ACCESS_CONTROL_REQUEST_METHOD, "POST") + .header(ACCESS_CONTROL_REQUEST_HEADERS, "content-type,x-grpc-web") + .send() + .await + .unwrap() + } + + /// The path is not served, so the response is a trailers-only `Unimplemented`. + async fn send_grpc_web_call(addr: SocketAddr, origin: &str) -> reqwest::Response { + reqwest::Client::new() + .post(format!("http://{addr}/{TEST_PATH}")) + .header(ORIGIN, origin) + .header(CONTENT_TYPE, "application/grpc-web+proto") + .header("x-grpc-web", "1") + // Empty gRPC-Web frame: 1-byte compression flag + 4-byte length. + .body(vec![0u8; 5]) + .send() + .await + .unwrap() + } + + fn header(response: &reqwest::Response, name: HeaderName) -> Option<&str> { + response.headers().get(name).map(|v| v.to_str().unwrap()) + } + + #[tokio::test] + async fn test_grpc_web_cors_preflight() { + let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await; + + let response = send_preflight(addr, "https://example.com").await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*")); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_HEADERS), Some("*")); + // Lists the method of the actual call, not `OPTIONS`. + assert_eq!( + header(&response, ACCESS_CONTROL_ALLOW_METHODS), + Some("POST") + ); + } + + #[tokio::test] + async fn test_grpc_web_cors_expose_headers() { + let (_server, addr) = start_test_grpc_server(Some(Vec::new())).await; + + let response = send_grpc_web_call(addr, "https://example.com").await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), Some("*")); + // Premise: the response is trailers-only. + assert_eq!( + header(&response, HeaderName::from_static("grpc-status")), + Some("12") + ); + let expose_headers = header(&response, ACCESS_CONTROL_EXPOSE_HEADERS).unwrap(); + for name in ["grpc-status", "grpc-message", "grpc-status-details-bin"] { + assert!(expose_headers.contains(name), "{expose_headers}"); + } + } + + #[tokio::test] + async fn test_grpc_web_cors_custom_origins() { + let (_server, addr) = + start_test_grpc_server(Some(vec!["https://example.com".to_string()])).await; + + let response = send_preflight(addr, "https://example.com").await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), + Some("https://example.com") + ); + + // A disallowed origin still gets a 200, just without allow-origin; the + // browser does the blocking. + let response = send_preflight(addr, "https://notallowed.com").await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None); + } + + #[tokio::test] + async fn test_grpc_web_cors_disabled() { + let (_server, addr) = start_test_grpc_server(None).await; + + let response = send_preflight(addr, "https://example.com").await; + + // Without CORS, `GrpcWebLayer` rejects the preflight. + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(header(&response, ACCESS_CONTROL_ALLOW_ORIGIN), None); + } } diff --git a/src/servers/src/grpc/builder.rs b/src/servers/src/grpc/builder.rs index 99513dcceb3..eaa7b628961 100644 --- a/src/servers/src/grpc/builder.rs +++ b/src/servers/src/grpc/builder.rs @@ -97,6 +97,7 @@ pub struct GrpcServerBuilder { >, >, memory_limiter: ServerMemoryLimiter, + cors_allowed_origins: Option>, } impl GrpcServerBuilder { @@ -112,9 +113,16 @@ impl GrpcServerBuilder { tls_config: None, otel_arrow_service: None, memory_limiter, + cors_allowed_origins: None, } } + /// Enables CORS for gRPC-Web clients in browsers. An empty list allows any origin. + pub fn with_cors(mut self, allowed_origins: Vec) -> Self { + self.cors_allowed_origins = Some(allowed_origins); + self + } + /// Set a global memory limiter for all server protocols. pub fn with_memory_limiter(mut self, limiter: ServerMemoryLimiter) -> Self { self.memory_limiter = limiter; @@ -257,6 +265,7 @@ impl GrpcServerBuilder { bind_addr: None, name: self.name, config: self.config, + cors_allowed_origins: self.cors_allowed_origins, } } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index a33e52a1d40..e7cb715fa40 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -1061,20 +1061,7 @@ impl HttpServer { Method::DELETE, Method::HEAD, ]) - .allow_origin(if self.options.cors_allowed_origins.is_empty() { - AllowOrigin::from(Any) - } else { - AllowOrigin::from( - self.options - .cors_allowed_origins - .iter() - .map(|s| { - HeaderValue::from_str(s.as_str()) - .context(InvalidHeaderValueSnafu) - }) - .collect::>>()?, - ) - }) + .allow_origin(cors_allow_origin(&self.options.cors_allowed_origins)?) .allow_headers(Any), ) } else { @@ -1612,6 +1599,18 @@ impl HttpServer { pub const HTTP_SERVER: &str = "HTTP_SERVER"; pub const HTTP_API_SERVER: &str = "HTTP_API_SERVER"; +/// An empty list allows any origin. +pub(crate) fn cors_allow_origin(origins: &[String]) -> Result { + if origins.is_empty() { + return Ok(AllowOrigin::any()); + } + let origins = origins + .iter() + .map(|origin| HeaderValue::from_str(origin).context(InvalidHeaderValueSnafu)) + .collect::>>()?; + Ok(AllowOrigin::list(origins)) +} + #[async_trait] impl Server for HttpServer { async fn shutdown(&self) -> Result<()> { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 9b6ef9d8ae2..b1cfeac1ba1 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2835,6 +2835,8 @@ flight_compression = "arrow_ipc" runtime_size = 8 http2_keep_alive_interval = "10s" http2_keep_alive_timeout = "3s" +enable_cors = false +cors_allowed_origins = [] [grpc.tls] mode = "disable" From 0d502e1fc17c50dec869a39eb7d064a09e286aaa Mon Sep 17 00:00:00 2001 From: Yingwen Date: Thu, 1 Oct 2026 07:04:10 +0000 Subject: [PATCH 084/133] ci: upgrade EC2 runner action to restore runner registration (#9430) ci: update EC2 runner action for minimum runner version enforcement Signed-off-by: evenyag --- .github/actions/start-runner/action.yml | 2 +- .github/actions/stop-runner/action.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/actions/start-runner/action.yml b/.github/actions/start-runner/action.yml index 78568e26d76..44a5ff3feb3 100644 --- a/.github/actions/start-runner/action.yml +++ b/.github/actions/start-runner/action.yml @@ -56,7 +56,7 @@ runs: - name: Start EC2 runner if: startsWith(inputs.runner, 'ec2') - uses: machulav/ec2-github-runner@v2.3.8 + uses: machulav/ec2-github-runner@v2.6.1 id: start-linux-arm64-ec2-runner with: mode: start diff --git a/.github/actions/stop-runner/action.yml b/.github/actions/stop-runner/action.yml index e25a1d8dad7..fb84f2b7181 100644 --- a/.github/actions/stop-runner/action.yml +++ b/.github/actions/stop-runner/action.yml @@ -33,7 +33,7 @@ runs: - name: Stop EC2 runner if: ${{ inputs.label && inputs.ec2-instance-id }} - uses: machulav/ec2-github-runner@v2.3.8 + uses: machulav/ec2-github-runner@v2.6.1 with: mode: stop label: ${{ inputs.label }} From 3e86422ef863482b1fd7705d418d9eb41525d778 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Fri, 2 Oct 2026 03:45:42 +0000 Subject: [PATCH 085/133] chore: bump version and start new 1.4 release cycle (#9403) chore: start new 1.4 release cycle --- Cargo.lock | 154 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 2 +- 2 files changed, 78 insertions(+), 78 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 834b0b541bb..e985f483efb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -230,7 +230,7 @@ checksum = "d301b3b94cb4b2f23d7917810addbbaff90738e0ca2be692bd027e70d7e0330c" [[package]] name = "api" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow-schema 59.2.0", "common-base", @@ -959,7 +959,7 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "auth" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -1618,7 +1618,7 @@ dependencies = [ [[package]] name = "cache" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "catalog", "common-error", @@ -1684,7 +1684,7 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "catalog" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow 59.2.0", @@ -2040,7 +2040,7 @@ checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" [[package]] name = "cli" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-stream", "async-trait", @@ -2095,7 +2095,7 @@ dependencies = [ [[package]] name = "client" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arc-swap", @@ -2127,7 +2127,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", - "substrait 1.3.0-beta.1", + "substrait 1.4.0-alpha.0", "tokio", "tokio-stream", "tonic 0.14.2", @@ -2175,7 +2175,7 @@ dependencies = [ [[package]] name = "cmd" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2311,7 +2311,7 @@ dependencies = [ [[package]] name = "common-base" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "anymap2", @@ -2332,7 +2332,7 @@ dependencies = [ [[package]] name = "common-batcher" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "dashmap", "futures", @@ -2341,14 +2341,14 @@ dependencies = [ [[package]] name = "common-catalog" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "const_format", ] [[package]] name = "common-config" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-error", @@ -2371,7 +2371,7 @@ dependencies = [ [[package]] name = "common-datasource" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "arrow-schema 59.2.0", @@ -2408,7 +2408,7 @@ dependencies = [ [[package]] name = "common-decimal" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "bigdecimal 0.4.8", "common-error", @@ -2421,7 +2421,7 @@ dependencies = [ [[package]] name = "common-error" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-macro", "http 1.5.0", @@ -2433,7 +2433,7 @@ dependencies = [ [[package]] name = "common-event-recorder" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2458,7 +2458,7 @@ dependencies = [ [[package]] name = "common-frontend" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2477,7 +2477,7 @@ dependencies = [ [[package]] name = "common-function" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -2539,7 +2539,7 @@ dependencies = [ [[package]] name = "common-greptimedb-telemetry" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "common-runtime", @@ -2556,7 +2556,7 @@ dependencies = [ [[package]] name = "common-grpc" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", @@ -2591,7 +2591,7 @@ dependencies = [ [[package]] name = "common-grpc-expr" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "common-catalog", @@ -2609,7 +2609,7 @@ dependencies = [ [[package]] name = "common-macro" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "greptime-proto", "once_cell", @@ -2620,7 +2620,7 @@ dependencies = [ [[package]] name = "common-mem-prof" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "anyhow", "common-error", @@ -2636,7 +2636,7 @@ dependencies = [ [[package]] name = "common-memory-manager" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-error", "common-macro", @@ -2648,7 +2648,7 @@ dependencies = [ [[package]] name = "common-meta" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "anymap2", "api", @@ -2719,7 +2719,7 @@ dependencies = [ [[package]] name = "common-options" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-grpc", @@ -2730,11 +2730,11 @@ dependencies = [ [[package]] name = "common-plugins" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" [[package]] name = "common-pprof" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-error", "common-macro", @@ -2745,7 +2745,7 @@ dependencies = [ [[package]] name = "common-procedure" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-stream", @@ -2774,7 +2774,7 @@ dependencies = [ [[package]] name = "common-procedure-test" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "common-procedure", @@ -2783,7 +2783,7 @@ dependencies = [ [[package]] name = "common-query" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -2813,7 +2813,7 @@ dependencies = [ [[package]] name = "common-recordbatch" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arc-swap", "common-base", @@ -2838,7 +2838,7 @@ dependencies = [ [[package]] name = "common-runtime" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "catio", @@ -2866,7 +2866,7 @@ dependencies = [ [[package]] name = "common-session" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "serde", "strum 0.27.1", @@ -2874,7 +2874,7 @@ dependencies = [ [[package]] name = "common-sql" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow-schema 59.2.0", "common-base", @@ -2893,7 +2893,7 @@ dependencies = [ [[package]] name = "common-stat" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-runtime", @@ -2907,7 +2907,7 @@ dependencies = [ [[package]] name = "common-telemetry" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "backtrace", "common-base", @@ -2937,7 +2937,7 @@ dependencies = [ [[package]] name = "common-test-util" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-grpc", "common-query", @@ -2949,7 +2949,7 @@ dependencies = [ [[package]] name = "common-time" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "chrono", @@ -2967,7 +2967,7 @@ dependencies = [ [[package]] name = "common-version" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "cargo-manifest", "const_format", @@ -2977,7 +2977,7 @@ dependencies = [ [[package]] name = "common-wal" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-base", "common-error", @@ -3000,7 +3000,7 @@ dependencies = [ [[package]] name = "common-workload" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "common-telemetry", "serde", @@ -4473,7 +4473,7 @@ dependencies = [ [[package]] name = "datanode" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", @@ -4546,7 +4546,7 @@ checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" [[package]] name = "datatypes" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arrow 59.2.0", "arrow-array 59.2.0", @@ -5278,7 +5278,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "file-engine" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -5419,7 +5419,7 @@ checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" [[package]] name = "flow" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow 59.2.0", @@ -5471,7 +5471,7 @@ dependencies = [ "snafu 0.8.6", "sql", "store-api", - "substrait 1.3.0-beta.1", + "substrait 1.4.0-alpha.0", "table", "tokio", "tokio-stream", @@ -5568,7 +5568,7 @@ dependencies = [ [[package]] name = "frontend" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", @@ -6946,7 +6946,7 @@ dependencies = [ [[package]] name = "index" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "async-trait", @@ -8097,7 +8097,7 @@ checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "log-query" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "chrono", "common-error", @@ -8109,7 +8109,7 @@ dependencies = [ [[package]] name = "log-store" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-stream", "async-trait", @@ -8461,7 +8461,7 @@ dependencies = [ [[package]] name = "meta-client" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -8490,7 +8490,7 @@ dependencies = [ [[package]] name = "meta-srv" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -8581,7 +8581,7 @@ dependencies = [ [[package]] name = "metric-engine" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", @@ -8681,7 +8681,7 @@ dependencies = [ [[package]] name = "mito-codec" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "bytes", @@ -8706,7 +8706,7 @@ dependencies = [ [[package]] name = "mito2" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", @@ -9549,7 +9549,7 @@ dependencies = [ [[package]] name = "object-store" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "anyhow", "bytes", @@ -10094,7 +10094,7 @@ dependencies = [ [[package]] name = "operator" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -10156,7 +10156,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-beta.1", + "substrait 1.4.0-alpha.0", "table", "tokio", "tokio-util", @@ -10490,7 +10490,7 @@ checksum = "e3c406c9e2aa74554e662d2c2ee11cd3e73756988800be7e6f5eddb16fed4699" [[package]] name = "partition" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "common-base", @@ -10884,7 +10884,7 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pipeline" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -11039,7 +11039,7 @@ dependencies = [ [[package]] name = "plugins" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "auth", "catalog", @@ -11382,7 +11382,7 @@ dependencies = [ [[package]] name = "promql" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "async-trait", @@ -11799,7 +11799,7 @@ dependencies = [ [[package]] name = "puffin" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-compression", "async-trait", @@ -11851,7 +11851,7 @@ dependencies = [ [[package]] name = "query" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -11918,7 +11918,7 @@ dependencies = [ "sql", "sqlparser", "store-api", - "substrait 1.3.0-beta.1", + "substrait 1.4.0-alpha.0", "table", "tokio", "toml 0.8.23", @@ -13533,7 +13533,7 @@ dependencies = [ [[package]] name = "servers" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -13676,7 +13676,7 @@ dependencies = [ [[package]] name = "session" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "ahash 0.8.12", "api", @@ -14036,7 +14036,7 @@ dependencies = [ [[package]] name = "sql" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-buffer 59.2.0", @@ -14090,7 +14090,7 @@ dependencies = [ [[package]] name = "sqlness-runner" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "clap", @@ -14369,7 +14369,7 @@ dependencies = [ [[package]] name = "standalone" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "async-trait", @@ -14419,7 +14419,7 @@ checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "store-api" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "aquamarine", @@ -14611,7 +14611,7 @@ dependencies = [ [[package]] name = "substrait" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "async-trait", "bytes", @@ -14747,7 +14747,7 @@ dependencies = [ [[package]] name = "table" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arc-swap", @@ -15027,7 +15027,7 @@ checksum = "8f50febec83f5ee1df3015341d8bd429f2d1cc62bcba7ea2076759d315084683" [[package]] name = "tests-fuzz" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "arbitrary", "async-trait", @@ -15075,7 +15075,7 @@ dependencies = [ [[package]] name = "tests-integration" -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-flight", @@ -15154,7 +15154,7 @@ dependencies = [ "sqlx", "standalone", "store-api", - "substrait 1.3.0-beta.1", + "substrait 1.4.0-alpha.0", "table", "tempfile", "time", diff --git a/Cargo.toml b/Cargo.toml index f6ca2d66562..cce9b502928 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,7 +76,7 @@ members = [ resolver = "2" [workspace.package] -version = "1.3.0-beta.1" +version = "1.4.0-alpha.0" edition = "2024" license = "Apache-2.0" From acebd6b35be27424385335ba71684f8dda333c2e Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 2 Oct 2026 09:50:08 +0000 Subject: [PATCH 086/133] fix(catalog): skip backend batch_get on full cache hits (#9431) fix(catalog): return cached batch values without an empty backend call Count backend batch_get invocations in regression tests and cover empty input. Signed-off-by: WenyXu --- src/catalog/src/kvbackend/client.rs | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/src/catalog/src/kvbackend/client.rs b/src/catalog/src/kvbackend/client.rs index 69b31d62322..402af1d9865 100644 --- a/src/catalog/src/kvbackend/client.rs +++ b/src/catalog/src/kvbackend/client.rs @@ -192,6 +192,10 @@ impl KvBackend for CachedKvBackend { } } + if miss_keys.is_empty() { + return Ok(BatchGetResponse { kvs }); + } + let batch_get_req = BatchGetRequest::new().with_keys(miss_keys.clone()); let pre_version = self.version(); @@ -503,6 +507,7 @@ mod tests { pub struct SimpleKvBackend { inner_map: DashMap, Vec>, get_execute_times: Arc, + batch_get_execute_times: Arc, } impl TxnService for SimpleKvBackend { @@ -520,6 +525,7 @@ mod tests { } async fn batch_get(&self, req: BatchGetRequest) -> Result { + self.batch_get_execute_times.fetch_add(1, Ordering::SeqCst); let mut kvs = Vec::with_capacity(req.keys.len()); for key in req.keys.iter() { if let Some(kv) = self.get(key).await? { @@ -571,6 +577,7 @@ mod tests { async fn test_cached_kv_backend() { let simple_kv = Arc::new(SimpleKvBackend::default()); let get_execute_times = simple_kv.get_execute_times.clone(); + let batch_get_execute_times = simple_kv.batch_get_execute_times.clone(); let cached_kv = CachedKvBackend::wrap(simple_kv); add_some_vals(&cached_kv).await; @@ -582,9 +589,11 @@ mod tests { assert_eq!(get_execute_times.load(Ordering::SeqCst), 0); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 2); assert_eq!(get_execute_times.load(Ordering::SeqCst), 2); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 1); } let batch_get_req = BatchGetRequest { @@ -594,14 +603,29 @@ mod tests { let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); for _ in 0..10 { - let _batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + let batch_get_resp = cached_kv.batch_get(batch_get_req.clone()).await.unwrap(); + assert_eq!(batch_get_resp.kvs.len(), 3); assert_eq!(get_execute_times.load(Ordering::SeqCst), 3); + assert_eq!(batch_get_execute_times.load(Ordering::SeqCst), 2); } } + #[tokio::test] + async fn test_cached_kv_backend_empty_batch_get() { + let simple_kv = Arc::new(SimpleKvBackend::default()); + let cached_kv = CachedKvBackend::wrap(simple_kv.clone()); + + let response = cached_kv.batch_get(BatchGetRequest::new()).await.unwrap(); + + assert!(response.kvs.is_empty()); + assert_eq!(simple_kv.batch_get_execute_times.load(Ordering::SeqCst), 0); + assert_eq!(simple_kv.get_execute_times.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_cached_kv_backend_rejects_writes_with_read_only_inner() { let inner = Arc::new(MemoryKvBackend::::new()); From a9ce6e2d01877466ab024a355c189864a2efc3dc Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Fri, 2 Oct 2026 10:19:50 +0000 Subject: [PATCH 087/133] ci: make query regression non-blocking for scheduled nightly releases (#9434) * ci: make query regression non-blocking for scheduled nightly releases Scheduled nightly releases no longer gate publishing on the query regression release test. The test still runs (validation policy stays 'all' for automatic releases) and uploads its report, but any result - failure, timeout, cancellation, or skip - no longer blocks image publishing or the GitHub release for schedule-triggered runs. Since reusable-workflow caller jobs cannot use continue-on-error, the gate is relaxed at the consumers instead: the if conditions of release-images-to-dockerhub and publish-github-release now also pass when github.event_name == 'schedule' regardless of the query-regression result. All other gates (runner allocation, artifact builds, prepare-release-validation) are unchanged. Tag-push and manual-dispatch releases remain fully blocking unless a skip policy is explicitly chosen via the release_validation input. Signed-off-by: Ning Sun * ci: preserve downstream nightly release jobs after regression failures Signed-off-by: evenyag --------- Signed-off-by: Ning Sun Signed-off-by: evenyag Co-authored-by: evenyag --- .github/workflows/release.yml | 45 +++++++++++++++++++++++++++++++---- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 5a689910aee..0fb22f23e61 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -84,7 +84,7 @@ on: default: false release_validation: type: choice - description: Manual policy for accepted perf regressions or runner failures; automatic releases always run all validation + description: Manual policy for accepted perf regressions or runner failures; automatic releases always run all validation (scheduled nightlies run query regression non-blocking) required: true default: all options: @@ -468,6 +468,9 @@ jobs: if [[ "${validation_policy}" != "all" ]]; then echo "- âš ï¸ Validation bypassed: compatibility=${run_compat}, query-regression=${run_query_regression}" fi + if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then + echo "- â„¹ï¸ Query-regression failure or cancellation is non-blocking for scheduled nightly releases" + fi } >> "${GITHUB_STEP_SUMMARY}" compat-release: @@ -551,6 +554,9 @@ jobs: release-images-to-dockerhub: name: Build and push images to DockerHub + # Scheduled nightly releases do not gate on query regression: the test + # still runs and reports, but a failure or cancellation does not block + # the nightly release. if: | always() && (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && @@ -562,7 +568,8 @@ jobs: (needs.compat-release.result == 'success' || (needs.prepare-release-validation.outputs.run-compat == 'false' && needs.compat-release.result == 'skipped')) && (needs.query-regression-release.result == 'success' || - (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) + (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped') || + github.event_name == 'schedule') needs: [ allocate-runners, build-linux-amd64-artifacts, @@ -598,7 +605,18 @@ jobs: release-cn-artifacts: name: Release artifacts to CN region - if: ${{ (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') }} + # A non-blocking query regression failure must not skip downstream releases. + if: | + always() && + (inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule') && + (github.event_name != 'schedule' || needs.allocate-runners.outputs.nightly-required == 'true') && + needs.allocate-runners.result == 'success' && + needs.build-linux-amd64-artifacts.result == 'success' && + needs.build-linux-arm64-artifacts.result == 'success' && + needs.build-linux-riscv64-artifacts.result == 'success' && + needs.build-macos-artifacts.result == 'success' && + needs.build-windows-artifacts.result == 'success' && + needs.release-images-to-dockerhub.result == 'success' needs: [ # The job have to wait for all the artifacts are built. allocate-runners, build-linux-amd64-artifacts, @@ -642,6 +660,8 @@ jobs: name: Create GitHub release and upload artifacts # Use always() to run even when optional jobs (macos, windows) are skipped. # Then check that required jobs succeeded and optional jobs didn't fail. + # Scheduled nightly releases do not gate on query regression either + # (see release-images-to-dockerhub); the test still runs and reports. if: | always() && (inputs.publish_github_release || github.event_name == 'push' || github.event_name == 'schedule') && @@ -651,7 +671,8 @@ jobs: (needs.compat-release.result == 'success' || (needs.prepare-release-validation.outputs.run-compat == 'false' && needs.compat-release.result == 'skipped')) && (needs.query-regression-release.result == 'success' || - (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) && + (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped') || + github.event_name == 'schedule') && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && needs.build-linux-riscv64-artifacts.result == 'success' && @@ -754,7 +775,21 @@ jobs: bump-downstream-repo-versions: name: Bump downstream repo versions - if: ${{ (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true') || ((github.event_name == 'push' || github.event_name == 'workflow_dispatch') && github.ref_type == 'tag' && !contains(github.ref_name, 'nightly')) && (needs.allocate-runners.outputs.is-current-version-stable == 'true' && needs.allocate-runners.outputs.is-current-version-latest == 'true') }} + # A non-blocking query regression failure must not skip downstream updates. + if: | + always() && + needs.allocate-runners.result == 'success' && + needs.publish-github-release.result == 'success' && + ( + (github.event_name == 'schedule' && needs.allocate-runners.outputs.nightly-required == 'true') || + ( + (github.event_name == 'push' || github.event_name == 'workflow_dispatch') && + github.ref_type == 'tag' && + !contains(github.ref_name, 'nightly') && + needs.allocate-runners.outputs.is-current-version-stable == 'true' && + needs.allocate-runners.outputs.is-current-version-latest == 'true' + ) + ) needs: [allocate-runners, publish-github-release] runs-on: ubuntu-latest # Permission reference: https://docs.github.com/en/actions/using-jobs/assigning-permissions-to-jobs From 3a4c24a863d3628d847b4aa054eae5ab78960350 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 2 Oct 2026 12:09:16 +0000 Subject: [PATCH 088/133] fix: temporarily disable v2 series scan by default (#9436) * fix: disable v2 series scan by default temporarily Signed-off-by: evenyag * test: explicitly enable v2 series scan in sqlness configs Signed-off-by: evenyag * test: regenerate sqlness results for legacy series scan default Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 4 +- config/datanode.example.toml | 2 +- config/standalone.example.toml | 2 +- src/mito2/src/config.rs | 11 +--- tests-integration/tests/http.rs | 2 +- .../promql/tsid_binary_join_regression.result | 56 +++++++++---------- .../tsid_histogram_quantile_regression.result | 2 +- .../tql-explain-analyze/tsid_column.result | 12 ++-- 8 files changed, 42 insertions(+), 49 deletions(-) diff --git a/config/config.md b/config/config.md index 24351b2b6fd..1a7e8af1fe3 100644 --- a/config/config.md +++ b/config/config.md @@ -217,7 +217,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

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

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | @@ -687,7 +687,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

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

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 3b8f226abb2..627a5d3288a 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -687,7 +687,7 @@ schedule_compaction_after_edit = true default_flat_format = true ## Whether to enable the experimental two-phase mode for eligible metric series scans. -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 4760e038f49..6b60c4482d7 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -891,7 +891,7 @@ schedule_compaction_after_edit = true default_flat_format = true ## Whether to enable the experimental two-phase mode for eligible metric series scans. -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 78b79aec81a..c0fa74281ca 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -262,7 +262,8 @@ impl Default for MitoConfig { min_compaction_interval: Duration::from_secs(0), schedule_compaction_after_edit: true, default_flat_format: true, - experimental_series_scan_v2: true, + // FIXME(#9435): Keep v2 opt-in while long-range scan memory usage is investigated. + experimental_series_scan_v2: false, gc: GcConfig::default(), }; @@ -498,14 +499,6 @@ mod tests { assert_eq!(config, restored); } } - - #[test] - fn test_experimental_series_scan_v2_config() { - assert!(MitoConfig::default().experimental_series_scan_v2); - - let config: MitoConfig = toml::from_str("experimental_series_scan_v2 = false").unwrap(); - assert!(!config.experimental_series_scan_v2); - } } /// Index build mode. diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index b1cfeac1ba1..eff6a8685f2 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2984,7 +2984,7 @@ scan_memory_on_exhausted = "fail" min_compaction_interval = "0s" schedule_compaction_after_edit = true default_flat_format = true -experimental_series_scan_v2 = true +experimental_series_scan_v2 = false [region_engine.mito.index] aux_path = "" diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 54fd5ad5e26..fc4bc6f01c2 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -127,13 +127,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -167,13 +167,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -210,19 +210,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -262,13 +262,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -308,13 +308,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -347,13 +347,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -386,13 +386,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -430,19 +430,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -485,25 +485,25 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -541,7 +541,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -549,13 +549,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -597,19 +597,19 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index e50fa43f10e..10b31cac1e1 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -118,7 +118,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, le@4 as le, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index c761267bd5a..3b509cdd20b 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -54,7 +54,7 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -82,7 +82,7 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, instance@3 as instance, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -120,7 +120,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -133,7 +133,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -172,7 +172,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -185,7 +185,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ From 97913b4b3a51e233d406b337ebf25b17809c5824 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 7 Oct 2026 07:18:25 +0000 Subject: [PATCH 089/133] fix: coalesce equal histogram quantile bucket boundaries (#9451) * fix: coalesce equal histogram quantile bucket boundaries Merge numerically equivalent bucket boundaries before monotonicity repair, preserving the allocation-free path for unique boundaries. Share coalescing with histogram_fraction and cover duplicate boundaries with unit and SQL regressions. Fixes #9443 Signed-off-by: evenyag * fix: keep duplicate infinity buckets in one histogram group Confirm classic histogram completion at a series/timestamp boundary or EOF instead of the first positive infinity. Retain lookahead across batches and reject incomplete optimistic groups so equivalent infinity labels reach numeric coalescing together. Add stream and SQL regressions for duplicate infinity labels, including established bucket counts and EOF. Correct older fixtures to use distinct keys for distinct histogram groups. Signed-off-by: evenyag * fix: avoid buffering histogram groups with invalid trailing bounds Discover candidate bucket counts at the first group boundary and let optimistic validation fall back to safe mode for malformed bounds. Document discovery and cover trailing invalid and null labels with a focused streaming regression. Signed-off-by: evenyag --------- Signed-off-by: evenyag --- .../src/extension_plan/histogram_fold.rs | 334 +++++++++++++++--- .../common/promql/simple_histogram.result | 89 +++++ .../common/promql/simple_histogram.sql | 47 +++ 3 files changed, 424 insertions(+), 46 deletions(-) diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 7fdab474f5d..3dc51e27bcd 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -846,6 +846,8 @@ impl HistogramFoldStream { self.find_first_complete_bucket(&batch) } + /// Returns the first group's candidate bucket count once a group boundary is + /// observed. Optimistic validation checks its bounds before folding. fn find_first_complete_bucket(&self, batch: &RecordBatch) -> DataFusionResult> { if batch.num_rows() == 0 { return Ok(None); @@ -853,21 +855,14 @@ impl HistogramFoldStream { let vectors = Helper::try_into_vectors(batch.columns()) .map_err(|e| DataFusionError::Execution(e.to_string()))?; - let le_array = batch.column(self.le_column_index); - let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); self.collect_tag_values(&vectors, 0, &mut tag_values_buf); - let mut group_start = 0usize; - for row in 0..batch.num_rows() { + for row in 1..batch.num_rows() { if !self.is_same_group(&vectors, row, &tag_values_buf) { - // new group begins - self.collect_tag_values(&vectors, row, &mut tag_values_buf); - group_start = row; - } - - if Self::is_positive_infinity(le_array, row) { - return Ok(Some(row - group_start + 1)); + // A new series/timestamp proves that the previous group is + // complete, including any numerically equal +Inf boundaries. + return Ok(Some(row)); } } @@ -888,7 +883,10 @@ impl HistogramFoldStream { let field_array = field_array.as_primitive::(); let mut tag_values_buf = Vec::with_capacity(self.normal_indices.len()); - while remaining_rows >= bucket_num && self.mode == FoldMode::Optimistic { + // Keep one row of lookahead: +Inf alone does not prove completion. + // A group ending at the batch boundary stays buffered until more input + // arrives, or flush_remaining completes it at EOF. + while remaining_rows > bucket_num && self.mode == FoldMode::Optimistic { self.collect_tag_values(&vectors, cursor, &mut tag_values_buf); if !self.validate_optimistic_group( &vectors, @@ -987,6 +985,9 @@ impl HistogramFoldStream { tag_values: &[ValueRef<'_>], ) -> bool { let inf_index = cursor + bucket_num - 1; + if self.is_same_group(vectors, inf_index + 1, tag_values) { + return false; + } if !Self::is_positive_infinity(le_array, inf_index) { return false; } @@ -1274,8 +1275,31 @@ impl HistogramFoldStream { } } - // check input value - if !bucket.windows(2).all(|w| w[0] <= w[1]) { + // Detect equal numeric boundaries during validation so unique buckets + // need neither another scan nor an allocation. + let has_duplicates = !bucket.windows(2).all(|w| w[0] < w[1]); + if has_duplicates && !bucket.windows(2).all(|w| w[0] <= w[1]) { + return Ok(f64::NAN); + } + // Counts must be coalesced before repairing monotonicity. + let coalesced = has_duplicates.then(|| { + let mut bounds = Vec::with_capacity(bucket.len()); + let mut counts = Vec::with_capacity(counter.len()); + for (&bound, &count) in bucket.iter().zip(counter) { + if bounds.last() == Some(&bound) { + *counts.last_mut().unwrap() += count; + } else { + bounds.push(bound); + counts.push(count); + } + } + (bounds, counts) + }); + let (bucket, counter) = match &coalesced { + Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), + None => (bucket, counter), + }; + if matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() < 2 { return Ok(f64::NAN); } let counter = match operation { @@ -1339,26 +1363,6 @@ impl HistogramFoldStream { } fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 { - let coalesced = bucket - .windows(2) - .any(|bounds| bounds[0] == bounds[1]) - .then(|| { - let mut bounds = Vec::with_capacity(bucket.len()); - let mut counts = Vec::with_capacity(counter.len()); - for (&bound, &count) in bucket.iter().zip(counter) { - if bounds.last() == Some(&bound) { - *counts.last_mut().unwrap() += count; - } else { - bounds.push(bound); - counts.push(count); - } - } - (bounds, counts) - }); - let (bucket, counter) = match &coalesced { - Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), - None => (bucket, counter), - }; let total = *counter.last().unwrap(); if total == 0.0 || lower.is_nan() || upper.is_nan() { return f64::NAN; @@ -1466,8 +1470,8 @@ mod test { // 12 items let host_column_1 = Arc::new(StringArray::from(vec![ - "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", "host_1", - "host_1", "host_1", "host_1", "host_1", + "host_1", "host_1", "host_1", "host_1", "host_1", "host_2", "host_2", "host_2", + "host_2", "host_2", "host_3", "host_3", ])) as _; let le_column_1 = Arc::new(StringArray::from(vec![ "0.001", "0.1", "10", "1000", "+Inf", "0.001", "0.1", "10", "1000", "+inf", "0.001", @@ -1478,14 +1482,14 @@ mod test { ])) as _; // 2 items - let host_column_2 = Arc::new(StringArray::from(vec!["host_1", "host_1"])) as _; + let host_column_2 = Arc::new(StringArray::from(vec!["host_3", "host_3"])) as _; let le_column_2 = Arc::new(StringArray::from(vec!["10", "1000"])) as _; let val_column_2 = Arc::new(Float64Array::from(vec![1.0, 1.0])) as _; // 11 items let host_column_3 = Arc::new(StringArray::from(vec![ - "host_1", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", "host_2", - "host_2", "host_2", "host_2", + "host_3", "host_4", "host_4", "host_4", "host_4", "host_4", "host_5", "host_5", + "host_5", "host_5", "host_5", ])) as _; let le_column_3 = Arc::new(StringArray::from(vec![ "+INF", "0.001", "0.1", "10", "1000", "+iNf", "0.001", "0.1", "10", "1000", "+Inf", @@ -1651,10 +1655,10 @@ mod test { | host | val | +--------+-------------------+ | host_1 | 257.5 | -| host_1 | 5.05 | -| host_1 | 0.0004 | -| host_2 | NaN | -| host_2 | 6.040000000000001 | +| host_2 | 5.05 | +| host_3 | 0.0004 | +| host_4 | NaN | +| host_5 | 6.040000000000001 | +--------+-------------------+", ); assert_eq!(result_literal, expected); @@ -1831,7 +1835,7 @@ mod test { Field::new("le", DataType::Utf8, true), Field::new("val", DataType::Float64, true), ])); - let host_column = Arc::new(StringArray::from(vec!["a", "a", "a", "a", "b", "b"])) as _; + let host_column = Arc::new(StringArray::from(vec!["a", "a", "b", "b", "c", "c"])) as _; let le_column = Arc::new(StringArray::from(vec![ "0.1", "+Inf", "0.1", "1.0", "0.1", "+Inf", ])) as _; @@ -1852,8 +1856,8 @@ mod test { | host | val | +------+-----+ | a | 0.1 | -| a | NaN | -| b | 0.1 | +| b | NaN | +| c | 0.1 | +------+-----+", ); assert_eq!(result_literal, expected); @@ -1930,6 +1934,79 @@ mod test { ); } + #[tokio::test] + async fn trailing_invalid_bounds_fall_back_before_eof() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("le", DataType::Utf8, true), + Field::new("val", DataType::Float64, false), + ])); + for trailing_bound in [Some("bad"), None] { + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "a", "a", "a", "b", "b", "b", "c", "c", "c", + ])), + Arc::new(StringArray::from( + [Some("1"), Some("+Inf"), trailing_bound].repeat(3), + )), + Arc::new(Float64Array::from([2.0, 4.0, 99.0].repeat(3))), + ], + ) + .unwrap(); + let fold = build_fold_exec_from_batches(vec![batch.clone()], schema.clone(), 0.5, 0); + let mut stream = HistogramFoldStream { + le_column_index: 1, + field_column_index: 2, + histogram_column_index: None, + operation: fold.operation, + normal_indices: vec![0], + bucket_size: None, + batch_size: 1, + output_schema: fold.output_schema.clone(), + input_schema: schema.clone(), + mode: FoldMode::Optimistic, + safe_group: None, + input_buffer: vec![], + input_buffered_rows: 0, + output_buffer: HistogramFoldStream::empty_output_buffer(&schema).unwrap(), + output_buffered_rows: 0, + input: fold + .input + .execute(0, SessionContext::default().task_ctx()) + .unwrap(), + metric: BaselineMetrics::new(&fold.metric, 0), + }; + + // Neither +Inf nor a batch boundary alone completes the group. + assert!(stream.fold_input(batch.slice(0, 2)).unwrap().is_none()); + assert!(stream.fold_input(batch.slice(2, 1)).unwrap().is_none()); + assert_eq!(stream.mode, FoldMode::Optimistic); + + for (offset, host) in [(3, "a"), (6, "b")] { + let output = stream + .fold_input(batch.slice(offset, 3)) + .unwrap() + .expect("a completed group should produce output before EOF") + .unwrap(); + assert_eq!(stream.mode, FoldMode::Safe); + assert_eq!(stream.input_buffered_rows, 0); + assert!(stream.input_buffer.is_empty()); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some(host)); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + } + + stream.flush_remaining().unwrap(); + let output = stream.take_output_buf().unwrap().unwrap(); + assert_eq!(output.num_rows(), 1); + assert_eq!(string_array_value_at_index(output.column(0), 0), Some("c")); + assert_eq!(output.column(1).as_primitive::().value(0), 1.0); + assert!(stream.take_output_buf().unwrap().is_none()); + } + } + #[tokio::test] async fn safe_mode_handles_misaligned_groups() { let schema = Arc::new(Schema::new(vec![ @@ -2045,6 +2122,171 @@ mod test { assert!((values[1] - 0.55).abs() < 1e-10, "{values:?}"); } + #[test] + fn evaluate_quantile_duplicate_bounds() { + let cases = [ + // Issue #9443: coalesce before repairing the count at 1000. + ( + vec![10.0, 10.0, 100.0, 100.0, 1000.0, f64::INFINITY], + vec![40.0, 40.0, 50.0, 50.0, 50.0, 100.0], + 0.95, + 77.5, + ), + // Repairing the split counts first would incorrectly inflate them. + ( + vec![10.0, 10.0, 100.0, 100.0, f64::INFINITY], + vec![30.0, 10.0, 40.0, 20.0, 60.0], + 0.5, + 7.5, + ), + ( + vec![10.0, f64::INFINITY, f64::INFINITY], + vec![40.0, 20.0, 30.0], + 0.5, + 6.25, + ), + ( + vec![f64::INFINITY, f64::INFINITY], + vec![20.0, 30.0], + 0.5, + f64::NAN, + ), + ( + vec![-0.0, 0.0, 10.0, f64::INFINITY], + vec![10.0, 10.0, 40.0, 40.0], + 0.75, + 5.0, + ), + ]; + for (bounds, counts, quantile, expected) in cases { + let actual = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(quantile.into()), + &bounds, + &counts, + ) + .unwrap(); + if expected.is_nan() { + assert!(actual.is_nan()); + } else { + assert!((actual - expected).abs() < 1e-10, "{actual} != {expected}"); + } + } + } + + #[tokio::test] + async fn duplicate_infinity_bounds_across_batches() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + let duplicate = [("10", 20.0), ("10.0", 20.0), ("+Inf", 25.0), ("+inf", 25.0)]; + let regular = [("10", 20.0), ("10.0", 20.0), ("+Inf", 50.0)]; + // Cover discovery, an already established optimistic bucket count, + // fallback, and the final group at EOF. + for groups in [ + vec![duplicate.as_slice()], + vec![duplicate.as_slice(), regular.as_slice()], + vec![regular.as_slice(), duplicate.as_slice(), regular.as_slice()], + ] { + let mut timestamps = Vec::new(); + let mut bounds = Vec::new(); + let mut counts = Vec::new(); + for (timestamp, group) in groups.iter().enumerate() { + for &(bound, count) in *group { + timestamps.push(timestamp as i64); + bounds.push(bound); + counts.push(count); + } + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(StringArray::from(bounds)), + Arc::new(Float64Array::from(counts)), + ], + ) + .unwrap(); + for batch_size in 1..=batch.num_rows() { + let batches = (0..batch.num_rows()) + .step_by(batch_size) + .map(|offset| batch.slice(offset, batch_size.min(batch.num_rows() - offset))) + .collect(); + let fold = build_fold_exec_from_batches(batches, schema.clone(), 0.5, 0); + let result = + datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![6.25; groups.len()], "batch size {batch_size}"); + } + } + } + + #[tokio::test] + async fn duplicate_bounds_across_batches_and_safe_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + // The first group establishes six buckets for optimistic folding. The + // shorter second group forces safe mode; both have duplicate bounds. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 1000, 1000, 1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000, 2000, + ])), + Arc::new(StringArray::from(vec![ + "10", "10.0", "100", "100.0", "1000", "+Inf", "10", "10.0", "100", "100.0", + "+Inf", + ])), + Arc::new(Float64Array::from(vec![ + 40.0, 40.0, 50.0, 50.0, 50.0, 100.0, 40.0, 40.0, 50.0, 50.0, 100.0, + ])), + ], + ) + .unwrap(); + // Split each pair of equal boundaries across an input batch boundary. + let batches = (0..batch.num_rows()).map(|i| batch.slice(i, 1)).collect(); + let fold = build_fold_exec_from_batches(batches, schema, 0.95, 0); + let result = datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values: Vec<_> = result + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![77.5, 77.5]); + } + #[test] fn evaluate_row_normal_case() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY]; diff --git a/tests/cases/standalone/common/promql/simple_histogram.result b/tests/cases/standalone/common/promql/simple_histogram.result index a794a30e8b9..95cc6b6aa30 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.result +++ b/tests/cases/standalone/common/promql/simple_histogram.result @@ -436,3 +436,92 @@ drop table histogram5_bucket; Affected Rows: 0 +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +Affected Rows: 0 + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +Affected Rows: 7 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + ++--------+---------------------+-----------------------------------+ +| le | ts | sum(equivalent_bounds_bucket.val) | ++--------+---------------------+-----------------------------------+ +| +Inf | 1970-01-01T00:50:00 | 100.0 | +| 10 | 1970-01-01T00:50:00 | 40.0 | +| 10.0 | 1970-01-01T00:50:00 | 40.0 | +| 100 | 1970-01-01T00:50:00 | 50.0 | +| 100.0 | 1970-01-01T00:50:00 | 50.0 | +| 1000.0 | 1970-01-01T00:50:00 | 50.0 | ++--------+---------------------+-----------------------------------+ + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + ++---------------------+-----------------------------------+ +| ts | sum(equivalent_bounds_bucket.val) | ++---------------------+-----------------------------------+ +| 1970-01-01T00:50:00 | 77.5 | ++---------------------+-----------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + ++----------+---------------------+-----------------------------------+ +| instance | ts | sum(equivalent_bounds_bucket.val) | ++----------+---------------------+-----------------------------------+ +| a | 1970-01-01T00:50:00 | 77.5 | +| b | 1970-01-01T00:50:00 | 77.5 | ++----------+---------------------+-----------------------------------+ + +DROP TABLE equivalent_bounds_bucket; + +Affected Rows: 0 + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +Affected Rows: 0 + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +Affected Rows: 4 + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + ++---------------------+------+ +| ts | val | ++---------------------+------+ +| 1970-01-01T00:50:00 | 6.25 | ++---------------------+------+ + +DROP TABLE infinity_bounds_bucket; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/simple_histogram.sql b/tests/cases/standalone/common/promql/simple_histogram.sql index ec0bf9b1439..6f73325a61f 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.sql +++ b/tests/cases/standalone/common/promql/simple_histogram.sql @@ -253,3 +253,50 @@ insert into histogram5_bucket values tql eval (3000, 3015, '3s') histogram_quantile(0.5, histogram5_bucket); drop table histogram5_bucket; + +-- Issue #9443: string labels stay distinct during aggregation, but equal +-- numeric bucket boundaries must be coalesced before monotonicity repair. +CREATE TABLE equivalent_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY (instance, le) +); + +INSERT INTO equivalent_bounds_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') sum by (le) (equivalent_bounds_bucket); + +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (le) (equivalent_bounds_bucket)); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') histogram_quantile(0.95, sum by (instance, le) (equivalent_bounds_bucket)); + +DROP TABLE equivalent_bounds_bucket; + +-- Equivalent infinity labels belong to the same histogram until the group ends. +CREATE TABLE infinity_bounds_bucket ( + ts TIMESTAMP TIME INDEX, + le STRING, + val DOUBLE, + PRIMARY KEY (le) +); + +INSERT INTO infinity_bounds_bucket VALUES + (3000000, '10', 20), + (3000000, '10.0', 20), + (3000000, '+Inf', 25), + (3000000, '+inf', 25); + +tql eval (3000, 3000, '1s') histogram_quantile(0.5, infinity_bounds_bucket); + +DROP TABLE infinity_bounds_bucket; From beed56c305bbf907b02e664339e5b03cbe8990ad Mon Sep 17 00:00:00 2001 From: yihong Date: Thu, 8 Oct 2026 02:21:45 +0000 Subject: [PATCH 090/133] chore: update jieba-rs (#9453) Signed-off-by: yihong0618 --- Cargo.lock | 53 ++++------------------- src/index/Cargo.toml | 4 +- src/index/src/fulltext_index/tokenizer.rs | 24 +++++----- 3 files changed, 23 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e985f483efb..6ec2a2d2764 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1774,15 +1774,6 @@ dependencies = [ "shlex 2.0.1", ] -[[package]] -name = "cedarwood" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0524a528a6a0288df1863c3c20fe92c301875b4941e7b6c4b394ab08c5a4c55" -dependencies = [ - "smallvec", -] - [[package]] name = "census" version = "0.4.2" @@ -1927,7 +1918,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7d8d1efd5109b9c1cd3b7966bd071cdfb53bb6eb0b22a473a68c2f70a11a1eb" dependencies = [ "parse-zoneinfo", - "phf_codegen 0.12.1", + "phf_codegen", "phf_shared 0.12.1", "uncased", ] @@ -7343,25 +7334,19 @@ dependencies = [ [[package]] name = "jieba-macros" -version = "0.10.1" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46adade69b634535a8f495cf87710ed893cff53e1dbc9dd750c2ab81c5defb82" -dependencies = [ - "phf_codegen 0.13.1", -] +checksum = "2987cf45b1bfdb4ff1f6a814b1aa2ad6361f4ff43bf353a9bd29f391a9568f70" [[package]] name = "jieba-rs" -version = "0.10.1" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11b53580aaa8ec8b713da271da434f8947409242c537a9ab3f7b76bdbb19e8a9" +checksum = "0a3a7d02ea31e6fc6f2fe10495dc7c5efcf381e1f21d7172f58ba792efa1c2fd" dependencies = [ "bytecount", - "cedarwood", "include-flate", "jieba-macros", - "phf 0.13.1", - "regex", "rustc-hash", ] @@ -10774,20 +10759,10 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "efbdcb6f01d193b17f0b9c3360fa7e0e620991b193ff08702f78b3ce365d7e61" dependencies = [ - "phf_generator 0.12.1", + "phf_generator", "phf_shared 0.12.1", ] -[[package]] -name = "phf_codegen" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49aa7f9d80421bca176ca8dbfebe668cc7a2684708594ec9f3c0db0805d5d6e1" -dependencies = [ - "phf_generator 0.13.1", - "phf_shared 0.13.1", -] - [[package]] name = "phf_generator" version = "0.12.1" @@ -10798,23 +10773,13 @@ dependencies = [ "phf_shared 0.12.1", ] -[[package]] -name = "phf_generator" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737" -dependencies = [ - "fastrand", - "phf_shared 0.13.1", -] - [[package]] name = "phf_macros" version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d713258393a82f091ead52047ca779d37e5766226d009de21696c4e667044368" dependencies = [ - "phf_generator 0.12.1", + "phf_generator", "phf_shared 0.12.1", "proc-macro2", "quote", @@ -14895,9 +14860,9 @@ dependencies = [ [[package]] name = "tantivy-jieba" -version = "0.20.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3392170e86f1c387170aba7d171a466ffdc98a8b55b006e19ac64b123a7b690a" +checksum = "25cf6c5eb9419d7da9aacf5e184863b9ed476c478de0e196729f16124e716dd1" dependencies = [ "jieba-rs", "lazy_static", diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index 078f9e42c56..2458a6591af 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -23,7 +23,7 @@ fst.workspace = true futures.workspace = true greptime-proto.workspace = true itertools.workspace = true -jieba-rs = "0.10" +jieba-rs = "0.11" lazy_static.workspace = true mockall.workspace = true pin-project.workspace = true @@ -37,7 +37,7 @@ serde_json.workspace = true snafu.workspace = true store-api.workspace = true tantivy = { version = "0.26", features = ["zstd-compression"] } -tantivy-jieba = "0.20" +tantivy-jieba = "0.21" tokio.workspace = true tokio-util.workspace = true uuid.workspace = true diff --git a/src/index/src/fulltext_index/tokenizer.rs b/src/index/src/fulltext_index/tokenizer.rs index e6752a40420..0a1cc47aa8e 100644 --- a/src/index/src/fulltext_index/tokenizer.rs +++ b/src/index/src/fulltext_index/tokenizer.rs @@ -262,6 +262,8 @@ mod tests { let text = "[2026/04/09/ 13:56:11.031]2026-04-09 13:56:11.031 - [ trace_id=340a6a44b0bd8e37bb7697ss7da61ff0 span_id=085ff5ttf1e0a23b trace_flags=01] - [http-nio-8081-exec-16] INFO c.h.p.xx.web.service.impl.CCCXForwardKKKServiceImpl.pushout(188) - 登录手机å·18888888888的动æ€key：829889AC8 ship_ship ship__ship _ __ __IDENTIFIER__ _ship ship_ EOF"; let tokens = tokenizer.tokenize(text); + // Jieba 0.11 preserves ASCII compounds and emits searchable alphabetic parts. + // The underscore fallback also retains leading, trailing, and repeated underscores. assert_eq!( tokens, vec![ @@ -271,36 +273,34 @@ mod tests { "13", "56", "11.031", - "2026-04", - "09", + "2026-04-09", "13", "56", "11.031", "trace", - "_", "id", + "trace_id", "340a6a44b0bd8e37bb7697ss7da61ff0", "span", - "_", "id", + "span_id", "085ff5ttf1e0a23b", "trace", - "_", "flags", + "trace_flags", "01", "http", - "nio-8081", - "exec-16", + "nio", + "exec", + "http-nio-8081-exec-16", "INFO", - "c", - "h", - "p", "xx", "web", "service", "impl", "CCCXForwardKKKServiceImpl", "pushout", + "c.h.p.xx.web.service.impl.CCCXForwardKKKServiceImpl.pushout", "188", "登录", "手机", @@ -311,8 +311,8 @@ mod tests { "key", "829889AC8", "ship", - "_", "ship", + "ship_ship", "ship", "__", "ship", @@ -367,8 +367,8 @@ mod tests { "_", "id", "trace", - "_", "id", + "trace_id", "手机", "手机å·", "_", From b78dc4b33bf49210c38782a2879be6f042353cb2 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 02:30:37 +0000 Subject: [PATCH 091/133] test(sqlness): wait for mirrored inserts before flushing streaming flows (#9424) Signed-off-by: Dennis Zhuang --- tests/cases/standalone/common/flow/flow_null.result | 4 ++++ tests/cases/standalone/common/flow/flow_null.sql | 4 ++++ .../standalone/common/flow/show_create_flow.result | 12 ++++++++++++ .../standalone/common/flow/show_create_flow.sql | 12 ++++++++++++ 4 files changed, 32 insertions(+) diff --git a/tests/cases/standalone/common/flow/flow_null.result b/tests/cases/standalone/common/flow/flow_null.result index 2386d676cd5..a141b9e03ab 100644 --- a/tests/cases/standalone/common/flow/flow_null.result +++ b/tests/cases/standalone/common/flow/flow_null.result @@ -261,6 +261,8 @@ VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); @@ -290,6 +292,8 @@ VALUES Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); diff --git a/tests/cases/standalone/common/flow/flow_null.sql b/tests/cases/standalone/common/flow/flow_null.sql index 4cc71628092..6c1219df2db 100644 --- a/tests/cases/standalone/common/flow/flow_null.sql +++ b/tests/cases/standalone/common/flow/flow_null.sql @@ -164,6 +164,8 @@ VALUES (2, "name2", "2024-10-18 19:00:00"), (3, "name3", "2024-10-18 19:00:00"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); @@ -177,6 +179,8 @@ VALUES ("name2", "2024-10-18 19:00:00"), ("name3", "2024-10-18 19:00:00"); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('calc_nullable_pk'); diff --git a/tests/cases/standalone/common/flow/show_create_flow.result b/tests/cases/standalone/common/flow/show_create_flow.result index f529cef1b9f..d451c552ec3 100644 --- a/tests/cases/standalone/common/flow/show_create_flow.result +++ b/tests/cases/standalone/common/flow/show_create_flow.result @@ -116,6 +116,8 @@ INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -155,6 +157,8 @@ INSERT INTO numbers_input_show VALUES (4,4),(5,4),(10, 3),(11, 4); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -195,6 +199,8 @@ INSERT INTO numbers_input_show VALUES (3, 1),(4, 2),(10, 3),(11, 4); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -324,6 +330,8 @@ INSERT INTO numbers_input_show VALUES(-4,0), (-3,1), (-2,2), (-1,3); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -384,6 +392,8 @@ INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); Affected Rows: 3 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -426,6 +436,8 @@ INSERT INTO numbers_input_show VALUES (10, 6),(11, 8),(15, 7),(18, 3); Affected Rows: 4 +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); diff --git a/tests/cases/standalone/common/flow/show_create_flow.sql b/tests/cases/standalone/common/flow/show_create_flow.sql index 9856a6041f8..a2bcc4c06d9 100644 --- a/tests/cases/standalone/common/flow/show_create_flow.sql +++ b/tests/cases/standalone/common/flow/show_create_flow.sql @@ -45,6 +45,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's not replaced in flownode INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -59,6 +61,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's not replaced in flownode INSERT INTO numbers_input_show VALUES (4,4),(5,4),(10, 3),(11, 4); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -73,6 +77,8 @@ SELECT flow_name, table_catalog, flow_definition, source_table_names FROM INFORM -- makesure it's replaced in flownode INSERT INTO numbers_input_show VALUES (3, 1),(4, 2),(10, 3),(11, 4); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -118,6 +124,8 @@ SELECT * FROM out_num_cnt_show; INSERT INTO numbers_input_show VALUES(-4,0), (-3,1), (-2,2), (-1,3); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -145,6 +153,8 @@ CREATE FLOW filter_numbers_show SINK TO out_num_cnt_show AS SELECT number as n1 INSERT INTO numbers_input_show VALUES (10, 0),(15, 1),(16, 2); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); @@ -160,6 +170,8 @@ SELECT flow_definition, source_table_names FROM INFORMATION_SCHEMA.FLOWS WHERE f INSERT INTO numbers_input_show VALUES (10, 6),(11, 8),(15, 7),(18, 3); +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s -- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | ADMIN FLUSH_FLOW('filter_numbers_show'); From 8f062976b619058e5d373b557e794cf33dfcda66 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 02:40:26 +0000 Subject: [PATCH 092/133] docs: point vulnerability reports to GitHub private reporting (#9459) --- SECURITY.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/SECURITY.md b/SECURITY.md index ebb11e93c55..72029c7beb7 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -14,6 +14,9 @@ and cloud platform. If you come across a security vulnerability in GreptimeDB, we kindly request that you inform us immediately. We will thoroughly investigate all valid reports and make every effort to resolve the issue promptly. -To report any issues or vulnerabilities, please email us at info@greptime.com, rather than -posting publicly on GitHub. Be sure to provide us with the version identifier as well as details +Please do not report security vulnerabilities through public GitHub issues, discussions, or pull requests. + +Report them privately through GitHub's [private vulnerability reporting](https://github.com/GreptimeTeam/greptimedb/security/advisories/new) +(the "Report a vulnerability" button on the repository's Security tab). If you cannot use GitHub, +email info@greptime.com instead. Be sure to provide us with the version identifier as well as details on how the vulnerability can be exploited. From 329edc88ecd1f1946b8c167e595f100a1e4ab122 Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Thu, 8 Oct 2026 03:35:14 +0000 Subject: [PATCH 093/133] fix: accept a bool value for the CLI --auto-create-schema flag (#9450) * fix: accept a value for the CLI --auto-create-schema flag clap derives ArgAction::SetTrue for bool fields, so --auto-create-schema took no value and could only set true, which is already the default. Users could not turn automatic schema creation off from the CLI, unlike metasrv's auto_create_schema = false. Set action = clap::ArgAction::Set so the flag accepts an explicit true/false value while keeping true as the default. Signed-off-by: fengyue-xve <1759519137@qq.com> * chore: enhance CLI argument for auto_create_schema option Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Signed-off-by: fengyue-xve <1759519137@qq.com> Co-authored-by: dennis zhuang Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/cli/src/common/store.rs | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/cli/src/common/store.rs b/src/cli/src/common/store.rs index e974037ea4d..c3b8e477ed1 100644 --- a/src/cli/src/common/store.rs +++ b/src/cli/src/common/store.rs @@ -86,7 +86,13 @@ pub struct StoreConfig { /// Automatically create PostgreSQL schema if it doesn't exist (default: true). #[cfg(feature = "pg_kvbackend")] - #[clap(long, default_value_t = true)] + #[clap( + long, + default_value_t = true, + default_missing_value = "true", + num_args = 0..=1, + action = clap::ArgAction::Set + )] pub auto_create_schema: bool, /// TLS mode for backend store connections (etcd, PostgreSQL, MySQL) @@ -223,3 +229,27 @@ impl StoreConfig { } } } + +#[cfg(all(test, feature = "pg_kvbackend"))] +mod tests { + use clap::Parser; + + use super::StoreConfig; + + #[test] + fn test_auto_create_schema_accepts_explicit_value() { + let config = + StoreConfig::try_parse_from(["store", "--auto-create-schema", "false"]).unwrap(); + assert!(!config.auto_create_schema); + + let config = + StoreConfig::try_parse_from(["store", "--auto-create-schema", "true"]).unwrap(); + assert!(config.auto_create_schema); + } + + #[test] + fn test_auto_create_schema_defaults_to_true() { + let config = StoreConfig::try_parse_from(["store"]).unwrap(); + assert!(config.auto_create_schema); + } +} From d640ab8bc58760fd26310ab4d9da78ba713138f8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 04:11:04 +0000 Subject: [PATCH 094/133] fix(otlp): coerce empty and "-" trace attribute strings to null (#9461) * fix(otlp): coerce empty and "-" trace attribute strings to null Instrumentations fill numeric attributes from absent headers with an empty string (PSR-7 getHeaderLine returns ""). Coercing that into an existing numeric column failed and rejected the whole span. Signed-off-by: Dennis Zhuang * test: trim redundant row from empty-string rewrite test Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp/trace_types.rs | 37 ++++++++++--------- src/servers/src/otlp/coerce.rs | 28 ++++++++++++++ 2 files changed, 48 insertions(+), 17 deletions(-) diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 0295d12e685..6f37f14e34a 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -519,31 +519,34 @@ mod tests { ); } - /// The PHP instrumentation case: an empty string must be distinguishable - /// from any other unparsable value in the reported diagnostics. + /// The PHP instrumentation case: `getHeaderLine('Content-Length')` sends "" + /// when the header is absent. The span must be kept with a null attribute. #[test] - fn test_prepare_trace_column_rewrites_reports_empty_string_value() { - let rows = vec![Row { - values: vec![Value { - value_data: Some(ValueData::StringValue(String::new())), + fn test_prepare_trace_column_rewrites_nulls_empty_string_value() { + let mut rows = Rows { + schema: vec![ColumnSchema { + column_name: "span_attributes.http.response.body.size".to_string(), + datatype: ColumnDataType::String as i32, + ..Default::default() }], - }]; + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue(String::new())), + }], + }], + }; let pending_rewrites = vec![PendingTraceColumnRewrite { col_idx: 0, target_type: ColumnDataType::Int64, column_name: "span_attributes.http.response.body.size".to_string(), }]; - let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces") - .unwrap_err(); - assert!( - err.error.to_string().contains( - "'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \ - from String(\"\") to Int64" - ), - "unexpected error message: {}", - err.error - ); + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "opentelemetry_traces") + .unwrap() + .apply(&mut rows); + + assert_eq!(rows.schema[0].datatype, ColumnDataType::Int64 as i32); + assert_eq!(rows.rows[0].values[0].value_data, None); } #[test] diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index 2e7dd36ed7a..2f8d25247cf 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -30,6 +30,9 @@ pub enum TraceCoerceError { // - String to Int64 // - String to Float64 // - String to Boolean +// The empty string and "-" coerce to null instead of failing: instrumentations +// fill numeric attributes from absent headers this way (PSR-7 `getHeaderLine` +// returns ""), and rejecting the value would reject the whole span. // // Lossless signed-to-unsigned integer casts. These let the built-in data // models move from unsigned to signed integers while existing unsigned tables @@ -86,6 +89,13 @@ pub fn coerce_value_data( return Ok(None); }; + if let ValueData::StringValue(s) = v + && matches!(s.as_str(), "" | "-") + && is_supported_trace_coercion(request_type, target) + { + return Ok(None); + } + let Some(value) = coerce_non_null_value(target, request_type, v) else { return Err(TraceCoerceError::Unsupported); }; @@ -273,6 +283,24 @@ mod tests { assert_eq!(result, Err(TraceCoerceError::Unsupported)); } + #[test] + fn test_coerce_absent_value_strings_to_null() { + for s in ["", "-"] { + for target in [ + ColumnDataType::Int64, + ColumnDataType::Float64, + ColumnDataType::Boolean, + ] { + let result = coerce_value_data( + &Some(ValueData::StringValue(s.to_string())), + target, + ColumnDataType::String, + ); + assert_eq!(result, Ok(None), "{s:?} to {target:?}"); + } + } + } + #[test] fn test_coerce_float64_to_int64_not_supported() { let result = coerce_value_data( From 48a2ca63409dbdb7d46fef038af91baf61fbc008 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Thu, 8 Oct 2026 06:09:14 +0000 Subject: [PATCH 095/133] refactor: expose OTLP write and batching helpers for downstream routers (#9401) Signed-off-by: shuiyisong --- src/frontend/AGENTS.md | 4 +- src/frontend/src/instance/logical_batcher.rs | 34 +++--- src/frontend/src/instance/otlp.rs | 119 +++++++++++-------- src/frontend/src/lib.rs | 2 + src/frontend/src/server.rs | 8 +- src/servers/src/http.rs | 2 +- 6 files changed, 99 insertions(+), 70 deletions(-) diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 20c7bc323a5..e16cbfc9872 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -57,7 +57,9 @@ remote datanodes via `operator`/`client`. - **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. + writes. Downstream routers can initialize it when enabling replacement endpoints. + `Instance::handle_otlp_metric_row_inserts` shares OTLP eligibility and dispatch + between converters and falls back for incompatible tables. The schema adapter holds a weak instance reference to avoid an ownership cycle. - **Table batching** (`instance/builder.rs`): protocol entry points opt in through diff --git a/src/frontend/src/instance/logical_batcher.rs b/src/frontend/src/instance/logical_batcher.rs index e68dcf644d4..255766ac14d 100644 --- a/src/frontend/src/instance/logical_batcher.rs +++ b/src/frontend/src/instance/logical_batcher.rs @@ -26,23 +26,25 @@ use crate::frontend::FrontendOptions; use crate::instance::Instance; impl Instance { - pub(crate) fn init_logical_batcher(self: &Arc, options: &FrontendOptions) { + /// Initializes the shared logical batcher for enabled ingestion endpoints. + /// Disabled endpoints leave it uninitialized so downstream routers can enable it. + pub fn init_logical_batcher(self: &Arc, options: &FrontendOptions) { + let options_batcher = options.logical_batcher_options(); + let enabled = options_batcher + .protocols + .iter() + .any(|protocol| match protocol { + BatchingProtocol::Prom => options.prom_store.enable, + BatchingProtocol::Otlp => options.otlp.enable, + _ => false, + }); + if !options.prom_store.with_metric_engine + || !enabled + || !options_batcher.pending_rows_batching_enabled() + { + return; + } self.logical_batcher.get_or_init(|| { - 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(), diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 5132b9f8a8d..5d96138319a 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -18,6 +18,7 @@ pub mod trace_types; use std::sync::Arc; +use api::v1::RowInsertRequests; use async_trait::async_trait; use auth::{ OTLP_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, @@ -43,6 +44,7 @@ use servers::query_handler::{ MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; +use session::protocol_ctx::OtlpMetricCtx; use snafu::ResultExt; use table::requests::{ OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY, @@ -166,57 +168,9 @@ impl OpenTelemetryProtocolHandler for Instance { Arc::new(c) }; - let physical_table = ctx - .extension(PHYSICAL_TABLE_PARAM) - .unwrap_or(GREPTIME_PHYSICAL_TABLE) - .to_string(); - // The bulk path converts each request's time index unit to the - // destination table's unit during batch alignment, so no pre-gate - // alignment is needed here. - let batcher = self.logical_batcher().filter(|_| { - ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine - }); - let batcher = if batcher.is_some() - && self - .inserter - .can_batch_metric_rows(&requests, &ctx, &physical_table) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu)? - { - batcher - } else { - None - }; - - // OTLP tables have one sample field in both the legacy and physical paths. - 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)? - } else { - self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu)? - }; + let output = self + .handle_otlp_metric_row_inserts(requests, ctx.clone(), &metric_ctx) + .await?; outcome.write_cost = output.meta.cost; // Derived enrichment follows the accepted metric submission, which may @@ -362,6 +316,69 @@ impl OpenTelemetryProtocolHandler for Instance { } } +impl Instance { + /// Inserts converted OTLP metrics through the eligible batching or ordinary write path. + /// The caller must check table permissions and admit the complete request first. + pub async fn handle_otlp_metric_row_inserts( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + metric_ctx: &OtlpMetricCtx, + ) -> ServerResult { + let physical_table = ctx + .extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .to_string(); + // The bulk path converts each request's time index unit to the + // destination table's unit during batch alignment, so no pre-gate + // alignment is needed here. + let batcher = self.logical_batcher().filter(|_| { + ctx.logical_batching_enabled() && !metric_ctx.is_legacy && metric_ctx.with_metric_engine + }); + let batcher = if batcher.is_some() + && self + .inserter + .can_batch_metric_rows(&requests, &ctx, &physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu)? + { + batcher + } else { + None + }; + + // OTLP tables have one sample field in both the legacy and physical paths. + if let Some(batcher) = batcher { + let (rows, cost) = batcher + .submit_with(requests, ctx.clone(), |mut requests| { + let ctx = ctx.clone(); + async move { + Inserter::meter_row_inserts(&mut requests, &ctx) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + }) + .await?; + Ok(Output::new( + OutputData::AffectedRows(rows as usize), + OutputMeta::new_with_cost(cost as _), + )) + } else if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + self.handle_row_inserts(requests, ctx.clone(), false, true) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } else { + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/frontend/src/lib.rs b/src/frontend/src/lib.rs index 1468fe7c0b2..7c68e445242 100644 --- a/src/frontend/src/lib.rs +++ b/src/frontend/src/lib.rs @@ -23,3 +23,5 @@ pub(crate) mod metrics; pub mod server; pub mod service_config; mod stream_wrapper; + +pub use metrics::{OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index cf278f1f1d0..8de7bfe340f 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -455,7 +455,8 @@ fn effective_prom_store_options(opts: &FrontendOptions) -> PromStoreOptions { prom_store } -fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { +/// Resolves HTTP options, allowing synchronous batches enough time to flush. +pub fn effective_http_options(opts: &FrontendOptions) -> HttpOptions { effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled()) } @@ -555,6 +556,11 @@ mod tests { .await .unwrap(), ); + let mut disabled_options = options.clone(); + disabled_options.otlp.enable = false; + instance.init_logical_batcher(&disabled_options); + assert!(instance.logical_batcher().is_none()); + let services = Services::new(options.clone(), instance.clone(), Plugins::default()); let batcher = instance.logical_batcher().unwrap().clone(); instance.init_logical_batcher(&options); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index e7cb715fa40..34f214bae26 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -225,7 +225,7 @@ pub fn is_api_listener_path(path: &str) -> bool { struct LogicalBatchingProtocols(Vec); /// Sets a local-only write selector after authentication creates the context. -async fn set_http_write_batching( +pub async fn set_http_write_batching( State(protocol): State, mut req: Request, next: Next, From f9ab3c296b9711c59bb6ffc9d3753a318d7883d7 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 8 Oct 2026 06:33:40 +0000 Subject: [PATCH 096/133] feat(log-store): skip corrupted object store WAL segments per region (#9425) * feat(log-store): skip corrupted object store WAL segments per region A segment that does not decode, because its checksum does not match or its content disagrees with its footer entry, is fetched once more. If it still does not decode, the read follows `on_corrupted_segment`: `skip` records a WAL hole for the region with the object key, sequence and the footer's entry id range, increments `greptime_logstore_object_store_wal_skipped_segments_total`, logs a warning and continues with the next object; `fail` fails the read. A failed fetch fails the read in both modes. Holes are kept in memory per region and can be listed from the store. The standalone test builder allocates WAL options through `build_standalone_wal_provider`, and integration tests run a standalone instance on the object store WAL across restarts, on S3 when the `GT_S3_*` variables are set. The standalone shutdown runs every step and returns the first error. Signed-off-by: jeremyhi * test(cmd): cover the best-effort standalone stop Build a standalone instance whose leader services fail to stop and whose datanode servers fail afterwards, and check that `Instance::stop` still stops the object store WAL and returns the leader services error. Rename `skip_segment` to `record_skipped_segment` and the file store integration test to `test_standalone_object_store_wal_survives_restarts_on_file`. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/cmd/Cargo.toml | 1 + src/cmd/src/standalone.rs | 140 ++++++- src/datanode/src/datanode.rs | 7 + src/log-store/src/metrics.rs | 8 + src/log-store/src/object_store_wal.rs | 6 +- src/log-store/src/object_store_wal/store.rs | 370 ++++++++++++++++-- .../engine/object_store_wal_recovery_test.rs | 111 +++++- tests-integration/src/standalone.rs | 29 +- tests-integration/src/tests/test_util.rs | 14 +- tests-integration/tests/main.rs | 1 + tests-integration/tests/object_store_wal.rs | 188 +++++++++ 11 files changed, 786 insertions(+), 89 deletions(-) create mode 100644 tests-integration/tests/object_store_wal.rs diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 2c925de57c8..0ef2f35b371 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -125,6 +125,7 @@ api.workspace = true client = { workspace = true, features = ["testing"] } common-test-util.workspace = true common-version.workspace = true +datanode = { workspace = true, features = ["testing"] } file-engine.workspace = true metric-engine.workspace = true mito2.workspace = true diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index bc6a62cd4d5..76c95c26d32 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -48,8 +48,8 @@ use common_meta::wal_provider::{WalProvider, WalProviderRef, build_wal_provider} use common_options::plugin_options::StandaloneFlag; use common_procedure::ProcedureManagerRef; use common_query::prelude::set_default_prefix; -use common_telemetry::info; use common_telemetry::logging::{DEFAULT_LOGGING_DIR, TracingOptions}; +use common_telemetry::{info, warn}; use common_time::timezone::set_default_timezone; use common_version::{short_version, verbose_version}; use common_wal::config::DatanodeWalConfig; @@ -274,32 +274,53 @@ impl App for Instance { Ok(()) } + /// Stops every component in order. A failed step does not skip the later + /// ones; the first error is returned and the rest are logged. async fn stop(&mut self) -> Result<()> { - self.frontend - .shutdown() - .await - .context(error::ShutdownFrontendSnafu)?; + let mut first_error = None; + let mut record = |result: Result<()>| { + if let Err(err) = result { + if first_error.is_none() { + first_error = Some(err); + } else { + warn!(err; "Ignored a later shutdown error"); + } + } + }; - self.leader_services_controller - .stop( - self.procedure_manager.clone(), - self.datanode.region_server(), - ) - .await?; + record( + self.frontend + .shutdown() + .await + .context(error::ShutdownFrontendSnafu), + ); - self.datanode - .shutdown() - .await - .context(error::ShutdownDatanodeSnafu)?; + record( + self.leader_services_controller + .stop( + self.procedure_manager.clone(), + self.datanode.region_server(), + ) + .await, + ); - self.flownode - .shutdown() - .await - .context(error::ShutdownFlownodeSnafu)?; + record( + self.datanode + .shutdown() + .await + .context(error::ShutdownDatanodeSnafu), + ); + + record( + self.flownode + .shutdown() + .await + .context(error::ShutdownFlownodeSnafu), + ); info!("Datanode instance stopped."); - Ok(()) + first_error.map_or(Ok(()), Err) } } @@ -1067,10 +1088,89 @@ mod tests { use frontend::frontend::FrontendOptions; use object_store::config::{FileConfig, GcsConfig}; use servers::grpc::GrpcOptions; + use servers::server::ServerHandlers; + use store_api::logstore::LogStore; use super::*; use crate::options::GlobalOptions; + struct FailingLeaderServices; + + #[async_trait] + impl StandaloneLeaderServicesController for FailingLeaderServices { + async fn start(&self, _context: LeaderServicesContext) -> Result<()> { + Ok(()) + } + + async fn stop( + &self, + _procedure_manager: ProcedureManagerRef, + _region_server: RegionServer, + ) -> Result<()> { + error::IllegalConfigSnafu { + msg: "leader services failed to stop", + } + .fail() + } + } + + struct FailingServer; + + #[async_trait] + impl servers::server::Server for FailingServer { + async fn shutdown(&self) -> servers::error::Result<()> { + servers::error::InternalSnafu { + err_msg: "server failed to stop", + } + .fail() + } + + async fn start(&mut self, _listening: SocketAddr) -> servers::error::Result<()> { + Ok(()) + } + + fn name(&self) -> &str { + "failing" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + #[tokio::test] + async fn test_stop_runs_every_step_and_returns_the_first_error() { + common_telemetry::init_default_ut_logging(); + let data_home = create_temp_dir("standalone-stop"); + let mut opts = StandaloneOptions { + wal: DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()), + ..Default::default() + }; + opts.storage.data_home = data_home.path().to_str().unwrap().to_string(); + let creator = InstanceCreator::default() + .with_leader_services_controller(Box::new(FailingLeaderServices)); + let (mut instance, _) = StartCommand::build_with(opts, vec![], creator) + .await + .unwrap(); + let log_store = instance.datanode.object_store_log_store().unwrap(); + // The datanode fails to stop its servers after the leader services failed. + let mut services = ServerHandlers::default(); + services.insert((Box::new(FailingServer), "127.0.0.1:0".parse().unwrap())); + services.start_all().await.unwrap(); + instance.datanode.setup_services(services); + + let err = instance.stop().await.unwrap_err(); + + assert!( + matches!(&err, error::Error::IllegalConfig { msg, .. } if msg == "leader services failed to stop"), + "unexpected error: {err:?}" + ); + assert!(matches!( + log_store.append_batch(vec![]).await, + Err(log_store::error::Error::ObjectStoreWalStopped { .. }) + )); + } + #[tokio::test] async fn test_build_standalone_wal_provider() { let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 8a80083aef3..fdac0597572 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -179,6 +179,13 @@ impl Datanode { self.region_server.clone() } + /// Returns the object store log store when the datanode runs on the object + /// store WAL. + #[cfg(any(test, feature = "testing"))] + pub fn object_store_log_store(&self) -> Option> { + self.object_store_log_store.clone() + } + pub fn plugins(&self) -> Plugins { self.plugins.clone() } diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index f12ad3e815a..824058d70af 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -89,4 +89,12 @@ lazy_static! { &[LOGSTORE_LABEL, PARTITION_LABEL], ) .unwrap(); + + /// Counter of segments a read of the object store logstore skipped because + /// they did not decode. + pub static ref METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_skipped_segments_total", + "object store logstore skipped corrupted segments total", + ) + .unwrap(); } diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index 2ec36b30463..d4bfc096393 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -50,7 +50,9 @@ //! Recovery lists the objects, reads and verifies only the header, trailer and //! footer of each, and indexes the footers of the chain in sequence order to //! rebuild the object catalog, which rejects a sequence it already holds. -//! Segments are read and checksummed only when a read decodes them. +//! Segments are read and checksummed only when a read decodes them; a segment +//! that does not decode is fetched once more and then skipped as a hole of its +//! region or fails the read, as `on_corrupted_segment` says. //! //! Entry ids are object-sequence-major, see [`entry_id`]: the high bits of an //! id name the object that holds the entry, the low bits its position among @@ -63,4 +65,4 @@ mod io; mod store; pub use batch::entry_id; -pub use store::ObjectStoreLogStore; +pub use store::{ObjectStoreLogStore, WalHole}; diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 5a6454123dd..8750a61cc4b 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -23,8 +23,8 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use bytes::Bytes; -use common_telemetry::info; -use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig}; +use common_telemetry::{info, warn}; +use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; use futures::future::BoxFuture; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt}; @@ -46,6 +46,7 @@ use crate::error::{ StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, WalObjectSequenceUnsettledSnafu, }; +use crate::metrics::METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL; use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ @@ -87,13 +88,17 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// indexed, so an acknowledged entry never has a missing predecessor; in the /// `enqueued` mode it returns on admission and the object is created in the /// background. While [`MAX_SEALED_BATCHES`] batches -/// wait to become durable no append is admitted. +/// wait to become durable no append is admitted. A segment that a read cannot +/// decode is handled as `on_corrupted_segment` says. pub struct ObjectStoreLogStore { prefix: String, ack_mode: AckMode, + on_corrupted_segment: CorruptedSegmentAction, io: Arc, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, + /// Segments that reads skipped, per region and by object sequence. + wal_holes: WalHoles, terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, @@ -114,6 +119,20 @@ pub struct ObjectStoreLogStore { type ObsoleteEntryIds = Arc>>; +type WalHoles = Arc>>>; + +/// A segment a read skipped because it did not decode: the entries of one +/// region in one object are missing from every read of that region. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WalHole { + /// Key of the object that holds the segment. + pub path: String, + pub object_seq: u64, + /// Entry id range the footer records for the segment. + pub min_entry_id: EntryId, + pub max_entry_id: EntryId, +} + type TerminalError = Arc>>>; impl fmt::Debug for ObjectStoreLogStore { @@ -121,6 +140,7 @@ impl fmt::Debug for ObjectStoreLogStore { f.debug_struct("ObjectStoreLogStore") .field("prefix", &self.prefix) .field("ack_mode", &self.ack_mode) + .field("on_corrupted_segment", &self.on_corrupted_segment) .finish_non_exhaustive() } } @@ -260,9 +280,11 @@ impl ObjectStoreLogStore { Ok(Arc::new(Self { prefix, ack_mode: config.ack_mode, + on_corrupted_segment: config.on_corrupted_segment, io, catalog, obsolete_entry_ids, + wal_holes: WalHoles::default(), terminal_error, stopped, command_tx, @@ -297,6 +319,21 @@ impl ObjectStoreLogStore { catalog.region_max_entry_id(region_id).unwrap_or(0) } + /// Returns the segments of the provider's region that reads of this store + /// skipped because they did not decode, ordered by object sequence. Holes + /// are kept in memory only: a store opened later on the same prefix + /// records them again when its reads meet the segments. + pub fn wal_holes(&self, provider: &Provider) -> Result> { + let region_id = self.region_of(provider)?; + Ok(self + .wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(®ion_id) + .map(|holes| holes.values().cloned().collect()) + .unwrap_or_default()) + } + /// Returns the region of `provider`, which must select this store's prefix. fn region_of(&self, provider: &Provider) -> Result { let provider = @@ -477,6 +514,30 @@ impl ObjectStoreLogStore { pub fn begin_stop(&self) { self.stopped.store(true, Ordering::Release); } + + /// Returns the key of the object `object_seq` and the byte range the + /// segment of the provider's region occupies in it, so a test can damage + /// that segment alone, or `None` when the object holds no segment of the + /// region. + pub fn segment_location( + &self, + provider: &Provider, + object_seq: u64, + ) -> Result)>> { + let region_id = self.region_of(provider)?; + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + Ok(catalog + .objects_for_entry_range(region_id, 0, EntryId::MAX)? + .into_iter() + .find(|(seq, _)| *seq == object_seq) + .map(|(_, entry)| { + let start = entry.segment_offset; + ( + self.io.object_path(object_seq), + start..start + entry.segment_len, + ) + })) + } } #[async_trait::async_trait] @@ -534,6 +595,9 @@ impl LogStore for ObjectStoreLogStore { /// Reads the provider's region from `entry_id`, hiding obsolete entries. /// The catalog locates each segment, so a caller-supplied index is unnecessary. + /// A segment that does not decode fails the read or is skipped and recorded + /// as a hole, as `on_corrupted_segment` says; a failed fetch always fails + /// the read. async fn read( &self, provider: &Provider, @@ -568,15 +632,20 @@ impl LogStore for ObjectStoreLogStore { let io = self.io.clone(); let provider = provider.clone(); + let on_corrupted_segment = self.on_corrupted_segment; + let wal_holes = self.wal_holes.clone(); Ok(Box::pin(try_stream! { for (object_seq, footer_entry) in objects { - let bytes = io - .get_range(object_seq, footer_entry.segment_offset, footer_entry.segment_len) - .await?; - let records = decode_segment(&bytes, &footer_entry) - .with_context(|_| InvalidWalObjectSnafu { - path: io.object_path(object_seq), - })?; + let records = match fetch_segment(io.as_ref(), object_seq, &footer_entry).await { + Ok(records) => records, + Err(error @ Error::InvalidWalObject { .. }) + if on_corrupted_segment == CorruptedSegmentAction::Skip => + { + record_skipped_segment(&wal_holes, io.as_ref(), object_seq, &footer_entry, &error); + continue; + } + Err(error) => Err(error)?, + }; let entries = records .into_iter() .filter(|record| record.entry_id >= start_entry_id) @@ -1646,6 +1715,59 @@ fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) } +/// Fetches and decodes the segment `entry` describes in the object +/// `object_seq`. A segment that does not decode is fetched once more, since +/// the download rather than the object may be damaged; when it still does not +/// decode the error is returned as [`Error::InvalidWalObject`]. A failed fetch +/// is returned as is. +async fn fetch_segment( + io: &dyn WalObjectIo, + object_seq: u64, + entry: &FooterEntry, +) -> Result> { + let fetch = || io.get_range(object_seq, entry.segment_offset, entry.segment_len); + if let Ok(records) = decode_segment(&fetch().await?, entry) { + return Ok(records); + } + decode_segment(&fetch().await?, entry).with_context(|_| InvalidWalObjectSnafu { + path: io.object_path(object_seq), + }) +} + +/// Records the segment `entry` describes in the object `object_seq` as a hole +/// of its region, once per object, counts the skip and warns about it. +fn record_skipped_segment( + wal_holes: &WalHoles, + io: &dyn WalObjectIo, + object_seq: u64, + entry: &FooterEntry, + error: &Error, +) { + let hole = WalHole { + path: io.object_path(object_seq), + object_seq, + min_entry_id: entry.min_entry_id, + max_entry_id: entry.max_entry_id, + }; + warn!( + error; + "Skipped a corrupted WAL segment of region {}, object {} (sequence {}), entry ids {}..={}", + entry.region_id, + hole.path, + object_seq, + hole.min_entry_id, + hole.max_entry_id + ); + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.inc(); + wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(entry.region_id) + .or_default() + .entry(object_seq) + .or_insert(hole); +} + fn encode_batch(header: Header, entries: Vec) -> Result { let records = entries .into_iter() @@ -2071,6 +2193,10 @@ mod tests { const PREFIX: &str = "wal/datanodes/1/epochs/2"; const WAIT: Duration = Duration::from_secs(30); + /// Held by every test that reads the process-wide skipped segments + /// counter, so that tests running in one process do not interleave + /// their increments. + static SKIPPED_SEGMENTS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); fn memory_store() -> ObjectStore { ObjectStore::new(Memory::default()).unwrap() @@ -2093,6 +2219,16 @@ mod tests { } } + fn on_corruption( + on_corrupted_segment: CorruptedSegmentAction, + config: ObjectStoreWalConfig, + ) -> ObjectStoreWalConfig { + ObjectStoreWalConfig { + on_corrupted_segment, + ..config + } + } + /// Every append reaches the size limit, so it is persisted on its own. fn eager() -> ObjectStoreWalConfig { config(Duration::from_secs(3600), 1) @@ -2472,6 +2608,18 @@ mod tests { (trailer, footer) } + /// Asserts that recovery over `object_store` rejects the object at `path` + /// for `reason`, whatever a read does with a corrupted segment. + async fn assert_recovery_rejects(object_store: &ObjectStore, path: &str, reason: &str) { + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let config = on_corruption(on_corrupted_segment, eager()); + let error = ObjectStoreLogStore::try_new(object_store.clone(), &config, 1, 2) + .await + .unwrap_err(); + assert_invalid_object(&error, path, reason); + } + } + fn assert_invalid_object(error: &Error, path: &str, reason: &str) { match error { Error::InvalidWalObject { @@ -2553,10 +2701,7 @@ mod tests { let path = object_path(&object_store, 1); corrupt_object(&object_store, &path, *corrupt).await; - let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object(&error, &path, reason); + assert_recovery_rejects(&object_store, &path, reason).await; } } @@ -2575,14 +2720,12 @@ mod tests { .unwrap(); io.put_if_absent(5, encoded.bytes).await.unwrap(); - let error = ObjectStoreLogStore::try_new(object_store, &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object( - &error, + assert_recovery_rejects( + &object_store, &io.object_path(5), "header sequence 0 does not match key sequence 5", - ); + ) + .await; } /// Overwrites the byte range of footer entry `index` and refreshes the @@ -2635,10 +2778,7 @@ mod tests { }) .await; - let error = ObjectStoreLogStore::try_new(object_store.clone(), &eager(), 1, 2) - .await - .unwrap_err(); - assert_invalid_object(&error, &path, reason); + assert_recovery_rejects(&object_store, &path, reason).await; let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); let error = recover_by_decoding(&io).await.unwrap_err(); assert_invalid_object(&error, &path, reason); @@ -2872,9 +3012,11 @@ mod tests { let store = ObjectStoreLogStore { prefix: PREFIX.to_string(), ack_mode: AckMode::Durable, + on_corrupted_segment: CorruptedSegmentAction::Skip, io: Arc::new(ObjectStoreIo::new(memory_store(), PREFIX).unwrap()), catalog: Arc::default(), obsolete_entry_ids: ObsoleteEntryIds::default(), + wal_holes: WalHoles::default(), terminal_error: Arc::default(), stopped: Arc::new(AtomicBool::new(false)), command_tx, @@ -3267,11 +3409,12 @@ mod tests { store.stop().await.unwrap(); } - #[tokio::test] - async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { - let object_store = memory_store(); + /// Writes object 0 with one entry of region 1 and one of region 2 and + /// object 1 with an entry of region 2, then flips a byte in the segment of + /// region 2 in object 0. Returns the key of object 0 and its footer. + async fn write_corrupted_segment(object_store: &ObjectStore) -> (String, Vec) { put_records( - &object_store, + object_store, 0, &[ Record { @@ -3287,19 +3430,93 @@ mod tests { ], ) .await; - put_object(&object_store, 1, region(2), &[10]).await; - let path = object_path(&object_store, 0); + put_object(object_store, 1, region(2), &[10]).await; + let path = object_path(object_store, 0); let bytes = object_store.read(&path).await.unwrap().to_vec(); let footer = footer_of(&bytes).1; let corrupt = &footer[1]; - corrupt_object(&object_store, &path, |bytes| { + assert_eq!(region(2), corrupt.region_id); + corrupt_object(object_store, &path, |bytes| { bytes[(corrupt.segment_offset + corrupt.segment_len - 1) as usize] ^= 1; }) .await; + (path, footer) + } + + #[tokio::test] + async fn test_store_skips_a_corrupted_segment_and_records_a_hole() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let object_store = memory_store(); + let (path, footer) = write_corrupted_segment(&object_store).await; let (io, reads) = RecordingIo::over(object_store); - let store = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) - .await - .unwrap(); + let store = open_over(io, &eager()).await; + assert_eq!(CorruptedSegmentAction::Skip, store.on_corrupted_segment); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // Recovery indexes the object; a read of the other region and one of + // the later object of the damaged region are unaffected. + assert_eq!( + expected_entries(region(1), &[(1, "a1")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 10).await + ); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + + // A read across the damaged segment fetches it twice, continues with + // the next object and records the hole. + reads.lock().unwrap().clear(); + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 0).await + ); + let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); + assert_eq!(corrupt, reads.lock().unwrap()[0]); + assert_eq!(corrupt, reads.lock().unwrap()[1]); + let hole = WalHole { + path, + object_seq: 0, + min_entry_id: 1, + max_entry_id: 1, + }; + assert_eq!( + vec![hole.clone()], + store.wal_holes(&provider(region(2))).unwrap() + ); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped + 1, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + + // Every read across the segment skips it again; the hole is one. + assert_eq!( + expected_entries(region(2), &[(10, "e10")]), + read_entries(&store, region(2), 0).await + ); + assert_eq!(vec![hole], store.wal_holes(&provider(region(2))).unwrap()); + assert_eq!( + skipped + 2, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let object_store = memory_store(); + let (path, footer) = write_corrupted_segment(&object_store).await; + let (io, reads) = RecordingIo::over(object_store); + let config = on_corruption(CorruptedSegmentAction::Fail, eager()); + let store = open_over(io, &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); assert_eq!(1, latest(&store, region(1))); assert_eq!(10, latest(&store, region(2))); reads.lock().unwrap().clear(); @@ -3328,13 +3545,78 @@ mod tests { &path, &format!("segment of region {} checksum mismatch", region(2)), ); + let corrupt = (0, footer[1].segment_offset, footer[1].segment_len); + assert_eq!(vec![corrupt, corrupt], *reads.lock().unwrap()); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); assert_eq!( - vec![(0, corrupt.segment_offset, corrupt.segment_len)], - *reads.lock().unwrap() + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() ); store.stop().await.unwrap(); } + #[tokio::test] + async fn test_store_reads_a_segment_whose_first_download_was_damaged() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1, 2]).await; + let (io, reads) = RecordingIo::over(object_store); + let config = on_corruption(on_corrupted_segment, eager()); + let store = open_over(io.clone(), &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // The first download of the segment is damaged, the second is + // not: the read sees the entries and records no hole. + reads.lock().unwrap().clear(); + io.damage_next_range_read.store(true, Ordering::Relaxed); + assert_eq!( + expected_entries(region(1), &[(1, "e1"), (2, "e2")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!(2, reads.lock().unwrap().len()); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_failed_segment_fetch_fails_the_read() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { + let object_store = memory_store(); + put_object(&object_store, 0, region(1), &[1]).await; + let path = object_path(&object_store, 0); + let (io, _) = RecordingIo::over(object_store.clone()); + let config = on_corruption(on_corrupted_segment, eager()); + let store = open_over(io, &config).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + object_store.delete(&path).await.unwrap(); + let error = store + .read(&provider(region(1)), 0, None) + .await + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { path: actual, .. } if *actual == path), + "unexpected error: {error:?}" + ); + assert!(store.wal_holes(&provider(region(1))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + #[tokio::test] async fn test_store_assigns_ids_from_the_object_sequence() { let store = open(memory_store(), &manual()).await; @@ -5540,13 +5822,14 @@ mod tests { /// Object access that records every range read as (sequence, offset, /// length) and, on request, reports the next conditional create as failed - /// after it wrote the object, or fails it with a given error before it - /// writes. + /// after it wrote the object, fails it with a given error before it + /// writes, or damages the bytes the next range read returns. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, fail_after_next_put: AtomicBool, fail_next_put: Mutex>, + damage_next_range_read: AtomicBool, } impl RecordingIo { @@ -5557,6 +5840,7 @@ mod tests { reads: reads.clone(), fail_after_next_put: AtomicBool::new(false), fail_next_put: Mutex::new(None), + damage_next_range_read: AtomicBool::new(false), }; (Arc::new(io), reads) } @@ -5601,7 +5885,15 @@ mod tests { async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result { self.reads.lock().unwrap().push((object_seq, offset, len)); - self.inner.get_range(object_seq, offset, len).await + let bytes = self.inner.get_range(object_seq, offset, len).await?; + if self.damage_next_range_read.swap(false, Ordering::Relaxed) { + let mut damaged = bytes.to_vec(); + if let Some(last) = damaged.last_mut() { + *last ^= 1; + } + return Ok(Bytes::from(damaged)); + } + Ok(bytes) } async fn list(&self) -> Result> { diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index 6017f745d2e..917ad614a07 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -30,10 +30,12 @@ use common_base::readable_size::ReadableSize; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_recordbatch::RecordBatches; -use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig, STANDALONE_GENERATION}; +use common_wal::config::object_store::{ + AckMode, CorruptedSegmentAction, ObjectStoreWalConfig, STANDALONE_GENERATION, +}; use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use log_store::ObjectStoreLogStore; -use log_store::object_store_wal::entry_id; +use log_store::object_store_wal::{WalHole, entry_id}; use object_store::ObjectStore; use object_store::services::Memory; use rstest::rstest; @@ -412,6 +414,111 @@ async fn test_reopen_after_crash_replays_durable_entries_once(#[case] ack_mode: assert_eq!(3, engine.get_region_statistic(REGION_A).unwrap().num_rows); } +#[rstest] +#[case(CorruptedSegmentAction::Skip)] +#[case(CorruptedSegmentAction::Fail)] +#[tokio::test] +async fn test_reopen_with_a_corrupted_segment( + #[case] on_corrupted_segment: CorruptedSegmentAction, +) { + let mut env = TestEnv::with_prefix("object-store-wal-corrupted-segment").await; + let object_store = memory_store(); + let store = open_store(&object_store, AckMode::Durable).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Objects 1 and 2 each hold one entry of both regions; nothing is flushed. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 2, 4)), + (REGION_B, rows(&schema_b, 3, 5)), + ], + ) + .await; + let rows_b = scan_rows(&engine, REGION_B).await; + engine.stop().await.unwrap(); + + // The segment of region A in object 1 is damaged. + let (path, segment) = store + .segment_location(&provider(REGION_A), 1) + .unwrap() + .unwrap(); + let mut bytes = object_store.read(&path).await.unwrap().to_vec(); + bytes[segment.end as usize - 1] ^= 1; + object_store.write(&path, bytes).await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + let config = ObjectStoreWalConfig { + on_corrupted_segment, + ..store_config(AckMode::Durable) + }; + let store = + ObjectStoreLogStore::try_new(object_store.clone(), &config, 0, STANDALONE_GENERATION) + .await + .unwrap(); + let engine = new_engine(&mut env, store.clone()).await; + + // Region B replays both of its entries whatever the damage does to A. + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); + assert!(store.wal_holes(&provider(REGION_B)).unwrap().is_empty()); + + let opened = open_region(&engine, REGION_A, &table_dir_a, &[]).await; + match on_corrupted_segment { + // Region A replays the entry of object 2 alone and records the + // segment of object 1 as a hole. + CorruptedSegmentAction::Skip => { + opened.unwrap(); + assert_eq!( + RegionRecoveryState { + last_entry_id: entry_id(2, 1), + memtable_rows: 2, + ..EMPTY_RECOVERY_STATE + }, + region_recovery_state(&engine, REGION_A).await + ); + assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!( + vec![WalHole { + path, + object_seq: 1, + min_entry_id: entry_id(1, 1), + max_entry_id: entry_id(1, 1), + }], + store.wal_holes(&provider(REGION_A)).unwrap() + ); + } + // The read fails, so the region does not open. + CorruptedSegmentAction::Fail => { + let message = opened.unwrap_err().output_msg(); + assert!( + message.contains("checksum mismatch"), + "unexpected error: {message}" + ); + assert!(!engine.is_region_exists(REGION_A)); + assert!(store.wal_holes(&provider(REGION_A)).unwrap().is_empty()); + } + } +} + #[tokio::test] async fn test_open_region_rejects_mismatched_wal_prefix() { let mut env = TestEnv::with_prefix("object-store-wal-prefix-mismatch").await; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 3fa1d394489..5a36d495df8 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -39,14 +39,13 @@ use common_meta::procedure_executor::LocalProcedureExecutor; use common_meta::region_keeper::MemoryRegionKeeper; use common_meta::region_registry::LeaderRegionRegistry; use common_meta::sequence::SequenceBuilder; -use common_meta::wal_provider::build_wal_provider; use common_procedure::ProcedureManagerRef; use common_procedure::local::EventRecorderHandle; use common_procedure::options::ProcedureConfig; use common_telemetry::logging::SlowQueryOptions; use common_test_util::find_workspace_path; -use common_wal::config::{DatanodeWalConfig, MetasrvWalConfig}; -use datanode::datanode::DatanodeBuilder; +use common_wal::config::DatanodeWalConfig; +use datanode::datanode::{Datanode, DatanodeBuilder}; use flow::{FlownodeBuilder, FrontendClient, GrpcQueryHandlerWithBoxedError}; use frontend::frontend::Frontend; use frontend::instance::Instance; @@ -63,6 +62,8 @@ use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode pub struct GreptimeDbStandalone { /// Storage engine for assertions across protocol and storage boundaries. pub mito_engine: mito2::engine::MitoEngine, + /// Owns the region server and the log store, so a test can shut them down. + pub datanode: Datanode, pub frontend: Arc, pub opts: StandaloneOptions, pub guard: TestGuard, @@ -81,7 +82,6 @@ impl GreptimeDbStandalone { pub struct GreptimeDbStandaloneBuilder { instance_name: String, datanode_wal_config: DatanodeWalConfig, - metasrv_wal_config: MetasrvWalConfig, store_providers: Option>, default_store: Option, plugin: Option, @@ -102,7 +102,6 @@ impl GreptimeDbStandaloneBuilder { plugin: None, default_store: None, datanode_wal_config: DatanodeWalConfig::default(), - metasrv_wal_config: MetasrvWalConfig::default(), // Enable slow query log with 1s threshold by default for integration tests. slow_query_options: SlowQueryOptions { enable: true, @@ -199,12 +198,6 @@ impl GreptimeDbStandaloneBuilder { self } - #[must_use] - pub fn with_metasrv_wal_config(mut self, metasrv_wal_config: MetasrvWalConfig) -> Self { - self.metasrv_wal_config = metasrv_wal_config; - self - } - pub async fn build_with( &self, kv_backend: KvBackendRef, @@ -281,10 +274,13 @@ impl GreptimeDbStandaloneBuilder { .step(10) .build(), ); - let kafka_options = opts.wal.clone().try_into().unwrap(); - let wal_provider = build_wal_provider(&kafka_options, kv_backend.clone()) - .await - .unwrap(); + let wal_provider = cmd::standalone::build_standalone_wal_provider( + &opts.wal, + opts.datanode_options().node_id, + kv_backend.clone(), + ) + .await + .unwrap(); let wal_provider = Arc::new(wal_provider); let table_metadata_allocator = Arc::new(TableMetadataAllocator::new( table_id_allocator, @@ -363,6 +359,7 @@ impl GreptimeDbStandaloneBuilder { GreptimeDbStandalone { mito_engine: datanode.region_server().mito_engine().unwrap(), + datanode, frontend: Arc::new(frontend), opts, guard, @@ -399,7 +396,7 @@ impl GreptimeDbStandaloneBuilder { storage: opts.storage, procedure: procedure_config, metadata_store: kv_backend_config, - wal: self.metasrv_wal_config.clone().into(), + wal: opts.wal, grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), diff --git a/tests-integration/src/tests/test_util.rs b/tests-integration/src/tests/test_util.rs index 5ddcd0748a9..7c99e89ee08 100644 --- a/tests-integration/src/tests/test_util.rs +++ b/tests-integration/src/tests/test_util.rs @@ -263,15 +263,8 @@ pub(crate) async fn standalone_with_kafka_wal() -> Option>(); let test_name = uuid::Uuid::new_v4().to_string(); - let builder = GreptimeDbStandaloneBuilder::new(&test_name) - .with_datanode_wal_config(DatanodeWalConfig::Kafka(DatanodeKafkaConfig { - connection: KafkaConnectionConfig { - broker_endpoints: endpoints.clone(), - ..Default::default() - }, - ..Default::default() - })) - .with_metasrv_wal_config(MetasrvWalConfig::Kafka(MetasrvKafkaConfig { + let builder = GreptimeDbStandaloneBuilder::new(&test_name).with_datanode_wal_config( + DatanodeWalConfig::Kafka(DatanodeKafkaConfig { connection: KafkaConnectionConfig { broker_endpoints: endpoints, ..Default::default() @@ -282,7 +275,8 @@ pub(crate) async fn standalone_with_kafka_wal() -> Option Output { + SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap() +} + +async fn query_rows(standalone: &GreptimeDbStandalone) -> String { + execute_sql(standalone.fe_instance(), QUERY) + .await + .data + .pretty_print() + .await +} + +/// Returns the number of WAL objects under [`WAL_ROOT`] of the default store. +async fn count_wal_objects(standalone: &GreptimeDbStandalone) -> usize { + let store = match &standalone.opts.storage.store { + ObjectStoreConfig::File(_) => { + ObjectStore::new(Fs::default().root(&standalone.opts.storage.data_home)).unwrap() + } + ObjectStoreConfig::S3(s3) => ObjectStore::new(S3::from(&s3.connection)).unwrap(), + other => panic!("unexpected store {other:?}"), + }; + store + .list_with(&format!("{WAL_ROOT}/")) + .recursive(true) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.metadata().is_file()) + .count() +} + +/// Stops the procedure manager, the region server and the WAL of the +/// instance. +async fn shutdown(standalone: &mut GreptimeDbStandalone) { + standalone.procedure_manager.stop().await.unwrap(); + standalone.datanode.shutdown().await.unwrap(); +} + +/// Shuts the instance down and builds it again on its metadata, data home and +/// object store, like a process restart. +async fn restart( + builder: &GreptimeDbStandaloneBuilder, + mut standalone: GreptimeDbStandalone, +) -> GreptimeDbStandalone { + shutdown(&mut standalone).await; + let GreptimeDbStandalone { + opts, + guard, + kv_backend, + .. + } = standalone; + + let (procedure_manager, event_recorder_handle) = + standalone::build_procedure_manager(kv_backend.clone(), ProcedureConfig::default()); + builder + .build_with( + kv_backend, + guard, + opts, + procedure_manager, + event_recorder_handle, + true, + ) + .await +} + +/// Writes batches that land in separate WAL objects, restarts before and +/// after a flush, and checks that the rows survive both restarts. +async fn run_restarts_on_object_store_wal(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let wal_config = ObjectStoreWalConfig { + prefix: WAL_ROOT.to_string(), + ..Default::default() + }; + let builder = GreptimeDbStandaloneBuilder::new("object_store_wal") + .with_default_store_type(store_type) + .with_datanode_wal_config(DatanodeWalConfig::ObjectStore(wal_config)); + let standalone = builder.build().await; + + execute_sql( + standalone.fe_instance(), + r#" + CREATE TABLE cpu ( + hostname STRING PRIMARY KEY, + usage_user DOUBLE, + ts TIMESTAMP TIME INDEX + ) + "#, + ) + .await; + // Every insert returns once its entries are durable, so the batches land + // in separate WAL objects after the object that started the epoch. + for batch in 0..BATCHES { + let values = (0..ROWS_PER_BATCH) + .map(|row| { + let i = batch * ROWS_PER_BATCH + row; + format!( + "('host_{i}', {i}.0, {})", + 1_686_567_600_000 + i as i64 * 1000 + ) + }) + .collect::>() + .join(", "); + execute_sql( + standalone.fe_instance(), + &format!("INSERT INTO cpu VALUES {values}"), + ) + .await; + } + let expected = query_rows(&standalone).await; + assert_eq!(BATCHES * ROWS_PER_BATCH + 4, expected.lines().count()); + let objects = count_wal_objects(&standalone).await; + assert!( + objects > BATCHES, + "expected more than {BATCHES} WAL objects" + ); + // No Raft Engine log store is created as a fallback. + assert!( + !Path::new(&standalone.opts.storage.data_home) + .join("wal") + .exists() + ); + + // Nothing was flushed, so the rows come back from the WAL alone; the + // reopened store starts its epoch with another object. + let standalone = restart(&builder, standalone).await; + assert_eq!(expected, query_rows(&standalone).await); + assert!(count_wal_objects(&standalone).await > objects); + + // The flushed rows come back from the SST files. + execute_sql(standalone.fe_instance(), "ADMIN FLUSH_TABLE('cpu')").await; + let mut standalone = restart(&builder, standalone).await; + assert_eq!(expected, query_rows(&standalone).await); + shutdown(&mut standalone).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_object_store_wal_survives_restarts_on_file() { + run_restarts_on_object_store_wal(StorageType::File).await; +} + +/// Runs against the S3 bucket of the `GT_S3_*` environment variables, so it +/// is skipped unless they are set. +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_object_store_wal_survives_restarts_on_s3() { + if !StorageType::S3.test_on() { + return; + } + run_restarts_on_object_store_wal(StorageType::S3).await; +} From c82da8308b19851502e3500f6199e9113a198fcd Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 8 Oct 2026 06:33:56 +0000 Subject: [PATCH 097/133] feat: add bounded Metric logical-table batch DDL (#9427) * feat: add bounded Metric logical-table batch DDL Signed-off-by: jeremyhi * fix: enforce logical batch DDL interceptor boundaries Signed-off-by: jeremyhi * fix: bound rewritten DDL size through shared SQL formatting Signed-off-by: jeremyhi * fix: preserve logical DDL error diagnostics Signed-off-by: jeremyhi * test: verify logical DDL error logging Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- docs/rfcs/2026-09-11-metric-export-import.md | 30 +- src/frontend/AGENTS.md | 3 + src/frontend/src/instance.rs | 315 +++++++++++++ src/operator/src/statement/ddl.rs | 99 +++- src/servers/src/error.rs | 7 +- src/servers/src/http.rs | 112 ++++- src/servers/src/http/handler.rs | 146 +++++- src/servers/src/query_handler/sql.rs | 22 + src/sql/src/statements/create.rs | 325 ++++++++++--- src/sql/src/statements/option_map.rs | 71 ++- src/sql/src/util.rs | 22 +- tests-integration/tests/http.rs | 453 +++++++++++++++++++ 12 files changed, 1493 insertions(+), 112 deletions(-) diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 3d2c93e4695..8ddc31f40b0 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -390,14 +390,34 @@ validation, schema-option inheritance and target metadata resolution. Bound the request size and the importer's DDL buffering separately. The production CLI needs an authenticated server operation that checks CREATE -authorization for every member and reuses the ordinary DDL path. Transport, -capability discovery and request limits must be agreed before accepting this -interface. Initially bound a batch to 128 statements and 1 MiB of SQL, with an -explicit oversized-statement error; separately bound the importer's parsed DDL -buffer. A target without batch-DDL capability uses ordinary DDL if it supports +authorization for every member and reuses the ordinary DDL path. Bound the +importer's parsed DDL buffer separately from the server request limits below. +A target without batch-DDL capability uses ordinary DDL if it supports the snapshot's data layout. Do not retry an ambiguously completed batch through another path. +The server operation is `POST /v1/ddl/logical-tables` with an +`application/x-www-form-urlencoded` body containing `sql` and optional `db`. +`db` uses the same catalog/schema interpretation and defaults as `/v1/sql`. +Authenticated `GET /v1/capabilities` advertises `metric_batch_ddl: 1` when the +SQL handler supports the operation, independently of packed import/export. +The form body is limited to 4 MiB before decoding, or the configured HTTP body +limit when that limit is stricter. Decoded SQL is limited to 1 MiB and 1–128 +statements. Only explicit Metric logical CREATE TABLE statements in one +catalog/schema with one literal `on_physical_table` value are accepted; +duplicate normalized target names are rejected. + +All members pass SQL interception, authorization and ordinary CREATE preparation +before one logical-table batch procedure is submitted. Success uses the normal +GreptimeDB v1 response with one `affectedrows: 0` output per input statement in +order; this is CREATE's output, not a count of newly created tables. A failure +returns one error without a partial success list. Validation failures before +submission have no DDL side effects. Once submitted, the existing procedure may +have applied some or all changes even if the HTTP request times out or disconnects; +there is no cross-table rollback guarantee. `IF NOT EXISTS` keeps ordinary CREATE +semantics. The operation does not regroup requests or retry through single-table +CREATE. + ## Packed data COPY The CLI retains one `COPY DATABASE FROM` operation per `(chunk, schema)` task. diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index e16cbfc9872..bab27676655 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -39,6 +39,9 @@ remote datanodes via `operator`/`client`. - **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, permission checks, timeout/cancellation, and delegates planning/execution to `StatementExecutor`. Distributed scans enter through `region_query.rs`. +- **Logical-table batch DDL** (`instance.rs`): dedicated SQL handler → bounded + parse/interception → all-member permissions → operator CREATE preparation → + one logical-table batch procedure. Keep pre-submission validation side-effect free. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing, meter admission) → local `RegionServer` diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 6e531214127..d3e6f2addfc 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -782,6 +782,98 @@ impl Instance { query_interceptor.post_execute(Output { data, meta }, query_ctx) } + async fn batch_create_logical_tables( + &self, + query: &str, + query_ctx: QueryContextRef, + ) -> Result> { + use servers::query_handler::sql::{ + MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS, + }; + ensure!(!self.is_suspended(), error::SuspendedSnafu); + ensure!( + query.len() <= MAX_LOGICAL_TABLE_DDL_BYTES, + InvalidSqlSnafu { + err_msg: "logical-table DDL exceeds 1 MiB" + } + ); + let interceptor = self.plugins.get::>(); + let interceptor = interceptor.as_ref(); + let query = interceptor.pre_parsing(query, query_ctx.clone())?; + ensure!( + query.len() <= MAX_LOGICAL_TABLE_DDL_BYTES, + InvalidSqlSnafu { + err_msg: "logical-table DDL exceeds 1 MiB" + } + ); + let statements = parse_stmt(query.as_ref(), query_ctx.sql_dialect()).map_err(|_| { + InvalidSqlSnafu { + err_msg: "invalid logical-table DDL", + } + .build() + })?; + ensure!( + !statements.is_empty() && statements.len() <= MAX_LOGICAL_TABLE_DDL_STATEMENTS, + InvalidSqlSnafu { + err_msg: "logical-table batch requires 1 to 128 statements" + } + ); + let original_statements = interceptor.map(|_| statements.clone()); + let statements = interceptor.post_parsing(statements, query_ctx.clone())?; + ensure!( + !statements.is_empty() && statements.len() <= MAX_LOGICAL_TABLE_DDL_STATEMENTS, + InvalidSqlSnafu { + err_msg: "logical-table batch requires 1 to 128 statements" + } + ); + let measure_rewritten_sql = || { + let mut remaining = MAX_LOGICAL_TABLE_DDL_BYTES - (statements.len() - 1); + statements.iter().all(|statement| { + matches!(statement, Statement::CreateTable(create) if create.consume_sql_size_budget(&mut remaining)) + }) + }; + ensure!( + original_statements + .as_ref() + .is_none_or(|original| original == &statements) + || measure_rewritten_sql(), + InvalidSqlSnafu { + err_msg: "invalid or oversized rewritten logical-table DDL" + } + ); + for statement in &statements { + ensure!( + matches!(statement, Statement::CreateTable(_)), + InvalidSqlSnafu { + err_msg: "batch requires explicit Metric logical CREATE TABLE statements" + } + ); + interceptor.pre_execute(Some(statement), None, query_ctx.clone())?; + } + for statement in &statements { + self.check_sql_permission(statement, &query_ctx).await?; + check_permission(self.plugins.clone(), statement, &query_ctx)?; + } + let count = statements.len(); + let creates = statements + .into_iter() + .filter_map(|statement| { + if let Statement::CreateTable(create) = statement { + Some(create) + } else { + None + } + }) + .collect(); + self.statement_executor + .batch_create_logical_tables(creates, query_ctx.clone()) + .await + .context(TableOperationSnafu)?; + (0..count) + .map(|_| interceptor.post_execute(Output::new_with_affected_rows(0), query_ctx.clone())) + .collect() + } + #[tracing::instrument(skip_all, name = "SqlQueryHandler::do_query")] async fn do_query_inner(&self, query: &str, query_ctx: QueryContextRef) -> Vec> { if self.is_suspended() { @@ -1113,6 +1205,21 @@ impl Instance { #[async_trait] impl SqlQueryHandler for Instance { + fn supports_metric_batch_ddl(&self) -> bool { + true + } + + async fn create_logical_tables( + &self, + query: &str, + query_ctx: QueryContextRef, + ) -> server_error::Result> { + self.batch_create_logical_tables(query, query_ctx) + .await + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu) + } + async fn do_query( &self, query: &str, @@ -4145,6 +4252,214 @@ mod tests { Ok(()) } + #[tokio::test] + async fn logical_ddl_validates_and_authorizes_rewritten_batch_before_submission() + -> TestResult<()> { + struct Rewrite(&'static str); + impl SqlQueryInterceptor for Rewrite { + type Error = Error; + fn pre_parsing<'a>( + &self, + query: &'a str, + _: QueryContextRef, + ) -> Result> { + Ok(if self.0 == "sql_bytes" { + std::borrow::Cow::Owned(format!("{query}{}", " ".repeat(1024 * 1024 + 1))) + } else { + std::borrow::Cow::Borrowed(query) + }) + } + fn pre_execute( + &self, + statement: Option<&Statement>, + _: Option<&LogicalPlan>, + _: QueryContextRef, + ) -> Result<()> { + ensure!( + !(self.0 == "pre_execute" + && matches!(statement, + Some(Statement::CreateTable(create)) if create.name.to_string() == "allowed2")), + InvalidSqlSnafu { + err_msg: "second member rejected by interceptor" + } + ); + Ok(()) + } + fn post_parsing( + &self, + mut statements: Vec, + _: QueryContextRef, + ) -> Result> { + match self.0 { + "target" => { + statements[1] = parse_one_sql( + "CREATE TABLE denied (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')", + ) + } + "count" => statements = (0..129).map(|i| parse_one_sql(&format!( + "CREATE TABLE allowed_{i} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy')" + ))).collect(), + "type" => statements[1] = parse_one_sql("SELECT 1"), + "secret_bytes" | "array_bytes" | "small_option" => { + if let Statement::CreateTable(create) = &mut statements[1] { + let value = "x".repeat(if self.0 == "small_option" { 10 } else { 2 * 1024 * 1024 }); + if self.0 == "array_bytes" { + create.options.insert_options("non_scalar", vec![value.as_str()].into()); + } else { + create.options.insert("secret_access_key".to_string(), value); + } + } + } + "partition_shape" => { + if let Statement::CreateTable(create) = &mut statements[1] { + create.partitions = Some(sql::statements::create::Partitions { + column_list: vec![], + exprs: vec![sqlparser::ast::Expr::Identifier(sqlparser::ast::Ident::new("hidden"))], + }); + } + } + "vector_hidden" => { + statements[1] = parse_one_sql("CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX, v VECTOR(3)) ENGINE=metric WITH (on_physical_table='phy')"); + if let Statement::CreateTable(create) = &mut statements[1] { + create.columns[1].column_def.options.push(sqlparser::ast::ColumnOptionDef { + name: None, + option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(200_000)), + }); + } + } + mode if mode.starts_with("fulltext_") || mode.starts_with("skipping_") + || mode.starts_with("inverted_") || mode.starts_with("vector_") => { + if let Statement::CreateTable(create) = &mut statements[1] { + let (kind, value_type) = mode.split_once('_').unwrap(); + let value = "x".repeat(if value_type == "small" { 10 } else { 2 * 1024 * 1024 }); + let extensions = &mut create.columns[1].extensions; + let options = match kind { + "fulltext" => &mut extensions.fulltext_index_options, + "skipping" => &mut extensions.skipping_index_options, + "inverted" => &mut extensions.inverted_index_options, + _ => &mut extensions.vector_options, + }.get_or_insert_default(); + if kind == "vector" { + options.insert("dim".to_string(), "3".to_string()); + } + if value_type == "array" { + options.insert_options("non_scalar", vec![value.as_str()].into()); + } else { + options.insert("secret_access_key".to_string(), value); + } + } + } + "physical_null" | "physical_array" => { + let value = if self.0 == "physical_null" { "NULL" } else { "['x']" }; + statements[1] = parse_one_sql(&format!( + "CREATE TABLE allowed2 (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='phy', physical_metric_table={value})" + )); + } + "ast_bytes" => { + if let Statement::CreateTable(create) = &mut statements[1] { + create.name = + ObjectName(vec![sqlparser::ast::ObjectNamePart::Identifier( + sqlparser::ast::Ident::new("t".repeat(1024 * 1024)), + )]); + } + } + _ => {} + } + Ok(statements) + } + } + for mode in [ + "none", + "target", + "count", + "type", + "sql_bytes", + "ast_bytes", + "secret_bytes", + "array_bytes", + "small_option", + "physical_null", + "physical_array", + "fulltext_secret", + "fulltext_array", + "fulltext_small", + "skipping_secret", + "skipping_array", + "skipping_small", + "inverted_secret", + "inverted_array", + "inverted_small", + "vector_secret", + "vector_array", + "vector_small", + "vector_hidden", + "partition_shape", + "pre_execute", + "cross_catalog", + "valid", + ] { + let catalog = + catalog::memory::MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + let executor = Arc::new(MockProcedureExecutor::new(catalog.clone())); + let plugins = Plugins::new(); + plugins.insert::(Arc::new(RejectUnresolvedPermissionChecker)); + plugins.insert::>(Arc::new(Rewrite(mode))); + plugins.insert(QueryOptions { + disallow_cross_catalog_query: mode == "cross_catalog", + }); + let instance = test_instance_with_catalog_manager( + catalog, + test_table(1025, "target")?, + plugins, + None, + executor.clone(), + ) + .await?; + let create = |name| { + format!( + "CREATE TABLE {name} (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')" + ) + }; + let sql = format!( + "{};{}", + create("allowed"), + create(match mode { + "none" => "denied", + "cross_catalog" => "other.public.allowed2", + _ => "allowed2", + }) + ); + let result = + SqlQueryHandler::create_logical_tables(&instance, &sql, QueryContext::arc()).await; + let submits = matches!(mode, "valid" | "small_option") + || (mode.ends_with("_small") && !mode.starts_with("vector_")); + if matches!(mode, "none" | "target") { + let err = result.unwrap_err(); + assert_eq!( + err.status_code(), + StatusCode::PermissionDenied, + "{mode}: {err}" + ); + } else { + assert_eq!( + result.unwrap_err().status_code(), + match mode { + _ if submits => StatusCode::Unsupported, + "cross_catalog" => StatusCode::AccessDenied, + _ => StatusCode::InvalidArguments, + }, + "{mode}" + ); + } + assert_eq!( + executor.submitted.lock().unwrap().len(), + usize::from(submits), + "{mode}" + ); + } + Ok(()) + } + /// `DROP VIEW IF EXISTS` on a missing view through the direct gRPC ingress must /// succeed with 0 affected rows (no error, no DDL task submitted), instead of /// returning `TableNotFound`. diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 2e3c88af463..fa9da9f322e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -93,7 +93,9 @@ use sql::statements::create::{ }; use sql::statements::statement::Statement; use sqlparser::ast::{Expr, Ident, UnaryOperator, Value as ParserValue}; -use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use store_api::metric_engine_consts::{ + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; use store_api::mito_engine_options::APPEND_MODE_KEY; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::TableRef; @@ -341,7 +343,22 @@ impl StatementExecutor { #[tracing::instrument(skip_all)] pub async fn create_table(&self, stmt: CreateTable, ctx: QueryContextRef) -> Result { - let (catalog, schema, _table) = table_idents_to_full_name(&stmt.name, &ctx) + let mut create_expr = self.prepare_create_table(&stmt, &ctx).await?; + self.create_table_inner( + &mut create_expr, + stmt.partitions, + ctx, + TriggerReason::Manual, + ) + .await + } + + async fn prepare_create_table( + &self, + stmt: &CreateTable, + ctx: &QueryContextRef, + ) -> Result { + let (catalog, schema, _table) = table_idents_to_full_name(&stmt.name, ctx) .map_err(BoxedError::new) .context(error::ExternalSnafu)?; @@ -356,7 +373,7 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)? .map(|v| v.into_inner()); - let create_expr = &mut expr_helper::create_to_expr(&stmt, &ctx)?; + let mut create_expr = expr_helper::create_to_expr(stmt, ctx)?; // TTL and compaction options are handled separately; ingestion quota is database-only. if let Some(schema_options) = schema_options { for (key, value) in schema_options.extra_options.iter() { @@ -370,8 +387,72 @@ impl StatementExecutor { } } - self.create_table_inner(create_expr, stmt.partitions, ctx, TriggerReason::Manual) - .await + Ok(create_expr) + } + + /// Prepares every explicit logical CREATE before submitting one batch procedure. + pub async fn batch_create_logical_tables( + &self, + statements: Vec, + ctx: QueryContextRef, + ) -> Result<()> { + let mut expressions = Vec::with_capacity(statements.len()); + let mut targets = std::collections::HashSet::new(); + let mut group = None; + for statement in statements { + ensure!( + statement.engine == METRIC_ENGINE_NAME + && statement.options.get(LOGICAL_TABLE_METADATA_KEY).is_some() + && statement + .options + .value(PHYSICAL_TABLE_METADATA_KEY) + .is_none(), + InvalidSqlSnafu { + err_msg: "batch requires explicit Metric logical CREATE TABLE statements" + } + ); + let expr = self.prepare_create_table(&statement, &ctx).await?; + let key = ( + expr.catalog_name.clone(), + expr.schema_name.clone(), + expr.table_options.get(LOGICAL_TABLE_METADATA_KEY).cloned(), + ); + ensure!( + group.as_ref().is_none_or(|group| group == &key), + InvalidSqlSnafu { + err_msg: "batch must share catalog, schema and literal on_physical_table" + } + ); + group = Some(key); + ensure!( + targets.insert(expr.table_name.clone()), + InvalidSqlSnafu { + err_msg: "duplicate target in logical-table batch" + } + ); + self.validate_logical_create(&expr, statement.partitions.as_ref(), &ctx) + .await?; + expressions.push(expr); + } + self.create_logical_tables(&expressions, ctx, TriggerReason::Manual) + .await?; + Ok(()) + } + + async fn validate_logical_create( + &self, + expr: &CreateTableExpr, + partitions: Option<&Partitions>, + ctx: &QueryContextRef, + ) -> Result<()> { + ensure_table_definition_writable(&expr.schema_name, &expr.table_name)?; + if let Some(partitions) = partitions + && !partitions.exprs.is_empty() + { + self.validate_logical_table_partition_rule(expr, partitions, ctx) + .await?; + } + Ok(()) } #[tracing::instrument(skip_all)] @@ -472,12 +553,8 @@ impl StatementExecutor { .table_options .contains_key(LOGICAL_TABLE_METADATA_KEY) { - if let Some(partitions) = partitions.as_ref() - && !partitions.exprs.is_empty() - { - self.validate_logical_table_partition_rule(create_table, partitions, &query_ctx) - .await?; - } + self.validate_logical_create(create_table, partitions.as_ref(), &query_ctx) + .await?; // Create logical tables self.create_logical_tables( std::slice::from_ref(create_table), diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index af6b2499e7f..6e558c8234b 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -200,6 +200,9 @@ pub enum Error { location: Location, }, + #[snafu(display("Unsupported operation: {}", operation))] + UnsupportedOperation { operation: String }, + #[snafu(display("Invalid query: {}", reason))] InvalidQuery { reason: String, @@ -767,7 +770,7 @@ impl ErrorExt for Error { | InvalidPromRemoteReadQueryResult { .. } | OtlpMetricModeIncompatible { .. } => StatusCode::IllegalState, - UnsupportedDataType { .. } => StatusCode::Unsupported, + UnsupportedDataType { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported, #[cfg(not(windows))] UpdateJemallocMetrics { .. } => StatusCode::Internal, @@ -941,7 +944,7 @@ impl From for Error { } } -fn log_error_if_necessary(error: &Error) { +pub(crate) fn log_error_if_necessary(error: &Error) { if error.status_code().should_log_error() { error!(error; "Failed to handle HTTP request "); } else { diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 34f214bae26..4555b85ac69 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -681,7 +681,12 @@ impl HttpServerBuilder { } pub fn with_sql_handler(self, sql_handler: ServerSqlQueryHandlerRef) -> Self { - let sql_router = HttpServer::route_sql(ApiState { sql_handler }); + let batch_body_limit = if self.options.body_limit.0 == 0 { + 4 * 1024 * 1024 + } else { + (self.options.body_limit.0 as usize).min(4 * 1024 * 1024) + }; + let sql_router = HttpServer::route_sql(ApiState { sql_handler }, batch_body_limit); Self { router: self @@ -1390,16 +1395,24 @@ impl HttpServer { .with_state(log_state) } - fn route_sql(api_state: ApiState) -> Router { + fn route_sql(api_state: ApiState, batch_body_limit: usize) -> Router { Router::new() .route( "/capabilities", - routing::get(|| async { - axum::Json( - serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}), - ) + routing::get(|State(state): State| async move { + let mut capabilities = + serde_json::json!({"metric_packed_import": 1, "metric_packed_export": 1}); + if state.sql_handler.supports_metric_batch_ddl() { + capabilities["metric_batch_ddl"] = serde_json::json!(1); + } + axum::Json(capabilities) }), ) + .route( + "/ddl/logical-tables", + routing::post(handler::create_logical_tables) + .layer(DefaultBodyLimit::max(batch_body_limit)), + ) .route( "/sql", routing::get(handler::sql).post(handler::sql).layer( @@ -1939,6 +1952,15 @@ mod test { client.get("/v1/capabilities").send().await.status(), StatusCode::UNAUTHORIZED ); + assert_eq!( + client + .post("/v1/ddl/logical-tables") + .form(&[("sql", "SELECT 1")]) + .send() + .await + .status(), + StatusCode::UNAUTHORIZED + ); let response = client .get("/v1/capabilities") .header("Authorization", "Basic dXNlcjpwYXNzd29yZA==") @@ -1952,6 +1974,84 @@ mod test { } } + #[tokio::test] + async fn logical_ddl_form_limits_and_unsupported_handler() { + use axum::body::{Body, to_bytes}; + use tower::ServiceExt; + + async fn post_form(app: Router, body: String) -> (StatusCode, serde_json::Value) { + let response = app + .oneshot( + Request::builder() + .method(Method::POST) + .uri("/v1/ddl/logical-tables") + .header("Content-Type", "application/x-www-form-urlencoded") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + (status, serde_json::from_slice(&body).unwrap()) + } + + for limit in [0, 64] { + let (tx, _rx) = mpsc::channel(1); + let server = HttpServerBuilder::new(HttpOptions { + body_limit: ReadableSize(limit), + ..Default::default() + }) + .with_sql_handler(Arc::new(DummyInstance { _tx: tx })) + .build(); + let app = server.build(server.make_app()).unwrap(); + let client = TestClient::new(app.clone()).await; + assert_eq!( + client.get("/v1/ddl/logical-tables").send().await.status(), + StatusCode::METHOD_NOT_ALLOWED + ); + let small = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", "CREATE TABLE t")]) + .send() + .await; + assert_eq!( + small.json::().await["code"], + common_error::status_code::StatusCode::Unsupported as u32 + ); + let (_, body) = post_form( + app.clone(), + if limit == 0 { + format!("sql={}", "%20".repeat(1024 * 1024)) + } else { + "sql=".to_string() + &" ".repeat(65) + }, + ) + .await; + let code = body["code"].clone(); + assert_eq!( + code, + if limit == 0 { + common_error::status_code::StatusCode::Unsupported as u32 + } else { + common_error::status_code::StatusCode::InvalidArguments as u32 + } + ); + for body in [ + format!("sql={}", "x".repeat(1024 * 1024 + 1)), + format!("sql={}", "%C3%A9".repeat(512 * 1024 + 1)), + format!("sql=CREATE+TABLE+t&padding={}", "x".repeat(4 * 1024 * 1024)), + ] { + let (status, body) = post_form(app.clone(), body).await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!( + body["code"], + common_error::status_code::StatusCode::InvalidArguments as u32 + ); + } + } + } + #[tokio::test] pub async fn test_extra_router_respects_api_surface() { // Downstream routes added through `with_extra_router()` (e.g. enterprise diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index bd31ee35ace..8cbfc25ff07 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -44,7 +44,10 @@ use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use tokio::sync::{Notify, watch}; -use crate::error::{CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result}; +use crate::error::{ + CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result, + log_error_if_necessary, +}; use crate::http::header::collect_plan_metrics; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -229,6 +232,83 @@ pub async fn sql( resp.with_execution_time(start.elapsed().as_millis() as u64) } +/// Form parameters for the bounded logical-table DDL operation. +#[derive(Deserialize)] +pub struct LogicalTableDdl { + sql: Option, + db: Option, +} + +/// Creates one batch of Metric logical tables using the normal v1 response. +pub async fn create_logical_tables( + State(state): State, + Extension(mut query_ctx): Extension, + form: std::result::Result, FormRejection>, +) -> HttpResponse { + let start = Instant::now(); + let form = match form { + Ok(Form(form)) => form, + Err(_) => { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "invalid or oversized form body".to_string(), + )); + } + }; + let Some(sql) = form.sql else { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "sql parameter is required.".to_string(), + )); + }; + if sql.len() > crate::query_handler::sql::MAX_LOGICAL_TABLE_DDL_BYTES { + return HttpResponse::Error(ErrorResponse::from_error_message( + StatusCode::InvalidArguments, + "logical-table DDL exceeds 1 MiB".to_string(), + )); + } + if let Some(db) = form.db { + let (catalog, schema) = parse_catalog_and_schema_from_db_string(&db); + query_ctx.set_current_catalog(&catalog); + query_ctx.set_current_schema(&schema); + } + query_ctx.set_channel(Channel::HttpSql); + let query_ctx = Arc::new(query_ctx); + if let Some((status, msg)) = validate_schema(state.sql_handler.clone(), query_ctx.clone()).await + { + return HttpResponse::Error(ErrorResponse::from_error_message(status, msg)); + } + let outputs = match state + .sql_handler + .create_logical_tables(&sql, query_ctx) + .await + { + Ok(outputs) + if outputs + .iter() + .all(|output| matches!(output.data, OutputData::AffectedRows(0))) => + { + outputs.into_iter().map(Ok).collect() + } + other => { + let status = match other { + Err(error) => { + log_error_if_necessary(&error); + error.status_code() + } + Ok(_) => StatusCode::Internal, + }; + return HttpResponse::Error( + ErrorResponse::from_error_message(status, "logical-table batch failed".to_string()) + .with_execution_time(start.elapsed().as_millis() as u64), + ); + } + }; + GreptimedbV1Response::from_output(outputs) + .await + .with_execution_time(start.elapsed().as_millis() as u64) +} + /// Limits each statement's response before collecting batches and converting rows. fn limit_output_rows(output: Output, limit: usize) -> Result { let mut remaining = limit; @@ -994,13 +1074,16 @@ mod tests { use std::sync::atomic::AtomicUsize; use async_trait::async_trait; + use common_error::ext::{BoxedError, PlainError}; use common_query::OutputMeta; + use common_telemetry::tracing::instrument::WithSubscriber; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{BinaryVector, UInt32Vector, VectorRef}; use futures::stream; use query::query_engine::DescribeResult; + use snafu::IntoError; use super::*; use crate::query_handler::sql::SqlQueryHandler; @@ -1015,6 +1098,12 @@ mod tests { std::mem::take(&mut *self.outputs.lock().unwrap()) } + async fn create_logical_tables(&self, _: &str, _: QueryContextRef) -> Result> { + std::mem::take(&mut *self.outputs.lock().unwrap()) + .into_iter() + .collect() + } + async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result { unimplemented!() } @@ -1068,6 +1157,61 @@ mod tests { .await } + #[tokio::test] + async fn test_logical_ddl_logs_diagnostics_without_exposing_sql_or_response_details() { + let sql = "CREATE TABLE logs (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='request-only-value')"; + for (status, level) in [ + (StatusCode::InvalidArguments, "WARN"), + (StatusCode::Internal, "ERROR"), + ] { + let logs = tempfile::NamedTempFile::new().unwrap(); + let writer = logs.reopen().unwrap(); + let subscriber = common_telemetry::tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_writer(move || writer.try_clone().unwrap()) + .finish(); + let error = crate::error::ExecuteQuerySnafu.into_error(BoxedError::new( + PlainError::new("sensitive-option-value".to_string(), status), + )); + let response = create_logical_tables( + State(ApiState { + sql_handler: Arc::new(TestSqlQueryHandler { + outputs: Mutex::new(vec![Err(error)]), + }), + }), + Extension(QueryContext::with_db_name(None)), + Ok(Form(LogicalTableDdl { + sql: Some(sql.to_string()), + db: None, + })), + ) + .with_subscriber(subscriber) + .await + .into_response(); + let diagnostics = std::fs::read_to_string(logs.path()).unwrap(); + assert!(diagnostics.contains(level), "{diagnostics}"); + assert!( + diagnostics.contains("Failed to execute query"), + "{diagnostics}" + ); + assert!( + diagnostics.contains("sensitive-option-value"), + "{diagnostics}" + ); + assert!(!diagnostics.contains(sql), "{diagnostics}"); + assert!(!diagnostics.contains("request-only-value"), "{diagnostics}"); + assert!(!format!("{:?}", response.headers()).contains("sensitive-option-value")); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(body["code"], status as u32); + assert_eq!(body["error"], "logical-table batch failed"); + assert!(body.get("output").is_none()); + } + } + fn number_batches() -> RecordBatches { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "number", diff --git a/src/servers/src/query_handler/sql.rs b/src/servers/src/query_handler/sql.rs index 9bf99a44441..1b65ff002db 100644 --- a/src/servers/src/query_handler/sql.rs +++ b/src/servers/src/query_handler/sql.rs @@ -26,8 +26,30 @@ use crate::error::Result; pub type ServerSqlQueryHandlerRef = Arc; +/// Maximum decoded SQL size accepted by the logical-table batch operation. +pub const MAX_LOGICAL_TABLE_DDL_BYTES: usize = 1024 * 1024; +/// Maximum number of CREATE statements in one logical-table batch. +pub const MAX_LOGICAL_TABLE_DDL_STATEMENTS: usize = 128; + #[async_trait] pub trait SqlQueryHandler { + /// Whether this handler implements the bounded logical-table DDL operation. + fn supports_metric_batch_ddl(&self) -> bool { + false + } + + /// Validates and authorizes the whole batch before submitting one procedure. + async fn create_logical_tables( + &self, + _query: &str, + _query_ctx: QueryContextRef, + ) -> Result> { + crate::error::UnsupportedOperationSnafu { + operation: "metric batch DDL", + } + .fail() + } + async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec>; /// Executes an HTTP analyze-stream query. diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 2cb78431969..725743ee679 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::collections::HashMap; -use std::fmt::{Display, Formatter}; +use std::fmt::{Display, Formatter, Write}; use common_catalog::consts::FILE_ENGINE; use datatypes::json::{JSON2_DEFAULT_MAX_AUTO_EXPANDED_PATHS, JsonSettings}; @@ -68,6 +68,36 @@ fn format_table_constraint(constraints: &[TableConstraint]) -> String { constraints.iter().map(|c| format_indent!(c)).join(LINE_SEP) } +fn write_list( + f: &mut dyn Write, + values: &[T], + separator: &str, + prefix: &str, +) -> std::fmt::Result { + for (index, value) in values.iter().enumerate() { + if index > 0 { + f.write_str(separator)?; + } + write!(f, "{prefix}{value}")?; + } + Ok(()) +} + +// sqlparser's Custom Display joins all modifiers into an intermediate String. +fn write_data_type(f: &mut dyn Write, data_type: &DataType) -> std::fmt::Result { + if let DataType::Custom(name, modifiers) = data_type { + write!(f, "{name}")?; + if !modifiers.is_empty() { + f.write_char('(')?; + write_list(f, modifiers, COMMA_SEP, "")?; + f.write_char(')')?; + } + Ok(()) + } else { + write!(f, "{data_type}") + } +} + /// Table constraint for create table statement. #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub enum TableConstraint { @@ -81,7 +111,9 @@ impl Display for TableConstraint { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { TableConstraint::PrimaryKey { columns } => { - write!(f, "PRIMARY KEY ({})", format_list_comma!(columns)) + f.write_str("PRIMARY KEY (")?; + write_list(f, columns, COMMA_SEP, "")?; + f.write_str(")") } TableConstraint::TimeIndex { column } => { write!(f, "TIME INDEX ({})", column) @@ -143,12 +175,34 @@ pub struct Json2Options { impl Display for Json2Options { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let mut options = Vec::with_capacity(self.type_hints.len() + 1); + f.write_str("(\n ")?; if let Some(max) = self.max_auto_expanded_paths { - options.push(format!("max_auto_expanded_paths = {max}")); + write!(f, "max_auto_expanded_paths = {max}")?; } - options.extend(self.type_hints.iter().map(format_json_type_hint)); - write!(f, "(\n {}\n )", options.iter().join(",\n ")) + for (index, hint) in self.type_hints.iter().enumerate() { + if index > 0 || self.max_auto_expanded_paths.is_some() { + f.write_str(",\n ")?; + } + for (index, segment) in hint.path.iter().enumerate() { + if index > 0 { + f.write_char('.')?; + } + f.write_char('"')?; + for part in segment.split_inclusive('"') { + f.write_str(part)?; + if part.ends_with('"') { + f.write_char('"')?; + } + } + f.write_char('"')?; + } + f.write_char(' ')?; + write_data_type(f, &hint.data_type)?; + if hint.inverted_index { + f.write_str(" INVERTED INDEX")?; + } + } + f.write_str("\n )") } } @@ -203,14 +257,20 @@ impl Column { impl Display for Column { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - if let Some(vector_options) = &self.extensions.vector_options + self.write_sql(f, false) + } +} + +impl Column { + fn write_sql(&self, f: &mut dyn Write, complete: bool) -> std::fmt::Result { + if !complete + && let Some(vector_options) = &self.extensions.vector_options && let Some(dim) = vector_options.get(VECTOR_OPT_DIM) { - write!(f, "{} VECTOR({})", self.column_def.name, dim)?; - return Ok(()); + return write!(f, "{} VECTOR({})", self.column_def.name, dim); } - - write!(f, "{} {}", self.column_def.name, self.column_def.data_type)?; + write!(f, "{} ", self.column_def.name)?; + write_data_type(f, &self.column_def.data_type)?; if let Some(options) = &self.extensions.json2_options { write!(f, "{options}")?; } @@ -218,30 +278,18 @@ impl Display for Column { write!(f, " {option}")?; } - if let Some(fulltext_options) = &self.extensions.fulltext_index_options { - if !fulltext_options.is_empty() { - let options = fulltext_options.kv_pairs(); - write!(f, " FULLTEXT INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " FULLTEXT INDEX")?; - } - } - - if let Some(skipping_index_options) = &self.extensions.skipping_index_options { - if !skipping_index_options.is_empty() { - let options = skipping_index_options.kv_pairs(); - write!(f, " SKIPPING INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " SKIPPING INDEX")?; - } - } - - if let Some(inverted_index_options) = &self.extensions.inverted_index_options { - if !inverted_index_options.is_empty() { - let options = inverted_index_options.kv_pairs(); - write!(f, " INVERTED INDEX WITH({})", format_list_comma!(options))?; - } else { - write!(f, " INVERTED INDEX")?; + for (kind, options) in [ + ("FULLTEXT", &self.extensions.fulltext_index_options), + ("SKIPPING", &self.extensions.skipping_index_options), + ("INVERTED", &self.extensions.inverted_index_options), + ] { + if let Some(options) = options { + write!(f, " {kind} INDEX")?; + if !options.is_empty() { + f.write_str(" WITH(")?; + options.write_sql(f, COMMA_SEP, "", complete)?; + f.write_str(")")?; + } } } @@ -317,24 +365,6 @@ fn json_type_hint_sql_data_type(data_type: &ConcreteDataType) -> Result String { - let path = hint - .path - .iter() - .map(|segment| format_json_path_segment(segment)) - .join("."); - let inverted_index = if hint.inverted_index { - " INVERTED INDEX" - } else { - "" - }; - format!("{} {}{}", path, hint.data_type, inverted_index) -} - -fn format_json_path_segment(segment: &str) -> String { - format!("\"{}\"", segment.replace('"', "\"\"")) -} - /// Partition on columns or values. /// /// - `column_list` is the list of columns in `PARTITION ON COLUMNS` clause. @@ -359,12 +389,11 @@ impl Partitions { impl Display for Partitions { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if !self.column_list.is_empty() { - write!( - f, - "PARTITION ON COLUMNS ({}) (\n{}\n)", - format_list_comma!(self.column_list), - format_list_indent!(self.exprs), - )?; + f.write_str("PARTITION ON COLUMNS (")?; + write_list(f, &self.column_list, COMMA_SEP, "")?; + f.write_str(") (\n")?; + write_list(f, &self.exprs, LINE_SEP, " ")?; + f.write_str("\n)")?; } Ok(()) } @@ -372,6 +401,53 @@ impl Display for Partitions { impl Display for CreateTable { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.write_sql(f, false) + } +} + +impl CreateTable { + /// Debits the SQL size of a rewritten AST, including unredacted and non-scalar + /// options. Returns false on exhaustion or a shape hidden by normal Display. + /// No SQL or secret value is returned. Delegated sqlparser formatters retain + /// their own allocation behavior; this is not a general AST memory limit. + pub fn consume_sql_size_budget(&self, remaining: &mut usize) -> bool { + struct Budget<'a>(&'a mut usize); + impl Write for Budget<'_> { + fn write_str(&mut self, value: &str) -> std::fmt::Result { + *self.0 = self.0.checked_sub(value.len()).ok_or(std::fmt::Error)?; + Ok(()) + } + } + if self + .partitions + .as_ref() + .is_some_and(|p| p.column_list.is_empty() && !p.exprs.is_empty()) + { + return false; + } + for column in &self.columns { + if let Some(options) = &column.extensions.vector_options { + let DataType::Custom(name, tokens) = &column.column_def.data_type else { + return false; + }; + if options.len() != 1 + || name.0.len() != 1 + || tokens.len() != 1 + || !name.0[0] + .as_ident() + .is_some_and(|name| name.value.eq_ignore_ascii_case("VECTOR")) + || !tokens[0].parse::().ok().is_some_and(|dim| { + options.get(VECTOR_OPT_DIM) == Some(dim.to_string().as_str()) + }) + { + return false; + } + } + } + self.write_sql(&mut Budget(remaining), true).is_ok() + } + + fn write_sql(&self, f: &mut dyn Write, complete: bool) -> std::fmt::Result { write!(f, "CREATE ")?; if self.engine == FILE_ENGINE { write!(f, "EXTERNAL ")?; @@ -381,16 +457,25 @@ impl Display for CreateTable { write!(f, "IF NOT EXISTS ")?; } writeln!(f, "{} (", &self.name)?; - writeln!(f, "{},", format_list_indent!(self.columns))?; - writeln!(f, "{}", format_table_constraint(&self.constraints))?; + for (index, column) in self.columns.iter().enumerate() { + if index > 0 { + f.write_str(LINE_SEP)?; + } + f.write_str(" ")?; + column.write_sql(f, complete)?; + } + f.write_str(",\n")?; + write_list(f, &self.constraints, LINE_SEP, " ")?; + f.write_char('\n')?; writeln!(f, ")")?; if let Some(partitions) = &self.partitions { writeln!(f, "{partitions}")?; } writeln!(f, "ENGINE={}", &self.engine)?; if !self.options.is_empty() { - let options = self.options.kv_pairs(); - write!(f, "WITH(\n{}\n)", format_list_indent!(options))?; + f.write_str("WITH(\n")?; + self.options.write_sql(f, LINE_SEP, " ", complete)?; + f.write_str("\n)")?; } Ok(()) } @@ -623,6 +708,120 @@ mod tests { use crate::parser::{ParseOptions, ParserContext}; use crate::statements::statement::Statement; + #[test] + fn test_rewritten_create_sql_budget() { + let parse = |sql| { + let mut statements = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::CreateTable(create) = statements.remove(0) else { + panic!("CREATE expected") + }; + create + }; + let create = parse( + "CREATE TABLE demo (ts TIMESTAMP TIME INDEX, host STRING) ENGINE=metric WITH (on_physical_table='phy')", + ); + let size = create.to_string().len(); + let mut remaining = size; + assert!(create.consume_sql_size_budget(&mut remaining)); + assert_eq!(remaining, 0); + assert!(!create.consume_sql_size_budget(&mut (size - 1))); + + for location in ["table", "fulltext", "skipping", "inverted"] { + for value_type in ["secret", "array", "struct"] { + for length in [10, 8192] { + let mut create = create.clone(); + let extensions = &mut create.columns[1].extensions; + let map = match location { + "table" => &mut create.options, + "fulltext" => extensions.fulltext_index_options.get_or_insert_default(), + "skipping" => extensions.skipping_index_options.get_or_insert_default(), + _ => extensions.inverted_index_options.get_or_insert_default(), + }; + let value = "private-value".repeat(length); + match value_type { + "secret" => map.insert("secret_access_key".into(), value.clone()), + "array" => map.insert_options("payload", vec![value.as_str()].into()), + _ => map.insert_options( + "payload", + crate::util::OptionValue::try_new(Expr::Struct { + values: vec![Expr::Value( + sqlparser::ast::Value::SingleQuotedString(value.clone()).into(), + )], + fields: vec![], + }) + .unwrap(), + ), + } + assert!(!create.to_string().contains(&value)); + assert_eq!( + create.consume_sql_size_budget(&mut 4096), + length == 10, + "{location}/{value_type}" + ); + } + } + } + let mut vector = parse("CREATE TABLE demo (ts TIMESTAMP TIME INDEX, v VECTOR(3))"); + vector.columns[1].column_def.options.push(ColumnOptionDef { + name: None, + option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(8192)), + }); + assert!(!vector.to_string().contains("private-value")); + assert!(!vector.consume_sql_size_budget(&mut 4096)); + vector.columns[1].column_def.options.clear(); + assert!(vector.consume_sql_size_budget(&mut 4096)); + if let DataType::Custom(_, tokens) = &mut vector.columns[1].column_def.data_type { + tokens[0] = format!("{}3", "0".repeat(8192)); + } + assert!(!vector.consume_sql_size_budget(&mut 4096)); + vector.columns[1] + .extensions + .vector_options + .as_mut() + .unwrap() + .insert("extra".into(), "x".into()); + assert!(!vector.consume_sql_size_budget(&mut 4096)); + + let mut create = create; + create.columns[1].extensions.json2_options = Some(Json2Options { + max_auto_expanded_paths: Some(1), + type_hints: vec![JsonTypeHint { + path: vec!["x".repeat(8192)], + data_type: DataType::Text, + inverted_index: true, + }], + }); + assert!(!create.consume_sql_size_budget(&mut 4096)); + create.columns[1].extensions.json2_options = None; + create.partitions = Some(Partitions { + column_list: vec![], + exprs: vec![Expr::Identifier(Ident::new("hidden"))], + }); + assert!(!create.consume_sql_size_budget(&mut 4096)); + } + + #[test] + fn test_json2_streamed_quote_escaping() { + let path = "a\\\"\"b"; + let options = Json2Options { + max_auto_expanded_paths: None, + type_hints: vec![JsonTypeHint { + path: vec![path.into()], + data_type: DataType::Text, + inverted_index: false, + }], + }; + assert_eq!( + options.to_string(), + format!("(\n \"{}\" TEXT\n )", path.replace('"', "\"\"")) + ); + } + #[test] fn test_display_create_table() { let sql = r"create table if not exists demo( diff --git a/src/sql/src/statements/option_map.rs b/src/sql/src/statements/option_map.rs index d562a61ce04..5077088b67a 100644 --- a/src/sql/src/statements/option_map.rs +++ b/src/sql/src/statements/option_map.rs @@ -139,26 +139,44 @@ impl OptionMap { } pub fn kv_pairs(&self) -> Vec { - let mut result = Vec::with_capacity(self.options.len() + self.secrets.len()); - for (k, v) in self - .options - .iter() - .filter_map(|(k, v)| v.as_string().map(|v| (k, v))) - { - if k.contains(".") { - result.push(format!("'{k}' = '{}'", v.escape_debug())); - } else { - result.push(format!("{k} = '{}'", v.escape_debug())); + self.sql_pairs(false).map(|pair| pair.to_string()).collect() + } + + /// Streams the same pairs as `kv_pairs`, or all actual values for size validation. + pub(crate) fn write_sql( + &self, + f: &mut dyn std::fmt::Write, + separator: &str, + prefix: &str, + complete: bool, + ) -> std::fmt::Result { + for (index, pair) in self.sql_pairs(complete).enumerate() { + if index > 0 { + f.write_str(separator)?; } + write!(f, "{prefix}{pair}")?; } - for (k, _) in self.secrets.iter() { - if k.contains(".") { - result.push(format!("'{k}' = '******'")); + Ok(()) + } + + fn sql_pairs(&self, complete: bool) -> impl Iterator> { + let options = self.options.iter().filter_map(move |(key, value)| { + let value = match value.as_string() { + Some(value) => Either::Right(value), + None if complete => Either::Left(value), + None => return None, + }; + Some(SqlPair { key, value }) + }); + let secrets = self.secrets.iter().map(move |(key, value)| SqlPair { + key, + value: Either::Right(if complete { + value.expose_secret() } else { - result.push(format!("{k} = '******'")); - } - } - result + "******" + }), + }); + options.chain(secrets) } pub fn entries(&self) -> impl Iterator)> { @@ -174,6 +192,25 @@ impl OptionMap { } } +struct SqlPair<'a> { + key: &'a str, + value: Either<&'a OptionValue, &'a str>, +} + +impl std::fmt::Display for SqlPair<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.key.contains('.') { + write!(f, "'{}'", self.key)?; + } else { + f.write_str(self.key)?; + } + match self.value { + Either::Left(value) => write!(f, " = {value}"), + Either::Right(value) => write!(f, " = '{}'", value.escape_debug()), + } + } +} + impl> From for OptionMap { fn from(value: I) -> Self { let mut result = OptionMap::default(); diff --git a/src/sql/src/util.rs b/src/sql/src/util.rs index c0c237cb0c0..a33c2576bbe 100644 --- a/src/sql/src/util.rs +++ b/src/sql/src/util.rs @@ -16,7 +16,6 @@ use std::collections::HashSet; use std::fmt::{Display, Formatter}; use std::ops::ControlFlow; -use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp}; use promql_parser::parser::{ Expr as PromExpr, MatrixSelector as PromMatrixSelector, VectorSelector as PromVectorSelector, @@ -161,12 +160,20 @@ impl Display for OptionValue { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if let Some(s) = self.as_string() { write!(f, "'{s}'") - } else if let Some(s) = self.as_list() { - write!( - f, - "[{}]", - s.into_iter().map(|x| format!("'{x}'")).join(", ") - ) + } else if let Expr::Array(array) = &self.0 + && array + .elem + .iter() + .all(|expr| Self::expr_as_string(expr).is_some()) + { + f.write_str("[")?; + for (index, expr) in array.elem.iter().enumerate() { + if index > 0 { + f.write_str(", ")?; + } + write!(f, "'{}'", Self::expr_as_string(expr).unwrap())?; + } + f.write_str("]") } else { write!(f, "'{}'", self.0) } @@ -667,6 +674,7 @@ fn extract_tables_from_sql_query(query: &sqlparser::ast::Query, names: &mut Hash #[cfg(test)] mod tests { + use itertools::Itertools; use sqlparser::tokenizer::Token; use super::*; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index eff6a8685f2..54a4e2d67b8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -517,6 +517,459 @@ async fn test_sql_skip_wal(distributed: bool) { check_http_skip_wal("sql", &cases, distributed).await; } +fn logical_ddl(name: &str) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (on_physical_table='phy.dot')" + ) +} + +#[tokio::test(flavor = "multi_thread")] +async fn logical_ddl_http_standalone() { + check_logical_ddl_http(false).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn logical_ddl_http_distributed() { + check_logical_ddl_http(true).await; +} + +#[tokio::test(flavor = "multi_thread")] +#[ignore = "manual logical DDL throughput measurement"] +async fn logical_ddl_http_benchmark() { + common_telemetry::init_default_ut_logging(); + let count: usize = std::env::var("GT_LOGICAL_DDL_BENCH_TABLES") + .unwrap_or_else(|_| "1000".to_string()) + .parse() + .unwrap(); + for (name, endpoint, batch_size) in [ + ("single", "/v1/sql", 1), + ("multi_sql", "/v1/sql", 128), + ("batch", "/v1/ddl/logical-tables", 128), + ] { + let mut instance = + MockInstanceImpl::Standalone( + tests_integration::standalone::GreptimeDbStandaloneBuilder::new(&format!( + "logical_ddl_bench_{name}" + )) + .with_event_recorder_options( + tests_integration::test_util::test_event_recorder_options(), + ) + .build() + .await, + ); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(instance.frontend()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + let response = execute_sql(&client, "CREATE TABLE \"phy.dot\" (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let statements = (0..count) + .map(|i| logical_ddl(&format!("bench_{i}"))) + .collect::>(); + let start = std::time::Instant::now(); + for chunk in statements.chunks(batch_size) { + let response = client + .post(endpoint) + .form(&[("sql", chunk.join(";"))]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + assert_eq!( + response.json::().await["output"] + .as_array() + .unwrap() + .len(), + chunk.len() + ); + } + let elapsed = start.elapsed(); + let requests = count.div_ceil(batch_size); + let procedures = if name == "batch" { requests } else { count }; + wait_for_event_data(&client, + "SELECT count(DISTINCT procedure_id), count(DISTINCT table_name) FROM greptime_private.events WHERE type='create_logical_tables' AND json_get_string(procedure_trigger, 'type')='Submitted' AND table_name LIKE 'bench_%'", + &format!("[[{procedures},{count}]]")).await; + common_telemetry::info!( + "Logical DDL: mode={name}, tables={count}, concurrency=1, HTTP requests={requests}, procedures={procedures}, elapsed={elapsed:?}" + ); + instance.shutdown().await; + } +} + +async fn check_logical_ddl_http(distributed: bool) { + common_telemetry::init_default_ut_logging(); + let events = tests_integration::test_util::test_event_recorder_options(); + let mut instance = if distributed { + MockInstanceImpl::Distributed( + tests_integration::cluster::GreptimeDbClusterBuilder::new("batch_logical_ddl") + .await + .with_event_recorder_options(events) + .with_datanodes(3) + .build(true) + .await, + ) + } else { + MockInstanceImpl::Standalone( + tests_integration::standalone::GreptimeDbStandaloneBuilder::new("batch_logical_ddl") + .with_event_recorder_options(events) + .build() + .await, + ) + }; + let frontend = instance.frontend(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(frontend.clone()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + assert_eq!( + client + .get("/v1/capabilities") + .send() + .await + .json::() + .await["metric_batch_ddl"], + 1 + ); + let response = execute_sql(&client, "CREATE TABLE \"phy.dot\" (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + + for sql in [ + String::new(), + format!("{}; SELECT 1", logical_ddl("rejected")), + format!("{}; CREATE VIEW v AS SELECT 1", logical_ddl("rejected")), + format!( + "{}; CREATE TABLE copy LIKE batch_0", + logical_ddl("rejected") + ), + format!("{}; CREATE TABLE copy AS SELECT 1", logical_ddl("rejected")), + logical_ddl("rejected").replace("WITH (", "WITH (physical_metric_table='', "), + format!( + "{}; CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX)", + logical_ddl("rejected") + ), + format!( + "{};{}", + logical_ddl("rejected"), + logical_ddl("public.rejected") + ), + format!("{};{}", logical_ddl("rejected"), logical_ddl("other.t")), + format!( + "{};{}", + logical_ddl("rejected"), + logical_ddl("other_physical").replace("phy.dot", "other") + ), + (0..129) + .map(|i| logical_ddl(&format!("rejected_{i}"))) + .collect::>() + .join(";"), + ] { + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(!response.status().is_success(), "{}", response.text().await); + } + validate_data( + "batch_rejected", + &client, + "SELECT count(*) FROM information_schema.tables WHERE table_name LIKE 'rejected%'", + "[[0]]", + ) + .await; + + let sql = (0..128) + .map(|i| logical_ddl(&format!("batch_{i}"))) + .collect::>() + .join(";"); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + assert_eq!( + response.json::().await["output"], + json!(vec![json!({"affectedrows": 0}); 128]) + ); + wait_for_event_data(&client, + "SELECT count(DISTINCT procedure_id), count(DISTINCT table_name) FROM greptime_private.events WHERE type='create_logical_tables' AND json_get_string(procedure_trigger, 'type')='Submitted' AND table_name LIKE 'batch_%'", + "[[1,128]]").await; + + let physical_id = frontend + .catalog_manager() + .table("greptime", "public", "phy.dot", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let mut ids = std::collections::HashSet::new(); + for i in 0..128 { + let table = frontend + .catalog_manager() + .table("greptime", "public", &format!("batch_{i}"), None) + .await + .unwrap() + .unwrap(); + let id = table.table_info().table_id(); + assert!(ids.insert(id)); + assert_ne!(id, physical_id); + let bound_physical_id = frontend + .table_metadata_manager() + .table_route_manager() + .get_physical_table_id(id) + .await + .unwrap(); + assert_eq!(bound_physical_id, physical_id); + assert_eq!( + table + .table_info() + .meta + .options + .extra_options + .get("on_physical_table") + .map(String::as_str), + Some("phy.dot") + ); + assert_eq!(table.schema().column_schemas().len(), 3); + } + let sql = logical_ddl("one_at_byte_limit"); + let padded = format!("{sql}{}", " ".repeat(1024 * 1024 - sql.len())); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", padded)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + + let response = execute_sql(&client, "CREATE DATABASE inherited WITH ('skip_wal'='true', ttl='1h', 'compaction.type'='twcs'); CREATE TABLE inherited.phy (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let inherited = logical_ddl("inherited.ordinary").replace("phy.dot", "phy"); + let response = execute_sql(&client, &inherited).await; + assert!(response.status().is_success(), "{}", response.text().await); + let sql = format!( + "{};{}", + logical_ddl("batched").replace("phy.dot", "phy"), + logical_ddl("explicit") + .replace("phy.dot", "phy") + .replace("WITH (", "WITH (skip_wal='false', ") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql.as_str()), ("db", "greptime-inherited")]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + let catalog = frontend.catalog_manager(); + let ordinary = catalog + .table("greptime", "inherited", "ordinary", None) + .await + .unwrap() + .unwrap(); + let batched = catalog + .table("greptime", "inherited", "batched", None) + .await + .unwrap() + .unwrap(); + let explicit = catalog + .table("greptime", "inherited", "explicit", None) + .await + .unwrap() + .unwrap(); + assert_eq!( + ordinary.table_info().meta.options, + batched.table_info().meta.options + ); + assert!(batched.table_info().meta.options.skip_wal); + assert!(!explicit.table_info().meta.options.skip_wal); + assert!( + !batched + .table_info() + .meta + .options + .extra_options + .contains_key("compaction.type") + ); + + let response = execute_sql(&client, "CREATE TABLE partitioned (ts TIMESTAMP TIME INDEX, host STRING PRIMARY KEY, val DOUBLE) PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') ENGINE=metric WITH (physical_metric_table='')").await; + assert!(response.status().is_success(), "{}", response.text().await); + let partitioned = logical_ddl("bad_partition") + .replace("phy.dot", "partitioned") + .replace( + "ENGINE=", + "PARTITION ON COLUMNS (host) (host < 'n', host >= 'n') ENGINE=", + ); + let response = execute_sql(&client, &partitioned).await; + assert!(!response.status().is_success()); + let sql = format!( + "{};{partitioned}", + logical_ddl("before_bad_partition").replace("phy.dot", "partitioned") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(!response.status().is_success()); + assert!( + catalog + .table("greptime", "public", "before_bad_partition", None) + .await + .unwrap() + .is_none() + ); + + let matching = partitioned + .replace("bad_partition", "matching_partition") + .replace("'n'", "'m'"); + let response = execute_sql(&client, &matching).await; + assert!(response.status().is_success(), "{}", response.text().await); + let sql = format!( + "{};{}", + matching.replace("matching_partition", "batch_matching_partition"), + logical_ddl("batch_inherited_partition").replace("phy.dot", "partitioned") + ); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + + let original_id = frontend + .catalog_manager() + .table("greptime", "public", "batch_0", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let sql = format!("{};{}", logical_ddl("batch_0"), logical_ddl("batch_new")); + for _ in 0..2 { + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", &sql)]) + .send() + .await; + assert!(response.status().is_success(), "{}", response.text().await); + } + let old = frontend + .catalog_manager() + .table("greptime", "public", "batch_0", None) + .await + .unwrap() + .unwrap(); + assert_eq!(original_id, old.table_info().table_id()); + let response = execute_sql( + &client, + "INSERT INTO batch_0 (ts, host, val) VALUES (1, 'a', 2.0); INSERT INTO batch_new (ts, host, val) VALUES (2, 'b', 3.0)", + ) + .await; + assert!(response.status().is_success(), "{}", response.text().await); + validate_data( + "batch_values", + &client, + "SELECT val FROM batch_0 UNION ALL SELECT val FROM batch_new ORDER BY val", + "[[2.0],[3.0]]", + ) + .await; + let ctx = session::context::QueryContext::arc(); + let provider = + auth::static_user_provider_from_option("static_user_provider:cmd:readonly:ro=password") + .unwrap(); + let user = auth::UserProvider::authenticate( + &provider, + auth::Identity::UserId("readonly", None), + auth::Password::PlainText("password".to_string().into()), + ) + .await + .unwrap(); + ctx.set_current_user(user); + frontend + .plugins() + .get_or_insert::(auth::DefaultPermissionChecker::arc); + let result = servers::query_handler::sql::SqlQueryHandler::create_logical_tables( + frontend.as_ref(), + &logical_ddl("read_only_denied"), + ctx, + ) + .await; + assert!(result.is_err()); + assert!( + frontend + .catalog_manager() + .table("greptime", "public", "read_only_denied", None) + .await + .unwrap() + .is_none() + ); + + struct FailAfterExecution(&'static str); + impl servers::interceptor::SqlQueryInterceptor for FailAfterExecution { + type Error = frontend::error::Error; + fn post_execute( + &self, + _: common_query::Output, + _: QueryContextRef, + ) -> frontend::error::Result { + use common_query::Output; + match self.0 { + "rows" => Ok(Output::new_with_affected_rows(1)), + "records" => Ok(Output::new_with_record_batches( + common_recordbatch::RecordBatches::empty(), + )), + "stream" => Ok(Output::new_with_stream(Box::pin( + common_recordbatch::RecordBatchStreamWrapper::new( + Arc::new(datatypes::schema::Schema::new(vec![])), + futures::stream::iter([Err::( + common_recordbatch::error::CreateRecordBatchesSnafu { + reason: "sensitive-option-value", + } + .build(), + )]), + ), + ))), + _ => frontend::error::InvalidSqlSnafu { + err_msg: "sensitive-option-value", + } + .fail(), + } + } + } + frontend + .plugins() + .insert::>(Arc::new( + FailAfterExecution("error"), + )); + for mode in ["error", "rows", "records", "stream"] { + frontend + .plugins() + .map_mut::, _, _>( + |plugin| *plugin.unwrap() = Arc::new(FailAfterExecution(mode)), + ); + let name = format!("created_before_error_{mode}"); + let response = client + .post("/v1/ddl/logical-tables") + .form(&[("sql", logical_ddl(&name))]) + .send() + .await; + assert!(!response.status().is_success(), "{mode}"); + assert!(!format!("{:?}", response.headers()).contains("sensitive-option-value")); + let body = response.json::().await; + assert!(body.get("output").is_none()); + assert_eq!(body["error"], "logical-table batch failed"); + assert!( + frontend + .catalog_manager() + .table("greptime", "public", &name, None) + .await + .unwrap() + .is_some() + ); + } + instance.shutdown().await; +} + pub async fn test_sql_api(store_type: StorageType) { let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "sql_api").await; let client = TestClient::new(app).await; From 9c69e381ae28dd86954530526239bbc6d4f17940 Mon Sep 17 00:00:00 2001 From: Aether <98370028+Aetherance@users.noreply.github.com> Date: Thu, 8 Oct 2026 07:15:23 +0000 Subject: [PATCH 098/133] fix(meta-client): bound heartbeat handshake timeout (#9438) * fix(meta-client): bound heartbeat handshake timeout Signed-off-by: Aetherance * chore(meta-client): fix test license header Signed-off-by: Aetherance * test(meta-client): prefix heartbeat tests with test Signed-off-by: Aetherance --------- Signed-off-by: Aetherance --- src/meta-client/src/client/heartbeat.rs | 29 ++-- src/meta-client/src/error.rs | 11 ++ .../tests/heartbeat_handshake_timeout.rs | 140 ++++++++++++++++++ 3 files changed, 170 insertions(+), 10 deletions(-) create mode 100644 src/meta-client/tests/heartbeat_handshake_timeout.rs diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index 96fc398556d..6dcc14c13f2 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -19,7 +19,7 @@ use std::time::Duration; use api::v1::meta::heartbeat_client::HeartbeatClient; use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, RequestHeader, Role}; use common_grpc::channel_manager::ChannelManager; -use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::distributed_time_constants::{BASE_HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT}; use common_meta::util; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{info, warn}; @@ -249,17 +249,26 @@ impl Inner { })?; let receiver = ReceiverStream::new(receiver); - let mut stream = leader - .heartbeat(receiver) - .await - .map_err(error::Error::from)? - .into_inner(); + let handshake = async { + let mut stream = leader + .heartbeat(receiver) + .await + .map_err(error::Error::from)? + .into_inner(); - let res = stream - .message() + let res = stream + .message() + .await + .map_err(error::Error::from)? + .context(error::CreateHeartbeatStreamSnafu)?; + Ok::<_, error::Error>((stream, res)) + }; + + let (stream, res) = tokio::time::timeout(HEARTBEAT_TIMEOUT, handshake) .await - .map_err(error::Error::from)? - .context(error::CreateHeartbeatStreamSnafu)?; + .with_context(|_| error::HeartbeatHandshakeTimeoutSnafu { + timeout: HEARTBEAT_TIMEOUT, + })??; let config = HeartbeatConfig::from_response(&res); diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 4623cfee33b..e15c83be1d9 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -86,6 +86,15 @@ pub enum Error { location: Location, }, + #[snafu(display("Heartbeat handshake timeout after {timeout:?}"))] + HeartbeatHandshakeTimeout { + timeout: std::time::Duration, + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: tokio::time::error::Elapsed, + }, + #[snafu(display("Invalid response header"))] InvalidResponseHeader { #[snafu(implicit)] @@ -150,6 +159,7 @@ impl ErrorExt for Error { | Error::MissingQueryContext { .. } | Error::SendHeartbeat { .. } | Error::CreateHeartbeatStream { .. } + | Error::HeartbeatHandshakeTimeout { .. } | Error::CreateChannel { .. } | Error::RetryTimesExceeded { .. } | Error::ConvertMetaConfig { .. } => StatusCode::Internal, @@ -168,6 +178,7 @@ impl ErrorExt for Error { fn retry_hint(&self) -> RetryHint { match self { Error::MetaServer { retry_hint, .. } => *retry_hint, + Error::HeartbeatHandshakeTimeout { .. } => RetryHint::Retryable, Error::InvalidResponseHeader { source, .. } | Error::ConvertMetaRequest { source, .. } | Error::ConvertMetaResponse { source, .. } diff --git a/src/meta-client/tests/heartbeat_handshake_timeout.rs b/src/meta-client/tests/heartbeat_handshake_timeout.rs new file mode 100644 index 00000000000..2ee43fa74d3 --- /dev/null +++ b/src/meta-client/tests/heartbeat_handshake_timeout.rs @@ -0,0 +1,140 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::pin::Pin; +use std::time::Duration; + +use api::v1::meta::heartbeat_server::{Heartbeat, HeartbeatServer}; +use api::v1::meta::{ + AskLeaderRequest, AskLeaderResponse, HeartbeatRequest, HeartbeatResponse, Peer, ResponseHeader, +}; +use common_error::ext::{ErrorExt, RetryHint}; +use common_meta::distributed_time_constants::HEARTBEAT_TIMEOUT; +use futures::Stream; +use meta_client::error::Error; +use meta_client::{MetaClientOptions, MetaClientRef, MetaClientType, create_meta_client}; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::codec::CompressionEncoding; +use tonic::{Request, Response, Status, Streaming}; + +type HeartbeatStream = + Pin> + Send + 'static>>; + +struct MockMetasrv { + addr: String, + send_first_heartbeat_response: bool, +} + +#[async_trait::async_trait] +impl Heartbeat for MockMetasrv { + type HeartbeatStream = HeartbeatStream; + + async fn heartbeat( + &self, + request: Request>, + ) -> Result, Status> { + let send_first_heartbeat_response = self.send_first_heartbeat_response; + let stream = futures::stream::once(async move { + let mut request = request.into_inner(); + request + .message() + .await? + .ok_or_else(|| Status::invalid_argument("missing heartbeat request"))?; + + if send_first_heartbeat_response { + Ok(HeartbeatResponse::default()) + } else { + futures::future::pending::>().await + } + }); + Ok(Response::new(Box::pin(stream))) + } + + async fn ask_leader( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(AskLeaderResponse { + header: Some(ResponseHeader::default()), + leader: Some(Peer { + id: 1, + addr: self.addr.clone(), + }), + })) + } +} + +async fn create_client(send_first_heartbeat_response: bool) -> MetaClientRef { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let mock = MockMetasrv { + addr: addr.to_string(), + send_first_heartbeat_response, + }; + + tokio::spawn(async move { + let service = HeartbeatServer::new(mock) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd) + .send_compressed(CompressionEncoding::Gzip) + .send_compressed(CompressionEncoding::Zstd); + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + + let options = MetaClientOptions { + metasrv_addrs: vec![addr.to_string()], + connect_timeout: Duration::from_millis(500), + ..Default::default() + }; + create_meta_client( + MetaClientType::Datanode { member_id: 1 }, + &options, + None, + None, + ) + .await + .expect("the mock answers leader discovery") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_heartbeat_handshake_succeeds_when_the_server_responds() { + let client = create_client(true).await; + + tokio::time::timeout(HEARTBEAT_TIMEOUT, client.heartbeat()) + .await + .expect("heartbeat handshake must complete before its deadline") + .expect("heartbeat handshake must succeed"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_heartbeat_handshake_times_out_when_the_stream_stays_silent() { + let client = create_client(false).await; + + let result = tokio::time::timeout(2 * HEARTBEAT_TIMEOUT, client.heartbeat()) + .await + .expect("heartbeat handshake must be bounded"); + let error = match result { + Ok(_) => panic!("a silent heartbeat stream must fail the handshake"), + Err(error) => error, + }; + assert!(matches!( + &error, + Error::HeartbeatHandshakeTimeout { timeout, .. } if *timeout == HEARTBEAT_TIMEOUT + )); + assert_eq!(error.retry_hint(), RetryHint::Retryable); +} From e2fd8c569521014910e13b6b1a196d2f5b3f5be4 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 07:23:22 +0000 Subject: [PATCH 099/133] docs(config): align example configs and config.md with actual defaults (#9460) * docs(config): align example configs and config.md with actual defaults Signed-off-by: Dennis Zhuang * test(cmd): compare example defaults inside optional sections Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 102 ++++++------ config/datanode.example.toml | 20 ++- config/flownode.example.toml | 15 +- config/frontend.example.toml | 30 ++-- config/metasrv.example.toml | 12 +- config/standalone.example.toml | 45 +++--- src/cmd/tests/load_config_test.rs | 255 ++++++++++++++++++++++++++++++ 7 files changed, 390 insertions(+), 89 deletions(-) diff --git a/config/config.md b/config/config.md index 1a7e8af1fe3..c538559e7cf 100644 --- a/config/config.md +++ b/config/config.md @@ -16,17 +16,16 @@ | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | | `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | -| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | -| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | +| `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | +| `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `init_regions_in_background` | Bool | `false` | Initialize all regions in the background during the startup.
By default, it provides services after all regions have been initialized. | | `init_regions_parallelism` | Integer | `16` | Parallelism of initializing regions. | -| `max_concurrent_queries` | Integer | `0` | The maximum concurrent queries allowed to be executed. Zero means unlimited. | | `enable_telemetry` | Bool | `true` | Enable telemetry to collect anonymous usage data. Enabled by default. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | | `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | | `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | @@ -112,7 +111,7 @@ | `wal.enable_log_recycle` | Bool | `true` | Whether to reuse logically truncated log files.
**It's only used when the provider is `raft_engine`**. | | `wal.prefill_log_files` | Bool | `false` | Whether to pre-create log files on start up.
**It's only used when the provider is `raft_engine`**. | | `wal.sync_period` | String | `5s` | Duration for fsyncing log files.
**It's only used when the provider is `raft_engine`**. | -| `wal.recovery_parallelism` | Integer | `2` | Parallelism during WAL recovery. | +| `wal.recovery_parallelism` | Integer | Auto | Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. | | `wal.broker_endpoints` | Array | -- | The Kafka broker endpoints.
**It's only used when the provider is `kafka`**. | | `wal.connect_timeout` | String | `3s` | The connect timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.timeout` | String | `5s` | The total request timeout for kafka client.
**It's only used when the provider is `kafka`**. | @@ -142,7 +141,7 @@ | `procedure.max_running_procedures` | Integer | `128` | Max running procedures.
The maximum number of procedures that can be running at the same time.
If the number of running procedures exceeds this limit, the procedure will be rejected. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -176,7 +175,7 @@ | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | | `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | -| `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | +| `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | | `region_engine.mito.worker_request_batch_size` | Integer | `64` | Max batch size for a worker to handle requests. | | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | @@ -260,17 +259,18 @@ | `logging.tracing_sample_ratio` | -- | Unset | The percentage of tracing will be sampled and exported.
Valid range `[0, 1]`, 1 means all traces are sampled, 0 means all traces are not sampled, the default value is 1.
ratio > 1 are treated as 1. Fractions < 0 are treated as 0 | | `logging.tracing_sample_ratio.default_ratio` | Float | `1.0` | -- | | `slow_query` | -- | -- | The slow query log options. | -| `slow_query.enable` | Bool | `false` | Whether to enable slow query log. | -| `slow_query.record_type` | String | Unset | The record type of slow queries. It can be `system_table` or `log`. | -| `slow_query.threshold` | String | Unset | The threshold of slow query. | -| `slow_query.sample_ratio` | Float | Unset | The sampling ratio of slow query log. The value should be in the range of (0, 1]. | +| `slow_query.enable` | Bool | `true` | Whether to enable slow query log. | +| `slow_query.record_type` | String | `system_table` | The record type of slow queries. It can be `system_table` or `log`.
If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`.
If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. | +| `slow_query.threshold` | String | `30s` | The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. | +| `slow_query.sample_ratio` | Float | `1.0` | The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. | +| `slow_query.ttl` | String | `90d` | The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. | | `tracing` | -- | -- | The tracing options. Only effect when compiled with `tokio-console` feature. | | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `pipeline` | -- | -- | The pipeline options. | | `pipeline.cache_ttl` | String | `10s` | Time to live of the local pipeline cache. Default is `10s`. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | All | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -285,13 +285,13 @@ | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | | `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | -| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | -| `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | -| `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | +| `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL:
`static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | +| `max_in_flight_write_bytes` | String | 0 | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | +| `write_bytes_exhausted_policy` | String | `wait` | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When synchronous Prometheus, 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`. | @@ -302,9 +302,9 @@ | `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.
Defaults to `127.0.0.1:4006`. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:4001` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:4001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | -| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | | `grpc.enable_cors` | Bool | `false` | Enable CORS for gRPC-Web clients in browsers. Disabled by default. | | `grpc.cors_allowed_origins` | Array | Unset | Origins allowed by gRPC CORS. An empty list allows any origin. | | `grpc.max_connection_age` | String | Unset | The maximum connection age for gRPC connection.
The value can be a human-readable time string. For example: `10m` for ten minutes or `1h` for one hour.
Refer to https://grpc.io/docs/guides/keepalive/ for more details. | @@ -313,11 +313,11 @@ | `grpc.tls.cert_path` | String | Unset | Certificate file path. | | `grpc.tls.key_path` | String | Unset | Private key file path. | | `grpc.tls.watch` | Bool | `false` | Watch for Certificate and key file change and auto reload.
For now, gRPC tls config does not support auto reload. | -| `internal_grpc` | -- | -- | The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend. | -| `internal_grpc.bind_addr` | String | `127.0.0.1:4010` | The address to bind the gRPC server. | -| `internal_grpc.server_addr` | String | `127.0.0.1:4010` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `internal_grpc.bind_addr`. | +| `internal_grpc` | -- | -- | The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend.
The internal gRPC server starts only when this section is set. | +| `internal_grpc.bind_addr` | String | Unset | The address to bind the gRPC server.
If unset, it falls back to `127.0.0.1:4001`, the same address as `grpc.bind_addr`. | +| `internal_grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
Unlike `grpc.server_addr`, it is not detected automatically; set it explicitly. | | `internal_grpc.runtime_size` | Integer | `8` | The number of server worker threads. | -| `internal_grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `internal_grpc.flight_compression` | String | `arrow_ipc` | Compression mode for frontend side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | | `internal_grpc.tls` | -- | -- | internal gRPC server TLS options, see `mysql.tls` section. | | `internal_grpc.tls.mode` | String | `disable` | TLS mode. | | `internal_grpc.tls.cert_path` | String | Unset | Certificate file path. | @@ -385,7 +385,7 @@ | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | | `query.allow_query_fallback` | Bool | `false` | Whether to allow query fallback when push down optimize fails.
Default to false, meaning when push down optimize failed, return error msg | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -393,7 +393,7 @@ | `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | -| `datanode.client.connect_timeout` | String | `10s` | -- | +| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. | | `datanode.client.tcp_nodelay` | Bool | `true` | -- | | `datanode.client.max_recv_message_size` | String | `512MB` | The maximum receive message size for the gRPC client. | | `datanode.client.max_send_message_size` | String | `512MB` | The maximum send message size for the gRPC client. | @@ -426,7 +426,7 @@ | `pipeline.cache_ttl` | String | `10s` | Time to live of the frontend-local pipeline cache. A pipeline created or deleted on
another frontend takes effect on this one after at most this duration. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | All | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | ### Metasrv @@ -438,7 +438,7 @@ | `store_key_prefix` | String | `""` | If it's not empty, the metasrv will store all data with this key prefix. | | `backend` | String | `etcd_store` | The datastore for meta server.
Available values:
- `etcd_store` (default value)
- `memory_store`
- `postgres_store`
- `mysql_store` | | `meta_table_name` | String | `greptime_metakv` | Table name in RDS to store metadata. Effect when using a RDS kvbackend.
**Only used when backend is `postgres_store`.** | -| `meta_schema_name` | String | `greptime_schema` | Optional PostgreSQL schema for metadata table and election table name qualification.
When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public),
set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`.
**Only used when backend is `postgres_store`.** | +| `meta_schema_name` | String | Unset | Optional PostgreSQL schema for metadata table and election table name qualification.
When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public),
set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`.
**Only used when backend is `postgres_store`.**
Unset by default; this example sets it to `greptime_schema`. | | `auto_create_schema` | Bool | `true` | Automatically create PostgreSQL schema if it doesn't exist.
When enabled, the system will execute `CREATE SCHEMA IF NOT EXISTS `
before creating metadata tables. This is useful in production environments where
manual schema creation may be restricted.
Default is true.
Note: The PostgreSQL user must have CREATE SCHEMA permission for this to work.
**Only used when backend is `postgres_store`.** | | `meta_election_lock_id` | Integer | `1` | Advisory lock id in PostgreSQL for election. Effect when using PostgreSQL as kvbackend
Only used when backend is `postgres_store`. | | `selector` | String | `round_robin` | Datanode selector type.
- `round_robin` (default value)
- `lease_based`
- `load_based`
For details, please see "https://docs.greptime.com/developer-guide/metasrv/selector". | @@ -449,9 +449,9 @@ | `heartbeat_interval` | String | `3s` | Base heartbeat interval for calculating distributed time constants.
The frontend heartbeat interval is 6 times of the base heartbeat interval.
The flownode/datanode heartbeat interval is 1 times of the base heartbeat interval.
e.g., If the base heartbeat interval is 3s, the frontend heartbeat interval is 18s, the flownode/datanode heartbeat interval is 3s.
Heartbeat intervals are negotiated from metasrv during handshake; local node configs do not override this. | | `enable_telemetry` | Bool | `true` | Whether to enable greptimedb telemetry. Enabled by default. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `backend_tls` | -- | -- | TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends)
When using etcd, PostgreSQL, or MySQL as metadata store, you can configure TLS here

Note: if TLS is configured in both this section and the `store_addrs` connection string, the
settings here will override the TLS settings in `store_addrs`. | | `backend_tls.mode` | String | `prefer` | TLS mode, refer to https://www.postgresql.org/docs/current/libpq-ssl.html
- "disable" - No TLS
- "prefer" (default) - Try TLS, fallback to plain
- "require" - Require TLS
- "verify_ca" - Require TLS and verify CA
- "verify_full" - Require TLS and verify hostname | | `backend_tls.cert_path` | String | `""` | Path to client certificate file (for client authentication)
Like "/path/to/client.crt" | @@ -463,7 +463,7 @@ | `backend_client.connect_timeout` | String | `3s` | The connect timeout for backend client. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:3002` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:3002` | The communication server address for the frontend and datanode to connect to metasrv.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The communication server address for the frontend and datanode to connect to metasrv.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | @@ -485,7 +485,7 @@ | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | | `datanode.client.timeout` | String | `10s` | Operation timeout. | -| `datanode.client.connect_timeout` | String | `10s` | Connect server timeout. | +| `datanode.client.connect_timeout` | String | 1s | Connect server timeout. Defaults to `1s`; this example sets it to `10s`. | | `datanode.client.tcp_nodelay` | Bool | `true` | `TCP_NODELAY` option for accepted connections. | | `datanode.client.max_recv_message_size` | String | `512MB` | The maximum receive message size for the gRPC client. | | `datanode.client.max_send_message_size` | String | `512MB` | The maximum send message size for the gRPC client. | @@ -507,7 +507,7 @@ | `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.
**It's only used when the provider is `kafka`**. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | All | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `stats_persistence` | -- | -- | Configuration options for the stats persistence. | | `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.
Set to `0s` to disable stats persistence.
Default is `0s`.
If you want to enable stats persistence, set the TTL to a value greater than 0.
It is recommended to set a small value, e.g., `3h`. | | `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.
The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. | @@ -552,22 +552,22 @@ | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `grpc` | -- | -- | The gRPC server options. | | `grpc.bind_addr` | String | `127.0.0.1:3001` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:3001` | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | | `grpc.runtime_size` | Integer | `8` | The number of server worker threads. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | -| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for datanode side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `none` | +| `grpc.flight_compression` | String | `arrow_ipc` | Compression mode for datanode side Arrow IPC service. Available options:
- `none`: disable all compression
- `transport`: only enable gRPC transport compression (zstd)
- `arrow_ipc`: only enable Arrow IPC compression (lz4)
- `all`: enable all compression.
Default to `arrow_ipc`. | | `grpc.tls` | -- | -- | gRPC server TLS options, see `mysql.tls` section. | | `grpc.tls.mode` | String | `disable` | TLS mode. | | `grpc.tls.cert_path` | String | Unset | Certificate file path. | | `grpc.tls.key_path` | String | Unset | Private key file path. | | `grpc.tls.watch` | Bool | `false` | Watch for Certificate and key file change and auto reload.
For now, gRPC tls config does not support auto reload. | | `runtime` | -- | -- | The runtime options. | -| `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | -| `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | -| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | -| `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. | +| `runtime.global_rt_size` | Integer | Auto | The number of threads to execute the runtime for global read operations.
Defaults to max(num_cpus, 2). | +| `runtime.compact_rt_size` | Integer | Auto | The number of threads to execute compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.compact_rt_max_blocking_threads` | Integer | Auto | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.query_rt_size` | Integer | Auto | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | +| `runtime.ingest_rt_size` | Integer | Auto | The number of threads to execute datanode ingestion operations.
Defaults to max(num_cpus, 2). | | `runtime.experimental_workload_scheduler` | -- | -- | Experimental weighted, work-conserving query/write task scheduler. | | `runtime.experimental_workload_scheduler.enable` | Bool | `false` | Enable when concurrent queries and writes interfere with each other—for example, when long-running queries increase ingestion latency.
The weights set their relative runtime shares while both are backlogged. Disabled by default. | | `runtime.experimental_workload_scheduler.query_weight` | Integer | `2` | Relative query share while both query and write workloads are backlogged. | @@ -593,7 +593,7 @@ | `wal.enable_log_recycle` | Bool | `true` | Whether to reuse logically truncated log files.
**It's only used when the provider is `raft_engine`**. | | `wal.prefill_log_files` | Bool | `false` | Whether to pre-create log files on start up.
**It's only used when the provider is `raft_engine`**. | | `wal.sync_period` | String | `5s` | Duration for fsyncing log files.
**It's only used when the provider is `raft_engine`**. | -| `wal.recovery_parallelism` | Integer | `2` | Parallelism during WAL recovery. | +| `wal.recovery_parallelism` | Integer | Auto | Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. | | `wal.broker_endpoints` | Array | -- | The Kafka broker endpoints.
**It's only used when the provider is `kafka`**. | | `wal.connect_timeout` | String | `3s` | The connect timeout for kafka client.
**It's only used when the provider is `kafka`**. | | `wal.timeout` | String | `5s` | The total request timeout for kafka client.
**It's only used when the provider is `kafka`**. | @@ -611,7 +611,7 @@ | `wal.on_corrupted_segment` | String | `skip` | What a read does with a segment that still does not decode after a second fetch, because its checksum does not match or its content disagrees with its footer entry.
- `skip`: the segment is skipped and recorded as a WAL hole of its region, a metric is incremented and a warning is logged; the other regions of the object are unaffected (the default).
- `fail`: the read fails, so the region does not open.
**It's only used when the provider is `experimental_object_store`**. | | `query` | -- | -- | The query engine options. | | `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | @@ -644,7 +644,7 @@ | `storage.http_client.skip_ssl_validation` | Bool | `false` | To skip the ssl verification
**Security Notice**: Setting `skip_ssl_validation = true` disables certificate verification, making connections vulnerable to man-in-the-middle attacks. Only use this in development or trusted private networks. | | `[[region_engine]]` | -- | -- | The region engine options. You can configure multiple region engines.
Each engine type (mito, file, metric) may appear only once; duplicates cause startup to fail. | | `region_engine.mito` | -- | -- | The Mito engine options. | -| `region_engine.mito.num_workers` | Integer | `8` | Number of region workers. | +| `region_engine.mito.num_workers` | Integer | Auto | Number of region workers. Defaults to half of the available CPU cores, rounded up. | | `region_engine.mito.worker_channel_size` | Integer | `128` | Request channel size of each worker. | | `region_engine.mito.worker_request_batch_size` | Integer | `64` | Max batch size for a worker to handle requests. | | `region_engine.mito.manifest_checkpoint_distance` | Integer | `10` | Number of meta action updated to trigger a new checkpoint for the manifest. | @@ -761,9 +761,9 @@ | `flow.batching_mode.frontend_tls.client_cert_path` | String | Unset | Client Certificate file path. | | `flow.batching_mode.frontend_tls.client_key_path` | String | Unset | Client Private key file path. | | `grpc` | -- | -- | The gRPC server options. | -| `grpc.bind_addr` | String | `127.0.0.1:6800` | The address to bind the gRPC server. | -| `grpc.server_addr` | String | `127.0.0.1:6800` | The address advertised to the metasrv,
and used for connections from outside the host | -| `grpc.runtime_size` | Integer | `2` | The number of server worker threads. | +| `grpc.bind_addr` | String | 127.0.0.1:3004 | The address to bind the gRPC server. Defaults to `127.0.0.1:3004`; this example sets it to `127.0.0.1:6800`. | +| `grpc.server_addr` | String | Unset | The address advertised to the metasrv, and used for connections from outside the host.
If left empty or unset, the server will automatically use the IP address of the first network interface
on the host, with the same port number as the one specified in `grpc.bind_addr`. | +| `grpc.runtime_size` | Integer | 8 | The number of server worker threads. Defaults to 8; this example sets it to 2. | | `grpc.max_recv_message_size` | String | `512MB` | The maximum receive message size for gRPC server. | | `grpc.max_send_message_size` | String | `512MB` | The maximum send message size for gRPC server. | | `http` | -- | -- | The HTTP server options. | @@ -797,7 +797,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `query` | -- | -- | -- | | `query.parallelism` | Integer | `1` | Parallelism of the query engine for query sent by flownode.
Default to 1, so it won't use too much cpu or memory | -| `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | +| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. | | `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). | | `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. | | `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 627a5d3288a..413afb609c6 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -45,6 +45,7 @@ bind_addr = "127.0.0.1:3001" ## The address advertised to the metasrv, and used for connections from outside the host. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:3001" ## The number of server worker threads. runtime_size = 8 @@ -57,7 +58,7 @@ max_send_message_size = "512MB" ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" ## gRPC server TLS options, see `mysql.tls` section. @@ -80,16 +81,24 @@ watch = false ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## The number of threads to execute datanode query operations. ## Defaults to max(num_cpus - 1, 2). +## @toml2docs:none-default="Auto" #+ query_rt_size = 7 ## The number of threads to execute datanode ingestion operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ ingest_rt_size = 8 ## Experimental weighted, work-conserving query/write task scheduler. @@ -176,7 +185,8 @@ prefill_log_files = false ## **It's only used when the provider is `raft_engine`**. sync_period = "5s" -## Parallelism during WAL recovery. +## Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. +## @toml2docs:none-default="Auto" recovery_parallelism = 2 ## The Kafka broker endpoints. @@ -330,9 +340,10 @@ parallelism = 0 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -503,7 +514,8 @@ skip_ssl_validation = false ## The Mito engine options. [region_engine.mito] -## Number of region workers. +## Number of region workers. Defaults to half of the available CPU cores, rounded up. +## @toml2docs:none-default="Auto" #+ num_workers = 8 ## Request channel size of each worker. diff --git a/config/flownode.example.toml b/config/flownode.example.toml index b0a53bee54a..fab31fed510 100644 --- a/config/flownode.example.toml +++ b/config/flownode.example.toml @@ -44,12 +44,16 @@ node_id = 14 ## The gRPC server options. [grpc] -## The address to bind the gRPC server. +## The address to bind the gRPC server. Defaults to `127.0.0.1:3004`; this example sets it to `127.0.0.1:6800`. +## @toml2docs:none-default="127.0.0.1:3004" bind_addr = "127.0.0.1:6800" -## The address advertised to the metasrv, -## and used for connections from outside the host +## The address advertised to the metasrv, and used for connections from outside the host. +## If left empty or unset, the server will automatically use the IP address of the first network interface +## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:6800" -## The number of server worker threads. +## The number of server worker threads. Defaults to 8; this example sets it to 2. +## @toml2docs:none-default="8" runtime_size = 2 ## The maximum receive message size for gRPC server. max_recv_message_size = "512MB" @@ -153,9 +157,10 @@ parallelism = 1 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 2f2c15cfc57..21e7bbb40f3 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -20,8 +20,8 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's -## native password handshake; a MySQL client must send the password in cleartext for such users. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL: +## `static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats. ## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. ## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or ## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. @@ -33,22 +33,26 @@ default_column_prefix = "greptime" ## Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). ## Set to 0 to disable the limit. Default: "0" (unlimited) -## @toml2docs:none-default +## @toml2docs:none-default="0" #+ max_in_flight_write_bytes = "1GB" ## Policy when write bytes quota is exhausted. ## Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" -## @toml2docs:none-default #+ write_bytes_exhausted_policy = "wait" ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. @@ -86,6 +90,7 @@ bind_addr = "127.0.0.1:4001" ## The address advertised to the metasrv, and used for connections from outside the host. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:4001" ## The number of server worker threads. runtime_size = 8 @@ -94,7 +99,7 @@ runtime_size = 8 ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" ## Enable CORS for gRPC-Web clients in browsers. Disabled by default. #+ enable_cors = false @@ -125,12 +130,15 @@ key_path = "" watch = false ## The internal gRPC server options. Internal gRPC port for nodes inside cluster to access frontend. +## The internal gRPC server starts only when this section is set. [internal_grpc] ## The address to bind the gRPC server. +## If unset, it falls back to `127.0.0.1:4001`, the same address as `grpc.bind_addr`. +## @toml2docs:none-default bind_addr = "127.0.0.1:4010" ## The address advertised to the metasrv, and used for connections from outside the host. -## If left empty or unset, the server will automatically use the IP address of the first network interface -## on the host, with the same port number as the one specified in `internal_grpc.bind_addr`. +## Unlike `grpc.server_addr`, it is not detected automatically; set it explicitly. +## @toml2docs:none-default server_addr = "127.0.0.1:4010" ## The number of server worker threads. runtime_size = 8 @@ -139,7 +147,7 @@ runtime_size = 8 ## - `transport`: only enable gRPC transport compression (zstd) ## - `arrow_ipc`: only enable Arrow IPC compression (lz4) ## - `all`: enable all compression. -## Default to `none` +## Default to `arrow_ipc`. flight_compression = "arrow_ipc" ## internal gRPC server TLS options, see `mysql.tls` section. @@ -357,9 +365,10 @@ allow_query_fallback = false ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans (only applies to datanodes). +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -390,6 +399,8 @@ memory_pool_size = "50%" [datanode] ## Datanode client options. [datanode.client] +## Connect server timeout. Defaults to `1s`; this example sets it to `10s`. +## @toml2docs:none-default="1s" connect_timeout = "10s" tcp_nodelay = true ## The maximum receive message size for the gRPC client. @@ -494,4 +505,5 @@ ttl = "90d" ## Event types to record. Current available event type: `admin_function`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. +## @toml2docs:none-default="All" #+ event_types = ["admin_function"] diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index b923c5a549c..fe4c38d02f9 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -35,6 +35,8 @@ meta_table_name = "greptime_metakv" ## When PostgreSQL public schema is not writable (e.g., PostgreSQL 15+ with restricted public), ## set this to a writable schema. GreptimeDB will use `meta_schema_name`.`meta_table_name`. ## **Only used when backend is `postgres_store`.** +## Unset by default; this example sets it to `greptime_schema`. +## @toml2docs:none-default meta_schema_name = "greptime_schema" ## Automatically create PostgreSQL schema if it doesn't exist. @@ -88,11 +90,16 @@ node_max_idle_time = "24hours" ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends) @@ -140,6 +147,7 @@ bind_addr = "127.0.0.1:3002" ## The communication server address for the frontend and datanode to connect to metasrv. ## If left empty or unset, the server will automatically use the IP address of the first network interface ## on the host, with the same port number as the one specified in `grpc.bind_addr`. +## @toml2docs:none-default server_addr = "127.0.0.1:3002" ## The number of server worker threads. runtime_size = 8 @@ -206,7 +214,8 @@ acceptable_heartbeat_pause = "10000ms" ## Operation timeout. timeout = "10s" -## Connect server timeout. +## Connect server timeout. Defaults to `1s`; this example sets it to `10s`. +## @toml2docs:none-default="1s" connect_timeout = "10s" ## `TCP_NODELAY` option for accepted connections. @@ -332,6 +341,7 @@ ttl = "90d" ## `reconcile_database`, `reconcile_table`, `reconcile_logical_tables`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. +## @toml2docs:none-default="All" #+ event_types = ["region_migration"] ## Configuration options for the stats persistence. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 6b60c4482d7..7e1f8dcf6b0 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -20,8 +20,8 @@ default_column_prefix = "greptime" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", ## "mysql_native_password:", ## "pg_scram_sha256::::" -## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's -## native password handshake; a MySQL client must send the password in cleartext for such users. +## "pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but such users cannot log in over MySQL: +## `static_user_provider` and `watch_file_user_provider` only negotiate `mysql_native_password`, which cannot verify these formats. ## "mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all. ## PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or ## "mysql_native_password", PostgreSQL falls back to cleartext password auth for every user. @@ -33,12 +33,11 @@ default_column_prefix = "greptime" ## Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). ## Set to 0 to disable the limit. Default: "0" (unlimited) -## @toml2docs:none-default +## @toml2docs:none-default="0" #+ max_in_flight_write_bytes = "1GB" ## Policy when write bytes quota is exhausted. ## Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" -## @toml2docs:none-default #+ write_bytes_exhausted_policy = "wait" ## Initialize all regions in the background during the startup. @@ -48,20 +47,22 @@ init_regions_in_background = false ## Parallelism of initializing regions. init_regions_parallelism = 16 -## The maximum concurrent queries allowed to be executed. Zero means unlimited. -max_concurrent_queries = 0 - ## Enable telemetry to collect anonymous usage data. Enabled by default. #+ enable_telemetry = true ## The runtime options. #+ [runtime] ## The number of threads to execute the runtime for global read operations. +## Defaults to max(num_cpus, 2). +## @toml2docs:none-default="Auto" #+ global_rt_size = 8 ## The number of threads to execute compact operations. +## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. ## Defaults to max(num_cpus / 2, 2). +## @toml2docs:none-default="Auto" #+ compact_rt_max_blocking_threads = 4 ## Experimental weighted, work-conserving query/write task scheduler. @@ -344,7 +345,8 @@ prefill_log_files = false ## **It's only used when the provider is `raft_engine`**. sync_period = "5s" -## Parallelism during WAL recovery. +## Parallelism during WAL recovery. Defaults to the number of CPU cores; this example sets it to 2. +## @toml2docs:none-default="Auto" recovery_parallelism = 2 ## The Kafka broker endpoints. @@ -538,9 +540,10 @@ parallelism = 0 ## Memory pool size for query execution operators (aggregation, sorting, join). ## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%"). -## Setting it to 0 disables the limit (unbounded, default behavior). +## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`. ## When this limit is reached, queries will fail with ResourceExhausted error. ## NOTE: This does NOT limit memory used by table scans. +## @toml2docs:none-default="0" memory_pool_size = "50%" ## Experimental memory pool allocation policy: @@ -720,7 +723,8 @@ skip_ssl_validation = false ## The Mito engine options. [region_engine.mito] -## Number of region workers. +## Number of region workers. Defaults to half of the available CPU cores, rounded up. +## @toml2docs:none-default="Auto" #+ num_workers = 8 ## Request channel size of each worker. @@ -1062,19 +1066,21 @@ default_ratio = 1.0 ## The slow query log options. [slow_query] ## Whether to enable slow query log. -#+ enable = false +enable = true ## The record type of slow queries. It can be `system_table` or `log`. -## @toml2docs:none-default -#+ record_type = "system_table" +## If `system_table` is selected, the slow queries will be recorded in a system table `greptime_private.slow_queries`. +## If `log` is selected, the slow queries will be logged in a log file `greptimedb-slow-queries.*`. +record_type = "system_table" -## The threshold of slow query. -## @toml2docs:none-default -#+ threshold = "10s" +## The threshold of slow query. It can be human readable time string, for example: `10s`, `100ms`, `1s`. +threshold = "30s" -## The sampling ratio of slow query log. The value should be in the range of (0, 1]. -## @toml2docs:none-default -#+ sample_ratio = 1.0 +## The sampling ratio of slow query log. The value should be in the range of (0, 1]. For example, `0.1` means 10% of the slow queries will be logged and `1.0` means all slow queries will be logged. +sample_ratio = 1.0 + +## The TTL of the `slow_queries` system table. Default is `90d` when `record_type` is `system_table`. +ttl = "90d" ## The tracing options. Only effect when compiled with `tokio-console` feature. #+ [tracing] @@ -1099,6 +1105,7 @@ ttl = "90d" ## `reconcile_database`, `reconcile_table`, `reconcile_logical_tables`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. +## @toml2docs:none-default="All" #+ event_types = ["create_database"] ## The memory options. diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index b56299f9916..1f4c69c8938 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -476,6 +476,261 @@ fn test_load_standalone_example_config() { similar_asserts::assert_eq!(options, expected); } +/// Renders an example config with its `#+` options enabled, as `config.md` +/// lists them. With `drop_none_default`, options annotated with +/// `@toml2docs:none-default` are removed, because `config.md` does not present +/// their example value as the default. +fn render_example(example: &str, drop_none_default: bool) -> String { + let mut rendered = String::new(); + let mut none_default = false; + for line in example.lines() { + let line = line.trim_start(); + if line.starts_with("##") { + none_default |= line.contains("@toml2docs:none-default"); + continue; + } + let line = line.strip_prefix("#+").map_or(line, str::trim_start); + let is_option = !line.starts_with('[') && line.contains('='); + if !(is_option && drop_none_default && none_default) { + rendered.push_str(line); + rendered.push('\n'); + } + none_default = false; + } + rendered +} + +/// Flattens leaf values into dotted paths. Arrays of tables such as +/// `region_engine` are merged, since each element is keyed by its engine name. +fn leaf_values(value: &toml::Value, prefix: &str, leaves: &mut Vec<(String, toml::Value)>) { + match value { + toml::Value::Table(table) => { + for (key, value) in table { + let path = if prefix.is_empty() { + key.clone() + } else { + format!("{prefix}.{key}") + }; + leaf_values(value, &path, leaves); + } + } + toml::Value::Array(array) + if array.iter().all(toml::Value::is_table) && !array.is_empty() => + { + for element in array { + leaf_values(element, prefix, leaves); + } + } + _ => leaves.push((prefix.to_string(), value.clone())), + } +} + +/// Removes the option at a dotted path produced by [`leaf_values`]. +fn remove_leaf(value: &mut toml::Value, path: &str) -> bool { + match value { + toml::Value::Array(array) => array.iter_mut().any(|element| remove_leaf(element, path)), + toml::Value::Table(table) => match path.split_once('.') { + None => table.remove(path).is_some(), + Some((head, rest)) => table + .get_mut(head) + .is_some_and(|value| remove_leaf(value, rest)), + }, + _ => false, + } +} + +fn load_as_leaves( + config: Option<&str>, +) -> common_config::error::Result> +where + T: Configurable + serde::Serialize, +{ + let file = config.map(|config| { + let mut file = create_named_temp_file(); + file.write_all(config.as_bytes()).unwrap(); + file + }); + let options = GreptimeOptions::::load_layered_options( + file.as_ref().map(|file| file.path().to_str().unwrap()), + "", + )?; + let mut leaves = Vec::new(); + leaf_values(&toml::Value::try_from(&options).unwrap(), "", &mut leaves); + Ok(leaves.into_iter().collect()) +} + +/// Returns the default of an option inside the example's other settings, so +/// that options in sections absent by default still have one to compare with. +/// Removes the closest enclosing table that still loads without the option. +fn default_within_example(full: &toml::Value, path: &str) -> Option +where + T: Configurable + serde::Serialize, +{ + let mut removed = path; + loop { + let mut without = full.clone(); + assert!(remove_leaf(&mut without, removed)); + if let Ok(defaults) = load_as_leaves::(Some(&toml::to_string(&without).unwrap())) { + return defaults.get(path).cloned(); + } + removed = removed + .rsplit_once('.') + .unwrap_or_else(|| panic!("the example does not load without `{path}`")) + .0; + } +} + +/// Checks that every option in an example config exists, and that every value +/// `config.md` presents as a default is the actual default. +/// +/// `variants` lists tagged sections, e.g. `("wal", "provider", ["kafka"])`, +/// whose options only deserialize under a specific tag value. +fn check_example_defaults(name: &str, variants: &[(&str, &str, &[&str])]) -> Vec +where + T: Configurable + serde::Serialize, +{ + let example = std::fs::read_to_string(common_test_util::find_workspace_path(&format!( + "config/{name}" + ))) + .unwrap(); + let full: toml::Value = toml::from_str(&render_example(&example, false)).unwrap(); + let documented: toml::Value = toml::from_str(&render_example(&example, true)).unwrap(); + + // (options loaded from the example, options loaded without config) per variant + let mut loaded = vec![( + load_as_leaves::(Some(&toml::to_string(&full).unwrap())).unwrap(), + load_as_leaves::(None).unwrap(), + )]; + for (section, tag, values) in variants { + for value in *values { + let mut variant = full.clone(); + variant[section][tag] = toml::Value::String(value.to_string()); + loaded.push(( + load_as_leaves::(Some(&toml::to_string(&variant).unwrap())).unwrap(), + load_as_leaves::(Some(&format!("[{section}]\n{tag} = \"{value}\"\n"))).unwrap(), + )); + } + } + + // Deserialize-only fields, invisible in the serialized options. + const SKIP_SERIALIZING: &[&str] = &[ + "storage.access_key_id", + "storage.secret_access_key", + "storage.access_key_secret", + "storage.account_name", + "storage.account_key", + "storage.credential_path", + "storage.credential", + "region_engine.mito.inverted_index.intermediate_path", + ]; + + let mut errors = Vec::new(); + let mut full_leaves = Vec::new(); + leaf_values(&full, "", &mut full_leaves); + for (path, _) in full_leaves { + if SKIP_SERIALIZING.contains(&path.as_str()) + || (path == "tracing.tokio_console_addr" && !cfg!(feature = "tokio-console")) + { + continue; + } + if !loaded + .iter() + .any(|(example, _)| example.contains_key(&path)) + { + errors.push(format!("{name}: `{path}` is not a known option")); + } + } + + // Options left unset by default and resolved where they are used. + let resolved_at_runtime: [(&str, toml::Value); 7] = [ + ( + "logging.dir", + format!("{DEFAULT_DATA_HOME}/{DEFAULT_LOGGING_DIR}").into(), + ), + ("logging.otlp_endpoint", DEFAULT_OTLP_HTTP_ENDPOINT.into()), + ("logging.otlp_export_protocol", "http".into()), + ("logging.tracing_sample_ratio.default_ratio", 1.0.into()), + ("wal.sync_period", "5s".into()), + ( + "storage.copy_root", + format!("{DEFAULT_DATA_HOME}/copy").into(), + ), + ("flow.batching_mode.frontend_tls.enabled", false.into()), + ]; + + let mut documented_leaves = Vec::new(); + leaf_values(&documented, "", &mut documented_leaves); + for (path, documented_value) in documented_leaves { + if SKIP_SERIALIZING.contains(&path.as_str()) { + continue; + } + if let Some((_, resolved)) = resolved_at_runtime.iter().find(|(p, _)| *p == path) { + if &documented_value != resolved { + errors.push(format!( + "{name}: `{path}` is documented as default `{documented_value}`, but resolves to `{resolved}`" + )); + } + continue; + } + let mut pairs = loaded + .iter() + .filter_map(|(example, default)| Some((example.get(&path)?, default.get(&path)?))) + .map(|(example, default)| (example.clone(), default.clone())) + .collect::>(); + if pairs.is_empty() { + // The option sits in a section that is absent without config, or + // is itself optional. + match default_within_example::(&full, &path) { + Some(default) => pairs.push((loaded[0].0[&path].clone(), default)), + None => { + errors.push(format!( + "{name}: `{path}` has no default to compare with; annotate it with \ + `@toml2docs:none-default` or list its runtime default in this test" + )); + continue; + } + } + } + if pairs.iter().all(|(example, default)| example != default) { + let (example, default) = &pairs[0]; + errors.push(format!( + "{name}: `{path}` is documented as default `{example}`, but the default is `{default}`; \ + fix the value or annotate it with `@toml2docs:none-default`" + )); + } + } + errors +} + +// The examples document HDFS storage options. +#[cfg(feature = "hdfs-object-store")] +#[test] +fn test_example_configs_document_actual_defaults() { + const WAL: (&str, &str, &[&str]) = ("wal", "provider", &["kafka", "experimental_object_store"]); + const STORAGE: (&str, &str, &[&str]) = + ("storage", "type", &["S3", "Oss", "Azblob", "Gcs", "Hdfs"]); + + let mut errors = + check_example_defaults::("standalone.example.toml", &[WAL, STORAGE]); + errors.extend(check_example_defaults::( + "frontend.example.toml", + &[], + )); + errors.extend(check_example_defaults::( + "datanode.example.toml", + &[WAL, STORAGE], + )); + errors.extend(check_example_defaults::( + "metasrv.example.toml", + &[("wal", "provider", &["kafka"])], + )); + errors.extend(check_example_defaults::( + "flownode.example.toml", + &[], + )); + assert!(errors.is_empty(), "{}", errors.join("\n")); +} + #[test] fn test_load_removed_histogram_options() { for enabled in [false, true] { From 90f117b957a70180bf102a19ca47810314900958 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 07:30:12 +0000 Subject: [PATCH 100/133] fix(partition): preserve region mapping during repartition (#9458) Signed-off-by: Dennis Zhuang --- .../src/test_util/partition_manager.rs | 124 ++++++++++++++++++ src/partition/src/manager.rs | 26 ++-- 2 files changed, 141 insertions(+), 9 deletions(-) diff --git a/src/operator/src/test_util/partition_manager.rs b/src/operator/src/test_util/partition_manager.rs index e51ac49051e..f80e2cc4fe2 100644 --- a/src/operator/src/test_util/partition_manager.rs +++ b/src/operator/src/test_util/partition_manager.rs @@ -368,3 +368,127 @@ async fn test_logical_split_rows_rejects_physical_only_partition_column() { partition::error::Error::UndefinedColumn { ref column, .. } if column == "b" )); } + +#[tokio::test] +async fn test_partition_routing_with_empty_default_expression() { + use datatypes::arrow::array::{ArrayRef, RecordBatch, StringArray}; + use datatypes::value::Value; + use partition::expr::col; + use store_api::storage::RegionId; + + for (transition, target_count) in [(true, 1), (true, 3), (false, 3)] { + let backend = Arc::new(common_meta::kv_backend::memory::MemoryKvBackend::new()); + let route_cache = test_new_table_route_cache(backend.clone()); + let manager = PartitionRuleManager::new( + backend.clone(), + route_cache.clone(), + test_new_partition_info_cache(route_cache), + ); + let value = |s: &str| Value::String(s.into()); + let mut exprs = vec![ + col("host").lt(value("36z0")), + col("host") + .gt_eq(value("36z0")) + .and(col("host").lt(value("7Ez0"))), + col("host") + .gt_eq(value("7Ez0")) + .and(col("host").lt(value("BMz0"))), + col("host").gt_eq(value("BMz0")), + ]; + if target_count == 1 { + exprs = vec![ + col("host").lt(value("36z0")), + col("host").gt_eq(value("36z0")), + ]; + } + let routes = exprs + .into_iter() + .enumerate() + .map(|(i, expr)| RegionRoute { + region: Region { + id: RegionId::new(1027, i as u32), + partition_expr: if transition && i == 0 { + String::new() + } else { + expr.as_json_str().unwrap() + }, + ..Default::default() + }, + leader_peer: Some(Peer::new(1, "")), + ..Default::default() + }) + .collect(); + let info = new_test_table_info_with_columns( + 1027, + "physical", + vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ], + vec![0], + ); + let metadata_manager = TableMetadataManager::new(backend); + metadata_manager + .create_table_metadata( + info.clone(), + TableRouteValue::physical(routes), + new_test_region_wal_options((0..=target_count).collect()), + ) + .await + .unwrap(); + let mut logical_info = info.clone(); + logical_info.ident.table_id = 1028; + logical_info.name = "logical".into(); + metadata_manager + .create_table_metadata( + logical_info.clone(), + TableRouteValue::logical(1027), + HashMap::new(), + ) + .await + .unwrap(); + + let hosts = [None, Some("1"), Some("5"), Some("8"), Some("animi")]; + let expected_regions = if target_count == 1 { + [0, 0, 1, 1, 1] + } else { + [0, 0, 1, 2, 3] + }; + let batch = RecordBatch::try_from_iter([( + "host", + Arc::new(StringArray::from(hosts.to_vec())) as ArrayRef, + )]) + .unwrap(); + for table_info in [&info, &logical_info] { + let (rule, _) = manager.find_table_partition_rule(table_info).await.unwrap(); + for (host, expected) in hosts.into_iter().zip(expected_regions) { + assert_eq!( + rule.find_region(&[host.map(value).unwrap_or(Value::Null)]) + .unwrap(), + expected + ); + } + let splits = rule.split_record_batch(&batch).unwrap(); + assert_eq!(splits.len(), target_count as usize + 1); + for (region, mask) in splits { + let selected = mask + .array() + .iter() + .enumerate() + .filter_map(|(i, selected)| (selected == Some(true)).then_some(i)) + .collect::>(); + let expected = expected_regions + .iter() + .enumerate() + .filter_map(|(i, expected)| (*expected == region).then_some(i)) + .collect::>(); + assert_eq!(selected, expected); + } + } + } +} diff --git a/src/partition/src/manager.rs b/src/partition/src/manager.rs index 89bb61c4174..68c9586e908 100644 --- a/src/partition/src/manager.rs +++ b/src/partition/src/manager.rs @@ -206,17 +206,25 @@ impl PartitionRuleManager { .iter() .map(|r| (r.id.region_number(), r.partition_expr_version)) .collect::>>(); - let regions = partition_info + let (mut regions, exprs): (Vec<_>, Vec<_>) = partition_info .partitions .iter() - .map(|x| x.id.region_number()) - .collect::>(); - let exprs = partition_info - .partitions - .iter() - .filter_map(|x| x.partition_expr.as_ref()) - .cloned() - .collect::>(); + .filter_map(|partition| { + partition + .partition_expr + .as_ref() + .map(|expr| (partition.id.region_number(), expr.clone())) + }) + .unzip(); + // Empty expressions belong to the default region. Keep it after the + // explicit regions so expression indices still identify the right region. + regions.extend( + partition_info + .partitions + .iter() + .filter(|partition| partition.partition_expr.is_none()) + .map(|partition| partition.id.region_number()), + ); let partition_rule = Arc::new(MultiDimPartitionRule::try_new( partition_columns, regions, From 526dd48bda49ed2d6d19b71d1cf4da1766cb6225 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 8 Oct 2026 07:43:25 +0000 Subject: [PATCH 101/133] fix(promql): restore exact filtering for nullable label matchers (#9418) * fix(promql): restore exact filtering for nullable label matchers Signed-off-by: Dennis Zhuang * test: verify compound filter semantics and last-row selection Signed-off-by: Dennis Zhuang * test: consolidate label filter regression coverage Signed-off-by: Dennis Zhuang * refactor: separate filter evaluation and timestamp adaptation Signed-off-by: Dennis Zhuang * refactor: inline predicate and timestamp tests Signed-off-by: Dennis Zhuang * refactor: keep direct IN predicates outside simple filters Signed-off-by: Dennis Zhuang * fix: coerce partition bounds before simple filtering Signed-off-by: Dennis Zhuang * test: cover partition bound overflow after timestamp widening Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/recordbatch/src/filter.rs | 1215 ++--------------- .../recordbatch/src/filter/predicate.rs | 944 +++++++++++++ src/common/recordbatch/src/filter/tests.rs | 72 + .../recordbatch/src/filter/timestamp.rs | 439 ++++++ src/mito2/src/engine/prune_test.rs | 6 +- src/mito2/src/engine/scan_test.rs | 209 +++ src/mito2/src/memtable/time_series.rs | 8 +- src/mito2/src/read/scan_region.rs | 93 +- src/mito2/src/sst/parquet/prefilter.rs | 3 +- src/mito2/src/sst/parquet/reader.rs | 16 +- src/query/src/dummy_catalog.rs | 40 +- src/query/src/promql/planner.rs | 25 +- src/query/src/promql/planner/test/delta.rs | 30 +- .../standalone/common/promql/regex.result | 84 +- .../cases/standalone/common/promql/regex.sql | 10 + 15 files changed, 2064 insertions(+), 1130 deletions(-) create mode 100644 src/common/recordbatch/src/filter/predicate.rs create mode 100644 src/common/recordbatch/src/filter/tests.rs create mode 100644 src/common/recordbatch/src/filter/timestamp.rs diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs index 26c28aa3fa0..60ff399601f 100644 --- a/src/common/recordbatch/src/filter.rs +++ b/src/common/recordbatch/src/filter.rs @@ -12,215 +12,174 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Util record batch stream wrapper that can perform precise filter. +//! Precise single-column predicates and record batch filtering. + +mod predicate; +#[cfg(test)] +mod tests; +mod timestamp; use std::sync::Arc; -use common_time::timestamp::div_mod_units; use datafusion::error::Result as DfResult; use datafusion::logical_expr::{Expr, Literal, Operator}; use datafusion::physical_plan::PhysicalExpr; -use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datafusion_common::arrow::array::{ArrayRef, Datum}; use datafusion_common::arrow::buffer::BooleanBuffer; -use datafusion_common::arrow::compute::kernels::cmp; -use datafusion_common::cast::{as_boolean_array, as_null_array, as_string_array}; +use datafusion_common::cast::{as_boolean_array, as_null_array}; use datafusion_common::{DataFusionError, ScalarValue, internal_err}; -use datatypes::arrow::array::{ - Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray, - RecordBatch, StringArrayType, -}; +use datatypes::arrow::array::{Array, BooleanArray, BooleanBufferBuilder, RecordBatch}; use datatypes::arrow::compute::filter_record_batch; -use datatypes::arrow::datatypes::{DataType, TimeUnit, UInt32Type}; -use datatypes::arrow::error::ArrowError; -use datatypes::compute::or_kleene; +use datatypes::arrow::datatypes::DataType; use datatypes::data_type::{ConcreteDataType, DataType as _}; use datatypes::value::Value; use datatypes::vectors::VectorRef; -use regex::Regex; use snafu::ResultExt; -use crate::error::{ArrowComputeSnafu, Result, ToArrowScalarSnafu, UnsupportedOperationSnafu}; +pub use self::predicate::{regexp_is_match_dictionary, regexp_is_match_scalar}; +pub use self::timestamp::timestamp_scalar_value; +use crate::error::{Result, ToArrowScalarSnafu}; +use crate::filter::predicate::SimplePredicate; -/// An inplace expr evaluator for simple filter. Only support -/// - `col` `op` `literal` -/// - `literal` `op` `col` -/// -/// And the `op` is one of `=`, `!=`, `>`, `>=`, `<`, `<=`, -/// or regex operators: `~`, `~*`, `!~`, `!~*`. -/// -/// This struct contains normalized predicate expr. In the form of -/// `col` `op` `literal` where the `col` is provided from input. +/// Evaluates supported single-column predicates with precompiled regular expressions. +/// Compound predicates preserve SQL null semantics until conversion to a scan mask. #[derive(Debug, Clone)] pub struct SimpleFilterEvaluator { - /// Name of the referenced column. column_name: String, - /// The literal value. - literal: Scalar, - /// The operator. - op: Operator, - /// Only used when the operator is `Or`-chain. - literal_list: Vec>, - /// Pre-compiled regex. - /// Only used when the operator is regex operators. - /// If the regex is empty, it is also `None`. - regex: Option, - /// Whether the regex is negative. - regex_negative: bool, + predicate: SimplePredicate, } impl SimpleFilterEvaluator { pub fn new(column_name: String, lit: T, op: Operator) -> Option { - match op { + if !matches!( + op, Operator::Eq - | Operator::NotEq - | Operator::Lt - | Operator::LtEq - | Operator::Gt - | Operator::GtEq => {} - _ => return None, + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ) { + return None; } - - let Expr::Literal(val, _) = lit.lit() else { + let Expr::Literal(value, _) = lit.lit() else { return None; }; - Some(Self { column_name, - literal: val.to_scalar().ok()?, - op, - literal_list: vec![], - regex: None, - regex_negative: false, + predicate: SimplePredicate::Comparison { + literal: value.to_scalar().ok()?, + op, + regex: None, + regex_negative: false, + }, }) } pub fn try_new(predicate: &Expr) -> Option { - match predicate { - Expr::BinaryExpr(binary) => { - // check if the expr is in the supported form - match binary.op { - Operator::Eq - | Operator::NotEq - | Operator::Lt - | Operator::LtEq - | Operator::Gt - | Operator::GtEq - | Operator::RegexMatch - | Operator::RegexIMatch - | Operator::RegexNotMatch - | Operator::RegexNotIMatch => {} - Operator::Or => { - let lhs = Self::try_new(&binary.left)?; - let rhs = Self::try_new(&binary.right)?; - if lhs.column_name != rhs.column_name - || !matches!(lhs.op, Operator::Eq | Operator::Or) - || !matches!(rhs.op, Operator::Eq | Operator::Or) - { - return None; - } - let mut list = vec![]; - let placeholder_literal = lhs.literal.clone(); - // above check guarantees the op is either `Eq` or `Or` - if matches!(lhs.op, Operator::Or) { - list.extend(lhs.literal_list); - } else { - list.push(lhs.literal); - } - if matches!(rhs.op, Operator::Or) { - list.extend(rhs.literal_list); - } else { - list.push(rhs.literal); - } - return Some(Self { - column_name: lhs.column_name, - literal: placeholder_literal, - op: Operator::Or, - literal_list: list, - regex: None, - regex_negative: false, - }); - } - _ => return None, - } - - // swap the expr if it is in the form of `literal` `op` `col` - let mut op = binary.op; - let (lhs, rhs) = match (&*binary.left, &*binary.right) { - (Expr::Column(col), Expr::Literal(lit, _)) => (col, lit), - (Expr::Literal(lit, _), Expr::Column(col)) => { - // safety: The previous check ensures the operator is able to swap. - op = op.swap().unwrap(); - (col, lit) - } - _ => return None, - }; - - let (regex, regex_negative) = Self::maybe_build_regex(op, rhs).ok()?; - let literal = rhs.to_scalar().ok()?; - Some(Self { - column_name: lhs.name.clone(), - literal, - op, - literal_list: vec![], - regex, - regex_negative, - }) - } - _ => None, - } + Self::try_new_with_column_type(predicate, &|_| None) } - /// Get the name of the referenced column. + /// Recognizes single-column filters, including lossless string dictionary casts + /// when the caller can resolve the column's type. + pub fn try_new_with_column_type( + predicate: &Expr, + column_type: &impl Fn(&str) -> Option, + ) -> Option { + let columns = predicate.column_refs(); + if columns.len() != 1 { + return None; + } + let column_name = columns.into_iter().next()?.name.clone(); + let column_type = column_type(&column_name); + Some(Self { + column_name, + predicate: SimplePredicate::try_new(predicate, column_type.as_ref())?, + }) + } + + /// Returns the referenced column name. pub fn column_name(&self) -> &str { &self.column_name } pub fn is_eq(&self) -> bool { - matches!(self.op, Operator::Eq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Eq, + .. + } + ) } - pub fn is_not_eq(&self) -> bool { - matches!(self.op, Operator::NotEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::NotEq, + .. + } + ) } - pub fn is_lt(&self) -> bool { - matches!(self.op, Operator::Lt) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Lt, + .. + } + ) } - pub fn is_lt_eq(&self) -> bool { - matches!(self.op, Operator::LtEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::LtEq, + .. + } + ) } - pub fn is_gt(&self) -> bool { - matches!(self.op, Operator::Gt) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::Gt, + .. + } + ) } - pub fn is_gt_eq(&self) -> bool { - matches!(self.op, Operator::GtEq) + matches!( + self.predicate, + SimplePredicate::Comparison { + op: Operator::GtEq, + .. + } + ) } - /// Returns true if this filter represents an `OR` chain of equality comparisons, e.g. - /// `col = lit1 OR col = lit2 ...`. + /// Returns true for an OR chain of equality comparisons. pub fn is_or_eq_chain(&self) -> bool { - matches!(self.op, Operator::Or) + matches!(self.predicate, SimplePredicate::OrEqChain { .. }) } - /// Returns the literal as a [`Value`]. It returns `None` if the literal can't be converted. + /// Returns the literal of a bare comparison, without discarding compound conditions. pub fn literal_value(&self) -> Option { - let array = self.literal.get().0; - let scalar = ScalarValue::try_from_array(array, 0).ok()?; - Value::try_from(scalar).ok() + let SimplePredicate::Comparison { literal, .. } = &self.predicate else { + return None; + }; + Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok() } - /// Returns the literal list as a list of [`Value`]s. It returns `None` if any literal can't be - /// converted. + /// Returns equality-OR literals, or None for other predicates or unsupported values. pub fn literal_list_values(&self) -> Option> { - self.literal_list + let SimplePredicate::OrEqChain { literals } = &self.predicate else { + return None; + }; + literals .iter() - .map(|scalar| { - let array = scalar.get().0; - let scalar = ScalarValue::try_from_array(array, 0).ok()?; - Value::try_from(scalar).ok() + .map(|literal| { + Value::try_from(ScalarValue::try_from_array(literal.get().0, 0).ok()?).ok() }) .collect() } @@ -229,8 +188,7 @@ impl SimpleFilterEvaluator { let input = input .to_scalar() .with_context(|_| ToArrowScalarSnafu { v: input.clone() })?; - let result = self.evaluate_datum(&input, 1)?; - Ok(result.value(0)) + Ok(self.evaluate_datum(&input, 1)?.value(0)) } pub fn evaluate_array(&self, input: &ArrayRef) -> Result { @@ -238,200 +196,27 @@ impl SimpleFilterEvaluator { } pub fn evaluate_vector(&self, input: &VectorRef) -> Result { - self.evaluate_datum(&input.to_arrow_array(), input.len()) + self.evaluate_array(&input.to_arrow_array()) } - fn evaluate_datum(&self, input: &impl Datum, input_len: usize) -> Result { - let result = match self.op { - Operator::Eq => cmp::eq(input, &self.literal), - Operator::NotEq => cmp::neq(input, &self.literal), - Operator::Lt => cmp::lt(input, &self.literal), - Operator::LtEq => cmp::lt_eq(input, &self.literal), - Operator::Gt => cmp::gt(input, &self.literal), - Operator::GtEq => cmp::gt_eq(input, &self.literal), - Operator::RegexMatch => self.regex_match(input), - Operator::RegexIMatch => self.regex_match(input), - Operator::RegexNotMatch => self.regex_match(input), - Operator::RegexNotIMatch => self.regex_match(input), - Operator::Or => { - // OR operator stands for OR-chained EQs (or INLIST in other words) - let mut result: BooleanArray = vec![false; input_len].into(); - for literal in &self.literal_list { - let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?; - result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?; - } - Ok(result) - } - _ => { - return UnsupportedOperationSnafu { - reason: format!("{:?}", self.op), - } - .fail(); - } - }; - result - .context(ArrowComputeSnafu) - .map(|array| boolean_array_to_scan_mask(&array).values().clone()) + fn evaluate_datum(&self, input: &impl Datum, len: usize) -> Result { + let result = self.predicate.evaluate(input, len)?; + Ok(boolean_array_to_scan_mask(&result).values().clone()) } - /// Builds a regex pattern from a scalar value and operator. - /// Returns the `(regex, negative)` and if successful. - /// - /// Returns `Err` if - /// - the value is not a string - /// - the regex pattern is invalid - /// - /// The regex is `None` if - /// - the operator is not a regex operator - /// - the pattern is empty - fn maybe_build_regex( - operator: Operator, - value: &ScalarValue, - ) -> Result<(Option, bool), ArrowError> { - let (ignore_case, negative) = match operator { - Operator::RegexMatch => (false, false), - Operator::RegexIMatch => (true, false), - Operator::RegexNotMatch => (false, true), - Operator::RegexNotIMatch => (true, true), - _ => return Ok((None, false)), - }; - let flag = if ignore_case { Some("i") } else { None }; - let regex = value - .try_as_str() - .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))? - .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?; - let pattern = match flag { - Some(flag) => format!("(?{flag}){regex}"), - None => regex.to_string(), - }; - if pattern.is_empty() { - Ok((None, negative)) - } else { - Regex::new(pattern.as_str()) - .map_err(|e| { - ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}")) - }) - .map(|regex| (Some(regex), negative)) - } - } - - fn regex_match(&self, input: &impl Datum) -> std::result::Result { - let array = input.get().0; - - // Try to cast to StringArray first - if let Ok(string_array) = as_string_array(array) { - let mut result = regexp_is_match_scalar(string_array, self.regex.as_ref())?; - if self.regex_negative { - result = datatypes::compute::not(&result)?; - } - return Ok(result); - } - - // Try to cast to StringDictionaryArray - if let Some(dict_array) = array.as_any().downcast_ref::>() { - let mut result = regexp_is_match_dictionary(dict_array, self.regex.as_ref())?; - if self.regex_negative { - result = datatypes::compute::not(&result)?; - } - return Ok(result); - } - - Err(ArrowError::CastError(format!( - "Cannot cast {:?} to StringArray or StringDictionaryArray", - array.data_type() - ))) - } - - /// Casts the filter's timestamp literal into the unit of the `target` - /// timestamp type, so it can be evaluated against a column stored in that - /// unit (e.g. an old-unit SST after the time index unit was widened). - /// - /// When the literal is not representable in `target`'s unit (e.g. - /// `ts = 7_000_500us` against a millisecond column), the outcome keeps the - /// row set instead of rounding the literal: `=` prunes / `!=` matches, and - /// inequalities strengthen the operator (e.g. `>= 2_500_500us` becomes - /// `> 2500ms`) so the excluded boundary row stays excluded. - /// - /// Returns `None` when the filter doesn't compare against a tz-naive - /// timestamp literal, or `target` is not a timestamp type. + /// Adapts timestamp predicates to an older SST's timestamp unit without rounding + /// away boundary rows. Nonrepresentable equality literals cannot match a row. pub fn cast_timestamp_unit(&self, target: &ConcreteDataType) -> Option { - let target_unit = match target.as_arrow_type() { - DataType::Timestamp(unit, _) => unit, - _ => return None, + let DataType::Timestamp(target_unit, _) = target.as_arrow_type() else { + return None; }; - - // An `OR` chain of equalities (an IN list): convert each literal - // with `=` semantics; literals not representable in the target unit - // cannot match any row and simply drop out of the chain. - if self.op == Operator::Or { - let mut literal_list = Vec::with_capacity(self.literal_list.len()); - for literal in &self.literal_list { - let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; - let Some((value, unit)) = timestamp_scalar_parts(&scalar) else { - continue; - }; - let Some(value) = value else { - continue; - }; - let cast = div_mod_units(value, unit.into(), target_unit.into())?; - if cast.remainder == 0 { - literal_list.push(timestamp_scalar(cast.quotient, target_unit)?); - } - } - if literal_list.is_empty() { - return Some(TimestampUnitCast::Pruned); - } - let literal = literal_list[0].clone(); - return Some(TimestampUnitCast::Filter(Self { + Some(match self.predicate.cast_timestamp_unit(target_unit)? { + SimplePredicate::Constant(Some(true)) => TimestampUnitCast::Matched, + SimplePredicate::Constant(_) => TimestampUnitCast::Pruned, + predicate => TimestampUnitCast::Filter(Self { column_name: self.column_name.clone(), - literal, - op: Operator::Or, - literal_list, - regex: None, - regex_negative: false, - })); - } - - let scalar = ScalarValue::try_from_array(self.literal.get().0, 0).ok()?; - let (value, unit) = timestamp_scalar_parts(&scalar)?; - let Some(value) = value else { - // A NULL literal never compares true (arrow comparison - // semantics: null results are filtered out). - return Some(TimestampUnitCast::Pruned); - }; - if unit == target_unit { - return Some(TimestampUnitCast::Filter(self.clone())); - } - let cast = div_mod_units(value, unit.into(), target_unit.into())?; - let divisible = cast.remainder == 0; - let literal = timestamp_scalar(cast.quotient, target_unit)?; - let filter = |op: Operator| Self { - column_name: self.column_name.clone(), - literal: literal.clone(), - op, - literal_list: vec![], - regex: None, - regex_negative: false, - }; - - Some(match self.op { - Operator::Eq if divisible => TimestampUnitCast::Filter(filter(Operator::Eq)), - Operator::Eq => TimestampUnitCast::Pruned, - Operator::NotEq if divisible => TimestampUnitCast::Filter(filter(Operator::NotEq)), - Operator::NotEq => TimestampUnitCast::Matched, - // v > L holds exactly for v > quotient, whether or not L is - // representable in the target unit. - Operator::Gt => TimestampUnitCast::Filter(filter(Operator::Gt)), - Operator::GtEq if divisible => TimestampUnitCast::Filter(filter(Operator::GtEq)), - // L strictly between two target values: v >= L is v > quotient. - Operator::GtEq => TimestampUnitCast::Filter(filter(Operator::Gt)), - Operator::Lt if divisible => TimestampUnitCast::Filter(filter(Operator::Lt)), - // L strictly between two target values: v < L is v <= quotient. - Operator::Lt => TimestampUnitCast::Filter(filter(Operator::LtEq)), - // v <= L holds exactly for v <= quotient. - Operator::LtEq => TimestampUnitCast::Filter(filter(Operator::LtEq)), - // Regex predicates don't apply to timestamps. - _ => return None, + predicate, + }), }) } } @@ -449,35 +234,6 @@ pub enum TimestampUnitCast { Matched, } -/// Extracts the value and unit from a tz-naive timestamp scalar. -fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option, TimeUnit)> { - let (value, unit, timezone) = match scalar { - ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz), - ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz), - ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz), - ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz), - _ => return None, - }; - // A timezone-aware literal doesn't compare against a tz-naive column; - // leave it to the caller instead of guessing the intended semantics. - (timezone.is_none()).then_some((value, unit)) -} - -/// Builds a tz-naive timestamp literal for `value` in `unit`. -pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue { - match unit { - TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None), - TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None), - TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None), - TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None), - } -} - -/// Builds a one-element scalar array holding `value` in `unit`. -fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option> { - timestamp_scalar_value(value, unit).to_scalar().ok() -} - /// Evaluate the predicate on the input [RecordBatch], and return a new [RecordBatch]. /// Copy from datafusion::physical_plan::src::filter.rs pub fn batch_filter( @@ -520,726 +276,3 @@ fn boolean_array_to_scan_mask(array: &BooleanArray) -> BooleanArray { } BooleanArray::new(values.into(), None) } - -/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar()) -/// with pre-compiled regex. -/// See for the implementation details. -pub fn regexp_is_match_scalar<'a, S>( - array: &'a S, - regex: Option<&Regex>, -) -> Result -where - &'a S: StringArrayType<'a>, -{ - let null_bit_buffer = array.nulls().map(|x| x.inner().sliced()); - let mut result = BooleanBufferBuilder::new(array.len()); - - if let Some(re) = regex { - for i in 0..array.len() { - let value = array.value(i); - result.append(re.is_match(value)); - } - } else { - result.append_n(array.len(), true); - } - - let buffer = result.into(); - let data = unsafe { - ArrayData::new_unchecked( - DataType::Boolean, - array.len(), - None, - null_bit_buffer, - 0, - vec![buffer], - vec![], - ) - }; - - Ok(BooleanArray::from(data)) -} - -/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray. -/// Iterates through dictionary keys to get string values and applies regex matching. -pub fn regexp_is_match_dictionary( - dict_array: &DictionaryArray, - regex: Option<&Regex>, -) -> Result { - // Get the string values from the dictionary - let string_values = dict_array - .values() - .as_any() - .downcast_ref::() - .ok_or_else(|| { - ArrowError::CastError("Dictionary values must be StringArray".to_string()) - })?; - - // Dictionary logical nulls include both null keys and keys whose dictionary value is null. - let logical_nulls = dict_array.logical_nulls(); - let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced()); - let mut result = BooleanBufferBuilder::new(dict_array.len()); - - if let Some(re) = regex { - let keys = dict_array.keys().values(); - for i in 0..dict_array.len() { - if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) { - result.append(false); - } else { - let key = keys[i] as usize; - let string_value = string_values.value(key); - result.append(re.is_match(string_value)); - } - } - } else { - result.append_n(dict_array.len(), true); - } - - let buffer = result.into(); - let data = unsafe { - ArrayData::new_unchecked( - DataType::Boolean, - dict_array.len(), - None, - null_bit_buffer, - 0, - vec![buffer], - vec![], - ) - }; - - Ok(BooleanArray::from(data)) -} - -#[cfg(test)] -mod test { - - use std::sync::Arc; - - use datafusion::execution::context::ExecutionProps; - use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; - use datafusion::logical_expr::{BinaryExpr, col, lit}; - use datafusion::physical_expr::create_physical_expr; - use datafusion_common::{Column, DFSchema}; - use datatypes::arrow::array::{TimestampMillisecondArray, TimestampNanosecondArray}; - use datatypes::arrow::datatypes::{DataType, Field, Schema}; - - use super::*; - - #[test] - fn unsupported_filter_op() { - // `+` is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Plus, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // two literal is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(1.lit()), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // two column is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(Expr::Column(Column::from_name("bar"))), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - - // compound expr is not supported - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1.lit()), - })), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); - } - - #[test] - fn supported_filter_op() { - // equal - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1.lit()), - }); - let _ = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - // swap operands - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(1.lit()), - op: Operator::Lt, - right: Box::new(Expr::Column(Column::from_name("foo"))), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - assert_eq!(evaluator.op, Operator::Gt); - assert_eq!(evaluator.column_name, "foo".to_string()); - } - - #[test] - fn run_on_array() { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Eq, - right: Box::new(1i64.lit()), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _; - let result = evaluator.evaluate_array(&input_1).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![true, false, false])); - - let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _; - let result = evaluator.evaluate_array(&input_2).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![true, true, true])); - - let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _; - let result = evaluator.evaluate_array(&input_3).unwrap(); - assert_eq!(result, BooleanBuffer::from(vec![])); - } - - #[test] - fn run_on_scalar() { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("foo"))), - op: Operator::Lt, - right: Box::new(1i64.lit()), - }); - let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); - - let input_1 = ScalarValue::Int64(Some(1)); - let result = evaluator.evaluate_scalar(&input_1).unwrap(); - assert!(!result); - - let input_2 = ScalarValue::Int64(Some(0)); - let result = evaluator.evaluate_scalar(&input_2).unwrap(); - assert!(result); - - let input_3 = ScalarValue::Int64(None); - let result = evaluator.evaluate_scalar(&input_3).unwrap(); - assert!(!result); - } - - #[test] - fn batch_filter_test() { - let expr = col("ts").gt(lit(123456u64)); - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("ts", DataType::UInt64, false), - ]); - let df_schema = DFSchema::try_from(schema.clone()).unwrap(); - let props = ExecutionProps::new(); - let physical_expr = create_physical_expr( - &expr, - &df_schema, - &props, - &PhysicalPlanningContext::default(), - ) - .unwrap(); - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])), - Arc::new(datatypes::arrow::array::UInt64Array::from(vec![ - 123456, 123457, 123458, - ])), - ], - ) - .unwrap(); - let new_batch = batch_filter(&batch, &physical_expr).unwrap(); - assert_eq!(new_batch.num_rows(), 2); - let first_column_values = new_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]); - assert_eq!(first_column_values, &expected); - } - - #[test] - fn test_complex_filter_expression() { - // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D' - let col_eq_b = col("col").eq(lit("B")); - let col_eq_c = col("col").eq(lit("C")); - let col_eq_d = col("col").eq(lit("D")); - - // Build the OR chain - let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d); - - // Check that SimpleFilterEvaluator can handle OR chain - let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap(); - assert_eq!(or_evaluator.column_name, "col"); - assert_eq!(or_evaluator.op, Operator::Or); - assert_eq!(or_evaluator.literal_list.len(), 3); - assert_eq!( - format!("{:?}", or_evaluator.literal_list), - "[Scalar(StringArray\n[\n \"B\",\n]), Scalar(StringArray\n[\n \"C\",\n]), Scalar(StringArray\n[\n \"D\",\n])]" - ); - - // Create a schema and batch for testing - let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]); - let df_schema = DFSchema::try_from(schema.clone()).unwrap(); - let props = ExecutionProps::new(); - let physical_expr = create_physical_expr( - &col_or_expr, - &df_schema, - &props, - &PhysicalPlanningContext::default(), - ) - .unwrap(); - - // Create test data - let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![ - "B", "C", "E", "B", "C", "D", "F", - ])); - let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap(); - let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]); - - // Filter the batch - let filtered_batch = batch_filter(&batch, &physical_expr).unwrap(); - - // Expected: rows with col in ("B", "C", "D") - // That would be rows 0, 1, 3, 4, 5 - assert_eq!(filtered_batch.num_rows(), 5); - - let col_filtered = filtered_batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(col_filtered, &expected); - } - - #[test] - fn test_maybe_build_regex() { - // Test case for RegexMatch (case sensitive, non-negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(!negative); - assert!(regex.unwrap().is_match("axxb")); - - // Test case for RegexIMatch (case insensitive, non-negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexIMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(!negative); - assert!(regex.unwrap().is_match("AxxB")); - - // Test case for RegexNotMatch (case sensitive, negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexNotMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(negative); - - // Test case for RegexNotIMatch (case insensitive, negative) - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexNotIMatch, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_some()); - assert!(negative); - - // Test with empty regex pattern - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("".to_string())), - ) - .unwrap(); - assert!(regex.is_none()); - assert!(!negative); - - // Test with non-regex operator - let (regex, negative) = SimpleFilterEvaluator::maybe_build_regex( - Operator::Eq, - &ScalarValue::Utf8(Some("a.*b".to_string())), - ) - .unwrap(); - assert!(regex.is_none()); - assert!(!negative); - - // Test with invalid regex pattern - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(Some("a(b".to_string())), - ); - assert!(result.is_err()); - - // Test with non-string value - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Int64(Some(123)), - ); - assert!(result.is_err()); - - // Test with null value - let result = SimpleFilterEvaluator::maybe_build_regex( - Operator::RegexMatch, - &ScalarValue::Utf8(None), - ); - assert!(result.is_err()); - } - - #[test] - fn test_regex_match_dictionary_array() { - use datatypes::arrow::array::StringDictionaryBuilder; - - // Create a StringDictionaryArray - let mut builder = StringDictionaryBuilder::::new(); - builder.append("apple").unwrap(); - builder.append("banana").unwrap(); - builder.append("apple").unwrap(); - builder.append("cherry").unwrap(); - let dict_array = builder.finish(); - - // Test regex that matches "apple" - let regex = regex::Regex::new(r"app.*").unwrap(); - let result = regexp_is_match_dictionary(&dict_array, Some(®ex)).unwrap(); - - // Should match indices 0 and 2 (both "apple") - assert_eq!(result.len(), 4); - assert!(result.value(0)); // "apple" - assert!(!result.value(1)); // "banana" - assert!(result.value(2)); // "apple" - assert!(!result.value(3)); // "cherry" - - // Test regex that matches "banana" - let regex2 = regex::Regex::new(r"ban.*").unwrap(); - let result2 = regexp_is_match_dictionary(&dict_array, Some(®ex2)).unwrap(); - - assert!(!result2.value(0)); // "apple" - assert!(result2.value(1)); // "banana" - assert!(!result2.value(2)); // "apple" - assert!(!result2.value(3)); // "cherry" - - // Test with no regex (should match all) - let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap(); - assert!(result3.value(0)); - assert!(result3.value(1)); - assert!(result3.value(2)); - assert!(result3.value(3)); - } - - #[test] - fn test_regex_scan_masks_preserve_sql_null_semantics() { - let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![ - Some("api"), - Some("API"), - Some("db"), - None, - ])) as ArrayRef; - assert_regex_scan_masks( - &plain, - [ - vec![true, false, false, false], - vec![true, true, false, false], - vec![false, true, true, false], - vec![false, false, true, false], - ], - ); - - let dictionary = DictionaryArray::new( - datatypes::arrow::array::UInt32Array::from(vec![ - Some(0), - Some(1), - Some(2), - None, - Some(3), - ]), - Arc::new(datatypes::arrow::array::StringArray::from(vec![ - Some("api"), - Some("API"), - Some("db"), - None, - ])), - ); - let raw = - regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap(); - assert!(raw.is_null(3)); // null dictionary key - assert!(raw.is_null(4)); // non-null key referencing a null dictionary value - - let dictionary = Arc::new(dictionary) as ArrayRef; - assert_regex_scan_masks( - &dictionary, - [ - vec![true, false, false, false, false], - vec![true, true, false, false, false], - vec![false, true, true, false, false], - vec![false, false, true, false, false], - ], - ); - - let negative = regex_evaluator(Operator::RegexNotMatch); - assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap()); - } - - #[test] - fn test_nullable_boolean_predicate_becomes_scan_mask() { - let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]); - assert_eq!( - BooleanArray::from(vec![true, false, false]), - boolean_array_to_scan_mask(&predicate) - ); - } - - fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec; 4]) { - for (op, expected) in [ - Operator::RegexMatch, - Operator::RegexIMatch, - Operator::RegexNotMatch, - Operator::RegexNotIMatch, - ] - .into_iter() - .zip(expected) - { - assert_eq!( - BooleanBuffer::from(expected), - regex_evaluator(op).evaluate_array(input).unwrap(), - "{op:?}" - ); - } - } - - fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator { - let expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("host"))), - op, - right: Box::new("^api$".lit()), - }); - SimpleFilterEvaluator::try_new(&expr).unwrap() - } - - fn ts_us(v: i64) -> ScalarValue { - ScalarValue::TimestampMicrosecond(Some(v), None) - } - - fn cast_to_ms(expr: &Expr) -> Option { - SimpleFilterEvaluator::try_new(expr) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype()) - } - - /// Evaluates a cast `Filter` outcome against millisecond values and - /// returns the mask. - fn eval_ms_mask(cast: Option, values: &[i64]) -> Vec { - let filter = match cast.expect("cast must apply") { - TimestampUnitCast::Filter(filter) => filter, - other => panic!("expected Filter outcome, got {other:?}"), - }; - let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef; - filter.evaluate_array(&array).unwrap().iter().collect() - } - - #[test] - fn cast_timestamp_unit_converts_representable_literal() { - // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms. - assert_eq!( - vec![false, true, false], - eval_ms_mask( - cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))), - &[6_999, 7_000, 7_001] - ) - ); - // != complements =. - assert_eq!( - vec![true, false, true], - eval_ms_mask( - cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))), - &[6_999, 7_000, 7_001] - ) - ); - } - - #[test] - fn cast_timestamp_unit_prunes_non_representable_equality() { - // 7_000_500us is not a whole millisecond: no ms row can equal it, - // and every ms row satisfies `!=`. - assert!(matches!( - cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))), - Some(TimestampUnitCast::Pruned) - )); - assert!(matches!( - cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))), - Some(TimestampUnitCast::Matched) - )); - } - - #[test] - fn cast_timestamp_unit_strengthens_inequalities() { - // 2_500_500us is strictly between 2500ms and 2501ms: the cast must - // not round the literal to a boundary the original predicate - // excludes (>= 2500ms would wrongly match 2500ms). - let values = vec![2_499, 2_500, 2_501]; - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values) - ); - - // Representable boundary keeps the original operator. - let values = vec![2_499, 2_500, 2_501]; - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![false, true, true], - eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![true, false, false], - eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values) - ); - } - - #[test] - fn cast_timestamp_unit_handles_negative_instants() { - // -2_500_500us floors to -2501ms with remainder 500us: only rows - // with instant >= -2.5005ms may match >= / >. - let values = vec![-2_502, -2_501, -2_500]; - assert_eq!( - vec![false, false, true], - eval_ms_mask( - cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))), - &values - ) - ); - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values) - ); - assert_eq!( - vec![true, true, false], - eval_ms_mask( - cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))), - &values - ) - ); - // Exactly representable negative literal: -2_500_000us == -2500ms. - assert_eq!( - vec![false, false, true], - eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values) - ); - } - - #[test] - fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() { - // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match. - let expr = col("ts") - .eq(lit(ts_us(7_000_500))) - .or(col("ts").eq(lit(ts_us(7_000_000)))); - assert_eq!( - vec![false, true], - eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000]) - ); - - // A chain where no literal is representable matches nothing. - let expr = col("ts") - .eq(lit(ts_us(7_000_500))) - .or(col("ts").eq(lit(ts_us(6_000_500)))); - assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned))); - } - - #[test] - fn cast_timestamp_unit_null_literal_prunes() { - // A NULL literal never compares true. - assert!(matches!( - cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))), - Some(TimestampUnitCast::Pruned) - )); - } - - #[test] - fn cast_timestamp_unit_same_unit_returns_filter_directly() { - // A literal already in the target unit is returned unchanged. - let expr = col("ts").eq(lit(ts_us(7_000_000))); - let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); - let cast = filter - .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype()) - .unwrap(); - let TimestampUnitCast::Filter(f) = cast else { - panic!("expected Filter, got {cast:?}") - }; - assert_eq!(filter.op, f.op); - assert_eq!(filter.column_name(), f.column_name()); - assert_eq!( - ts_us(7_000_000), - ScalarValue::try_from_array(f.literal.get().0, 0).unwrap() - ); - } - - #[test] - fn cast_timestamp_unit_rejects_non_timestamps() { - // Non-timestamp target type. - assert!( - SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1)))) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::int64_datatype()) - .is_none() - ); - // Non-timestamp literal against a timestamp target. - assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none()); - } - - #[test] - fn cast_timestamp_unit_to_finer_unit() { - // 5 seconds against a nanosecond column becomes 5_000_000_000ns. - let filter = SimpleFilterEvaluator::try_new( - &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))), - ) - .unwrap() - .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype()) - .and_then(|cast| match cast { - TimestampUnitCast::Filter(filter) => Some(filter), - _ => None, - }) - .unwrap(); - let array = Arc::new(TimestampNanosecondArray::from(vec![ - 4_999_999_999, - 5_000_000_000, - ])) as ArrayRef; - assert_eq!( - vec![false, true], - filter - .evaluate_array(&array) - .unwrap() - .iter() - .collect::>() - ); - } -} diff --git a/src/common/recordbatch/src/filter/predicate.rs b/src/common/recordbatch/src/filter/predicate.rs new file mode 100644 index 00000000000..e283633dc02 --- /dev/null +++ b/src/common/recordbatch/src/filter/predicate.rs @@ -0,0 +1,944 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Predicate recognition, execution, and precompiled regular expressions. + +use datafusion::logical_expr::{Expr, Operator}; +use datafusion_common::ScalarValue; +use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datafusion_common::arrow::buffer::BooleanBuffer; +use datafusion_common::arrow::compute::kernels::cmp; +use datafusion_common::cast::as_string_array; +use datatypes::arrow::array::{ + Array, ArrayAccessor, ArrayData, BooleanArray, BooleanBufferBuilder, DictionaryArray, + StringArrayType, +}; +use datatypes::arrow::datatypes::{DataType, UInt32Type}; +use datatypes::arrow::error::ArrowError; +use datatypes::compute::or_kleene; +use regex::Regex; +use snafu::ResultExt; + +use crate::error::{ArrowComputeSnafu, Result, UnsupportedOperationSnafu}; + +#[derive(Debug, Clone)] +pub(crate) enum SimplePredicate { + Comparison { + literal: Scalar, + op: Operator, + regex: Option, + regex_negative: bool, + }, + OrEqChain { + literals: Vec>, + }, + IsNull { + negated: bool, + }, + Constant(Option), + And(Box, Box), + Or(Box, Box), +} + +impl SimplePredicate { + pub(crate) fn try_new(expr: &Expr, column_type: Option<&DataType>) -> Option { + match expr { + Expr::BinaryExpr(binary) if matches!(binary.op, Operator::And | Operator::Or) => { + let left = Self::try_new(&binary.left, column_type)?; + let right = Self::try_new(&binary.right, column_type)?; + if binary.op == Operator::And { + return Some(Self::And(Box::new(left), Box::new(right))); + } + // Preserve the encoded-key IN-list fast path for pure equality disjunctions. + if let (Some(mut lhs), Some(rhs)) = + (left.equality_literals(), right.equality_literals()) + { + lhs.extend(rhs); + return Some(Self::OrEqChain { literals: lhs }); + } + Some(Self::Or(Box::new(left), Box::new(right))) + } + Expr::BinaryExpr(binary) => { + let mut op = binary.op; + if !matches!( + op, + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + | Operator::RegexMatch + | Operator::RegexIMatch + | Operator::RegexNotMatch + | Operator::RegexNotIMatch + ) { + return None; + } + let literal = if let Expr::Literal(literal, _) = &*binary.right { + Self::check_column(&binary.left, column_type)?; + literal + } else if let Expr::Literal(literal, _) = &*binary.left { + Self::check_column(&binary.right, column_type)?; + op = op.swap()?; + literal + } else { + return None; + }; + let (regex, regex_negative) = Self::maybe_build_regex(op, literal).ok()?; + Some(Self::Comparison { + literal: literal.to_scalar().ok()?, + op, + regex, + regex_negative, + }) + } + Expr::IsNull(column) | Expr::IsNotNull(column) => { + Self::check_column(column, column_type)?; + Some(Self::IsNull { + negated: matches!(expr, Expr::IsNotNull(_)), + }) + } + Expr::Literal(ScalarValue::Boolean(value), _) => Some(Self::Constant(*value)), + _ => None, + } + } + + fn check_column(expr: &Expr, column_type: Option<&DataType>) -> Option<()> { + match expr { + Expr::Column(_) => Some(()), + Expr::Cast(cast) if cast.field.data_type() == &DataType::Utf8 => { + let Expr::Column(_) = &*cast.expr else { + return None; + }; + match column_type? { + DataType::Utf8 => Some(()), + DataType::Dictionary(key, value) + if **key == DataType::UInt32 && **value == DataType::Utf8 => + { + Some(()) + } + _ => None, + } + } + _ => None, + } + } + + fn equality_literals(&self) -> Option>> { + match self { + Self::Comparison { + literal, + op: Operator::Eq, + .. + } => Some(vec![literal.clone()]), + Self::OrEqChain { literals } => Some(literals.clone()), + _ => None, + } + } + + pub(crate) fn evaluate(&self, input: &impl Datum, len: usize) -> Result { + match self { + Self::Comparison { + literal, + op, + regex, + regex_negative, + } => match op { + Operator::Eq => cmp::eq(input, literal), + Operator::NotEq => cmp::neq(input, literal), + Operator::Lt => cmp::lt(input, literal), + Operator::LtEq => cmp::lt_eq(input, literal), + Operator::Gt => cmp::gt(input, literal), + Operator::GtEq => cmp::gt_eq(input, literal), + Operator::RegexMatch + | Operator::RegexIMatch + | Operator::RegexNotMatch + | Operator::RegexNotIMatch => { + Self::regex_match(input, regex.as_ref(), *regex_negative) + } + _ => { + return UnsupportedOperationSnafu { + reason: format!("{op:?}"), + } + .fail(); + } + } + .context(ArrowComputeSnafu), + Self::OrEqChain { literals } => { + let mut result = BooleanArray::from(vec![false; len]); + for literal in literals { + let rhs = cmp::eq(input, literal).context(ArrowComputeSnafu)?; + result = or_kleene(&result, &rhs).context(ArrowComputeSnafu)?; + } + Ok(result) + } + Self::IsNull { negated } => { + let values = match input.get().0.logical_nulls() { + Some(nulls) if *negated => nulls.inner().clone(), + Some(nulls) => !nulls.inner(), + None if *negated => BooleanBuffer::new_set(len), + None => BooleanBuffer::new_unset(len), + }; + Ok(BooleanArray::new(values, None)) + } + Self::Constant(value) => Ok(BooleanArray::from(vec![*value; len])), + Self::And(left, right) => datatypes::compute::and_kleene( + &left.evaluate(input, len)?, + &right.evaluate(input, len)?, + ) + .context(ArrowComputeSnafu), + Self::Or(left, right) => { + or_kleene(&left.evaluate(input, len)?, &right.evaluate(input, len)?) + .context(ArrowComputeSnafu) + } + } + } + + /// Builds a regex pattern from a scalar value and operator. + /// Returns the `(regex, negative)` and if successful. + /// + /// Returns `Err` if + /// - the value is not a string + /// - the regex pattern is invalid + /// + /// The regex is `None` if + /// - the operator is not a regex operator + /// - the pattern is empty + fn maybe_build_regex( + operator: Operator, + value: &ScalarValue, + ) -> Result<(Option, bool), ArrowError> { + let (ignore_case, negative) = match operator { + Operator::RegexMatch => (false, false), + Operator::RegexIMatch => (true, false), + Operator::RegexNotMatch => (false, true), + Operator::RegexNotIMatch => (true, true), + _ => return Ok((None, false)), + }; + let flag = if ignore_case { Some("i") } else { None }; + let regex = value + .try_as_str() + .ok_or_else(|| ArrowError::CastError(format!("Cannot cast {:?} to str", value)))? + .ok_or_else(|| ArrowError::CastError("Regex should not be null".to_string()))?; + let pattern = match flag { + Some(flag) => format!("(?{flag}){regex}"), + None => regex.to_string(), + }; + if pattern.is_empty() { + Ok((None, negative)) + } else { + Regex::new(pattern.as_str()) + .map_err(|e| { + ArrowError::ComputeError(format!("Regular expression did not compile: {e:?}")) + }) + .map(|regex| (Some(regex), negative)) + } + } + + fn regex_match( + input: &impl Datum, + regex: Option<&Regex>, + negative: bool, + ) -> std::result::Result { + let array = input.get().0; + + // Try to cast to StringArray first + if let Ok(string_array) = as_string_array(array) { + let mut result = regexp_is_match_scalar(string_array, regex)?; + if negative { + result = datatypes::compute::not(&result)?; + } + return Ok(result); + } + + // Try to cast to StringDictionaryArray + if let Some(dict_array) = array.as_any().downcast_ref::>() { + let mut result = regexp_is_match_dictionary(dict_array, regex)?; + if negative { + result = datatypes::compute::not(&result)?; + } + return Ok(result); + } + + Err(ArrowError::CastError(format!( + "Cannot cast {:?} to StringArray or StringDictionaryArray", + array.data_type() + ))) + } +} + +/// The same as arrow [regexp_is_match_scalar()](datatypes::compute::kernels::regexp::regexp_is_match_scalar()) +/// with pre-compiled regex. +/// See for the implementation details. +pub fn regexp_is_match_scalar<'a, S>( + array: &'a S, + regex: Option<&Regex>, +) -> Result +where + &'a S: StringArrayType<'a>, +{ + let null_bit_buffer = array.nulls().map(|x| x.inner().sliced()); + let mut result = BooleanBufferBuilder::new(array.len()); + + if let Some(re) = regex { + for i in 0..array.len() { + let value = array.value(i); + result.append(re.is_match(value)); + } + } else { + result.append_n(array.len(), true); + } + + let buffer = result.into(); + let data = unsafe { + ArrayData::new_unchecked( + DataType::Boolean, + array.len(), + None, + null_bit_buffer, + 0, + vec![buffer], + vec![], + ) + }; + + Ok(BooleanArray::from(data)) +} + +/// Similar to [regexp_is_match_scalar] but for StringDictionaryArray. +/// Iterates through dictionary keys to get string values and applies regex matching. +pub fn regexp_is_match_dictionary( + dict_array: &DictionaryArray, + regex: Option<&Regex>, +) -> Result { + // Get the string values from the dictionary + let string_values = dict_array + .values() + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::CastError("Dictionary values must be StringArray".to_string()) + })?; + + // Dictionary logical nulls include both null keys and keys whose dictionary value is null. + let logical_nulls = dict_array.logical_nulls(); + let null_bit_buffer = logical_nulls.as_ref().map(|x| x.inner().sliced()); + let mut result = BooleanBufferBuilder::new(dict_array.len()); + + if let Some(re) = regex { + let keys = dict_array.keys().values(); + for i in 0..dict_array.len() { + if logical_nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) { + result.append(false); + } else { + let key = keys[i] as usize; + let string_value = string_values.value(key); + result.append(re.is_match(string_value)); + } + } + } else { + result.append_n(dict_array.len(), true); + } + + let buffer = result.into(); + let data = unsafe { + ArrayData::new_unchecked( + DataType::Boolean, + dict_array.len(), + None, + null_bit_buffer, + 0, + vec![buffer], + vec![], + ) + }; + + Ok(BooleanArray::from(data)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::execution::context::ExecutionProps; + use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; + use datafusion::logical_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; + use datafusion::physical_expr::create_physical_expr; + use datafusion_common::arrow::array::ArrayRef; + use datafusion_common::arrow::buffer::BooleanBuffer; + use datafusion_common::{Column, DFSchema, ScalarValue}; + use datatypes::arrow::array::{Array, DictionaryArray, RecordBatch}; + use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; + use datatypes::value::Value; + use regex::Regex; + + use crate::filter::predicate::{SimplePredicate, regexp_is_match_dictionary}; + use crate::filter::{SimpleFilterEvaluator, batch_filter}; + + #[test] + fn unsupported_filter_op() { + // `+` is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Plus, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // two literal is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(1.lit()), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // two column is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(Expr::Column(Column::from_name("bar"))), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + + // compound expr is not supported + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1.lit()), + })), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + for negated in [false, true] { + let membership = col("label").in_list(vec![lit(""), lit("lo")], negated); + assert!(SimpleFilterEvaluator::try_new(&membership).is_none()); + assert!( + SimpleFilterEvaluator::try_new(&col("label").is_null().or(membership)).is_none() + ); + } + for expr in [ + col("other").is_null().or(col("label").not_eq(lit("lo"))), + col("a.label") + .is_null() + .or(col("b.label").not_eq(lit("lo"))), + col("label").is_null().or(col("label").eq(col("other"))), + col("label") + .is_null() + .or(Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(col("label")), + DataType::Int64, + )) + .eq(lit("lo"))), + ] { + assert!(SimpleFilterEvaluator::try_new(&expr).is_none(), "{expr}"); + } + } + + #[test] + fn string_cast_requires_matching_column_type() { + let cast = Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(col("label")), + DataType::Utf8, + )); + let expr = col("label").is_null().or(cast.not_eq(lit("lo"))); + for data_type in [ + DataType::Int64, + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Binary, + ] { + assert!( + SimpleFilterEvaluator::try_new_with_column_type(&expr, &|_| Some( + data_type.clone() + )) + .is_none() + ); + } + assert!(SimpleFilterEvaluator::try_new(&expr).is_none()); + } + + #[tokio::test] + async fn single_column_predicates_match_datafusion() { + use datafusion::prelude::SessionContext; + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + use datatypes::arrow::array::{Int32Array, StringArray, UInt32Array}; + + let values = vec![ + None, + Some(""), + Some("lo"), + Some("eth0"), + Some("a\nb"), + Some("æœåŠ¡"), + None, + ]; + let strings: ArrayRef = Arc::new(StringArray::from(values)); + let dictionary: ArrayRef = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![ + None, + Some(0), + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ]), + Arc::new(StringArray::from(vec![ + Some(""), + Some("lo"), + Some("eth0"), + Some("a\nb"), + Some("æœåŠ¡"), + None, + ])), + )); + let null = || lit(ScalarValue::Utf8(None)); + let mut predicates = vec![ + col("label").is_null(), + col("label").is_not_null(), + col("label") + .is_null() + .or(col("label").eq(lit(""))) + .and(col("label").is_not_null()), + col("label") + .not_eq(lit("lo")) + .and(col("label").not_eq(lit("eth0"))), + col("label") + .eq(lit("lo")) + .or(col("label").eq(lit("eth0"))) + .or(col("label").eq(lit(""))), + col("label").eq(null()).or(col("label").eq(lit("eth0"))), + col("label").gt(lit("a")).and(col("label").lt(lit("z"))), + lit("lo").not_eq(col("label")), + ]; + for (op, literal) in [ + (Operator::Eq, ""), + (Operator::NotEq, "lo"), + (Operator::RegexNotMatch, "^(?:lo|vfat)$"), + (Operator::RegexMatch, "^(?:|eth0)$"), + ] { + let comparison = Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("label")), + op, + Box::new(lit(literal)), + )); + predicates.extend([ + col("label").is_null().or(comparison.clone()), + comparison.or(col("label").is_null()), + ]); + } + for input in [strings, dictionary] { + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![ + Field::new("label", input.data_type().clone(), true), + Field::new("id", DataType::Int32, false), + ])); + ctx.register_batch( + "labels", + RecordBatch::try_new( + schema, + vec![ + input.clone(), + Arc::new(Int32Array::from_iter_values(0..input.len() as i32)), + ], + ) + .unwrap(), + ) + .unwrap(); + for expr in &predicates { + let df = ctx + .table("labels") + .await + .unwrap() + .filter(expr.clone()) + .unwrap(); + let plan = df.clone().into_optimized_plan().unwrap(); + let expected = df + .collect() + .await + .unwrap() + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + let evaluate = |expr: &Expr| { + let evaluator = + SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + (name == "label").then(|| input.data_type().clone()) + }) + .unwrap_or_else(|| panic!("unsupported: {expr}")); + let mask = evaluator.evaluate_array(&input).unwrap(); + for (i, matched) in mask.iter().enumerate() { + let scalar = ScalarValue::try_from_array(&input, i).unwrap(); + assert_eq!( + evaluator.evaluate_scalar(&scalar).unwrap(), + matched, + "{expr}, row={i}" + ); + } + mask + }; + let selected = |mask: BooleanBuffer| { + mask.iter() + .enumerate() + .filter_map(|(i, matched)| matched.then_some(i as i32)) + .collect::>() + }; + assert_eq!(selected(evaluate(expr)), expected, "{expr}"); + let mut filters = Vec::new(); + plan.apply(|plan| { + if let datafusion::logical_expr::LogicalPlan::Filter(filter) = plan { + filters.push(filter.predicate.clone()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + if !filters.is_empty() { + let mut mask = BooleanBuffer::new_set(input.len()); + for filter in filters { + mask = &mask & &evaluate(&filter); + } + assert_eq!(selected(mask), expected, "optimized: {plan}"); + } + } + } + } + + #[test] + fn supported_filter_op() { + // equal + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1.lit()), + }); + let _ = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + // swap operands + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(1.lit()), + op: Operator::Lt, + right: Box::new(Expr::Column(Column::from_name("foo"))), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!(evaluator.is_gt()); + assert_eq!(evaluator.column_name, "foo".to_string()); + } + + #[test] + fn run_on_array() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Eq, + right: Box::new(1i64.lit()), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _; + let result = evaluator.evaluate_array(&input_1).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![true, false, false])); + + let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _; + let result = evaluator.evaluate_array(&input_2).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![true, true, true])); + + let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _; + let result = evaluator.evaluate_array(&input_3).unwrap(); + assert_eq!(result, BooleanBuffer::from(vec![])); + } + + #[test] + fn run_on_scalar() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("foo"))), + op: Operator::Lt, + right: Box::new(1i64.lit()), + }); + let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap(); + + let input_1 = ScalarValue::Int64(Some(1)); + let result = evaluator.evaluate_scalar(&input_1).unwrap(); + assert!(!result); + + let input_2 = ScalarValue::Int64(Some(0)); + let result = evaluator.evaluate_scalar(&input_2).unwrap(); + assert!(result); + + let input_3 = ScalarValue::Int64(None); + let result = evaluator.evaluate_scalar(&input_3).unwrap(); + assert!(!result); + } + + #[test] + fn test_complex_filter_expression() { + // Create an expression tree for: col = 'B' OR col = 'C' OR col = 'D' + let col_eq_b = col("col").eq(lit("B")); + let col_eq_c = col("col").eq(lit("C")); + let col_eq_d = col("col").eq(lit("D")); + + // Build the OR chain + let col_or_expr = col_eq_b.or(col_eq_c).or(col_eq_d); + + // Check that SimpleFilterEvaluator can handle OR chain + let or_evaluator = SimpleFilterEvaluator::try_new(&col_or_expr).unwrap(); + assert_eq!(or_evaluator.column_name, "col"); + assert!(or_evaluator.is_or_eq_chain()); + assert_eq!( + or_evaluator.literal_list_values().unwrap(), + vec![Value::from("B"), Value::from("C"), Value::from("D")] + ); + + // Create a schema and batch for testing + let schema = Schema::new(vec![Field::new("col", DataType::Utf8, false)]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + let props = ExecutionProps::new(); + let physical_expr = create_physical_expr( + &col_or_expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); + + // Create test data + let col_data = Arc::new(datatypes::arrow::array::StringArray::from(vec![ + "B", "C", "E", "B", "C", "D", "F", + ])); + let batch = RecordBatch::try_new(Arc::new(schema), vec![col_data]).unwrap(); + let expected = datatypes::arrow::array::StringArray::from(vec!["B", "C", "B", "C", "D"]); + + // Filter the batch + let filtered_batch = batch_filter(&batch, &physical_expr).unwrap(); + + // Expected: rows with col in ("B", "C", "D") + // That would be rows 0, 1, 3, 4, 5 + assert_eq!(filtered_batch.num_rows(), 5); + + let col_filtered = filtered_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col_filtered, &expected); + } + + #[test] + fn test_maybe_build_regex() { + // Test case for RegexMatch (case sensitive, non-negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(!negative); + assert!(regex.unwrap().is_match("axxb")); + + // Test case for RegexIMatch (case insensitive, non-negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexIMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(!negative); + assert!(regex.unwrap().is_match("AxxB")); + + // Test case for RegexNotMatch (case sensitive, negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexNotMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(negative); + + // Test case for RegexNotIMatch (case insensitive, negative) + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexNotIMatch, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_some()); + assert!(negative); + + // Test with empty regex pattern + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("".to_string())), + ) + .unwrap(); + assert!(regex.is_none()); + assert!(!negative); + + // Test with non-regex operator + let (regex, negative) = SimplePredicate::maybe_build_regex( + Operator::Eq, + &ScalarValue::Utf8(Some("a.*b".to_string())), + ) + .unwrap(); + assert!(regex.is_none()); + assert!(!negative); + + // Test with invalid regex pattern + let result = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Utf8(Some("a(b".to_string())), + ); + assert!(result.is_err()); + + // Test with non-string value + let result = SimplePredicate::maybe_build_regex( + Operator::RegexMatch, + &ScalarValue::Int64(Some(123)), + ); + assert!(result.is_err()); + + // Test with null value + let result = + SimplePredicate::maybe_build_regex(Operator::RegexMatch, &ScalarValue::Utf8(None)); + assert!(result.is_err()); + } + + #[test] + fn test_regex_match_dictionary_array() { + use datatypes::arrow::array::StringDictionaryBuilder; + + // Create a StringDictionaryArray + let mut builder = StringDictionaryBuilder::::new(); + builder.append("apple").unwrap(); + builder.append("banana").unwrap(); + builder.append("apple").unwrap(); + builder.append("cherry").unwrap(); + let dict_array = builder.finish(); + + // Test regex that matches "apple" + let regex = regex::Regex::new(r"app.*").unwrap(); + let result = regexp_is_match_dictionary(&dict_array, Some(®ex)).unwrap(); + + // Should match indices 0 and 2 (both "apple") + assert_eq!(result.len(), 4); + assert!(result.value(0)); // "apple" + assert!(!result.value(1)); // "banana" + assert!(result.value(2)); // "apple" + assert!(!result.value(3)); // "cherry" + + // Test regex that matches "banana" + let regex2 = regex::Regex::new(r"ban.*").unwrap(); + let result2 = regexp_is_match_dictionary(&dict_array, Some(®ex2)).unwrap(); + + assert!(!result2.value(0)); // "apple" + assert!(result2.value(1)); // "banana" + assert!(!result2.value(2)); // "apple" + assert!(!result2.value(3)); // "cherry" + + // Test with no regex (should match all) + let result3 = regexp_is_match_dictionary(&dict_array, None).unwrap(); + assert!(result3.value(0)); + assert!(result3.value(1)); + assert!(result3.value(2)); + assert!(result3.value(3)); + } + + #[test] + fn test_regex_scan_masks_preserve_sql_null_semantics() { + let plain = Arc::new(datatypes::arrow::array::StringArray::from(vec![ + Some("api"), + Some("API"), + Some("db"), + None, + ])) as ArrayRef; + assert_regex_scan_masks( + &plain, + [ + vec![true, false, false, false], + vec![true, true, false, false], + vec![false, true, true, false], + vec![false, false, true, false], + ], + ); + + let dictionary = DictionaryArray::new( + datatypes::arrow::array::UInt32Array::from(vec![ + Some(0), + Some(1), + Some(2), + None, + Some(3), + ]), + Arc::new(datatypes::arrow::array::StringArray::from(vec![ + Some("api"), + Some("API"), + Some("db"), + None, + ])), + ); + let raw = + regexp_is_match_dictionary(&dictionary, Some(&Regex::new("^api$").unwrap())).unwrap(); + assert!(raw.is_null(3)); // null dictionary key + assert!(raw.is_null(4)); // non-null key referencing a null dictionary value + + let dictionary = Arc::new(dictionary) as ArrayRef; + assert_regex_scan_masks( + &dictionary, + [ + vec![true, false, false, false, false], + vec![true, true, false, false, false], + vec![false, true, true, false, false], + vec![false, false, true, false, false], + ], + ); + + let negative = regex_evaluator(Operator::RegexNotMatch); + assert!(!negative.evaluate_scalar(&ScalarValue::Utf8(None)).unwrap()); + } + + fn assert_regex_scan_masks(input: &ArrayRef, expected: [Vec; 4]) { + for (op, expected) in [ + Operator::RegexMatch, + Operator::RegexIMatch, + Operator::RegexNotMatch, + Operator::RegexNotIMatch, + ] + .into_iter() + .zip(expected) + { + assert_eq!( + BooleanBuffer::from(expected), + regex_evaluator(op).evaluate_array(input).unwrap(), + "{op:?}" + ); + } + } + + fn regex_evaluator(op: Operator) -> SimpleFilterEvaluator { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name("host"))), + op, + right: Box::new("^api$".lit()), + }); + SimpleFilterEvaluator::try_new(&expr).unwrap() + } +} diff --git a/src/common/recordbatch/src/filter/tests.rs b/src/common/recordbatch/src/filter/tests.rs new file mode 100644 index 00000000000..77abf59b404 --- /dev/null +++ b/src/common/recordbatch/src/filter/tests.rs @@ -0,0 +1,72 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::execution::context::ExecutionProps; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion::logical_expr::{col, lit}; +use datafusion::physical_expr::create_physical_expr; +use datafusion_common::DFSchema; +use datatypes::arrow::datatypes::{DataType, Field, Schema}; + +use datatypes::arrow::array::{BooleanArray, RecordBatch}; + +use crate::filter::{batch_filter, boolean_array_to_scan_mask}; + +#[test] +fn batch_filter_test() { + let expr = col("ts").gt(lit(123456u64)); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("ts", DataType::UInt64, false), + ]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + let props = ExecutionProps::new(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(datatypes::arrow::array::Int32Array::from(vec![4, 5, 6])), + Arc::new(datatypes::arrow::array::UInt64Array::from(vec![ + 123456, 123457, 123458, + ])), + ], + ) + .unwrap(); + let new_batch = batch_filter(&batch, &physical_expr).unwrap(); + assert_eq!(new_batch.num_rows(), 2); + let first_column_values = new_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let expected = datatypes::arrow::array::Int32Array::from(vec![5, 6]); + assert_eq!(first_column_values, &expected); +} + +#[test] +fn test_nullable_boolean_predicate_becomes_scan_mask() { + let predicate = BooleanArray::from(vec![Some(true), None, Some(false)]); + assert_eq!( + BooleanArray::from(vec![true, false, false]), + boolean_array_to_scan_mask(&predicate) + ); +} diff --git a/src/common/recordbatch/src/filter/timestamp.rs b/src/common/recordbatch/src/filter/timestamp.rs new file mode 100644 index 00000000000..4c19aad546b --- /dev/null +++ b/src/common/recordbatch/src/filter/timestamp.rs @@ -0,0 +1,439 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Timestamp-unit adaptation for predicates evaluated against historical SSTs. + +use common_time::timestamp::div_mod_units; +use datafusion::logical_expr::Operator; +use datafusion_common::ScalarValue; +use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar}; +use datatypes::arrow::datatypes::TimeUnit; + +use crate::filter::predicate::SimplePredicate; + +/// Builds a tz-naive timestamp literal for `value` in `unit`. +pub fn timestamp_scalar_value(value: i64, unit: TimeUnit) -> ScalarValue { + match unit { + TimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None), + TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None), + TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None), + TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None), + } +} + +impl SimplePredicate { + pub(crate) fn cast_timestamp_unit(&self, target_unit: TimeUnit) -> Option { + // Constants represent scan outcomes here. As in the public Matched/Pruned + // result, UNKNOWN can be collapsed to false under these AND/OR-only trees. + match self { + Self::And(left, right) | Self::Or(left, right) => { + let left = left.cast_timestamp_unit(target_unit)?; + let right = right.cast_timestamp_unit(target_unit)?; + let and = matches!(self, Self::And(..)); + Some(match (left, right) { + (Self::Constant(Some(false)), _) | (_, Self::Constant(Some(false))) if and => { + Self::Constant(Some(false)) + } + (Self::Constant(Some(true)), _) | (_, Self::Constant(Some(true))) if !and => { + Self::Constant(Some(true)) + } + (Self::Constant(Some(true)), other) | (other, Self::Constant(Some(true))) + if and => + { + other + } + (Self::Constant(Some(false)), other) | (other, Self::Constant(Some(false))) + if !and => + { + other + } + (left, right) => { + if and { + Self::And(Box::new(left), Box::new(right)) + } else { + Self::Or(Box::new(left), Box::new(right)) + } + } + }) + } + Self::IsNull { .. } => Some(self.clone()), + Self::Constant(value) => Some(Self::Constant(Some(value.unwrap_or(false)))), + Self::OrEqChain { literals } => { + let mut converted = Vec::with_capacity(literals.len()); + for literal in literals { + let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; + let (value, unit) = timestamp_scalar_parts(&scalar)?; + let Some(value) = value else { + continue; + }; + let cast = div_mod_units(value, unit.into(), target_unit.into())?; + if cast.remainder == 0 { + converted.push(timestamp_scalar(cast.quotient, target_unit)?); + } + } + if converted.is_empty() { + return Some(Self::Constant(Some(false))); + } + Some(Self::OrEqChain { + literals: converted, + }) + } + Self::Comparison { literal, op, .. } => { + let scalar = ScalarValue::try_from_array(literal.get().0, 0).ok()?; + let (value, unit) = timestamp_scalar_parts(&scalar)?; + let Some(value) = value else { + return Some(Self::Constant(Some(false))); + }; + if unit == target_unit { + return Some(self.clone()); + } + let cast = div_mod_units(value, unit.into(), target_unit.into())?; + let divisible = cast.remainder == 0; + let literal = timestamp_scalar(cast.quotient, target_unit)?; + let filter = |op| Self::Comparison { + literal: literal.clone(), + op, + regex: None, + regex_negative: false, + }; + Some(match op { + Operator::Eq if divisible => filter(Operator::Eq), + Operator::Eq => Self::Constant(Some(false)), + Operator::NotEq if divisible => filter(Operator::NotEq), + Operator::NotEq => Self::Constant(Some(true)), + // The quotient is floor(L), including negative timestamps. + // v > L is v > quotient whether or not L is representable. + Operator::Gt => filter(Operator::Gt), + Operator::GtEq if divisible => filter(Operator::GtEq), + // Between target values, v >= L is v > quotient. + Operator::GtEq => filter(Operator::Gt), + Operator::Lt if divisible => filter(Operator::Lt), + // Between target values, v < L is v <= quotient. + Operator::Lt => filter(Operator::LtEq), + Operator::LtEq => filter(Operator::LtEq), + _ => return None, + }) + } + } + } +} + +/// Extracts the value and unit from a tz-naive timestamp scalar. +fn timestamp_scalar_parts(scalar: &ScalarValue) -> Option<(Option, TimeUnit)> { + let (value, unit, timezone) = match scalar { + ScalarValue::TimestampSecond(v, tz) => (*v, TimeUnit::Second, tz), + ScalarValue::TimestampMillisecond(v, tz) => (*v, TimeUnit::Millisecond, tz), + ScalarValue::TimestampMicrosecond(v, tz) => (*v, TimeUnit::Microsecond, tz), + ScalarValue::TimestampNanosecond(v, tz) => (*v, TimeUnit::Nanosecond, tz), + _ => return None, + }; + // A timezone-aware literal doesn't compare against a tz-naive column; + // leave it to the caller instead of guessing the intended semantics. + (timezone.is_none()).then_some((value, unit)) +} + +/// Builds a one-element scalar array holding `value` in `unit`. +fn timestamp_scalar(value: i64, unit: TimeUnit) -> Option> { + timestamp_scalar_value(value, unit).to_scalar().ok() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::execution::context::ExecutionProps; + use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; + use datafusion::logical_expr::{Expr, col, lit}; + use datafusion::physical_expr::create_physical_expr; + use datafusion_common::arrow::array::ArrayRef; + use datafusion_common::arrow::buffer::BooleanBuffer; + use datafusion_common::cast::as_boolean_array; + use datafusion_common::{DFSchema, ScalarValue}; + use datatypes::arrow::array::{ + RecordBatch, TimestampMillisecondArray, TimestampNanosecondArray, + }; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::data_type::ConcreteDataType; + + use crate::filter::{SimpleFilterEvaluator, TimestampUnitCast, boolean_array_to_scan_mask}; + + #[test] + fn compound_timestamp_filters_preserve_unit_conversion() { + use datatypes::arrow::array::TimestampMicrosecondArray; + + let null = || lit(ScalarValue::TimestampMicrosecond(None, None)); + let predicates = [ + col("ts") + .gt(lit(ts_us(-1500))) + .and(col("ts").lt(lit(ts_us(500)))), + col("ts") + .eq(lit(ts_us(500))) + .or(col("ts").gt_eq(lit(ts_us(1500)))), + col("ts") + .is_not_null() + .and(col("ts").not_eq(lit(ts_us(500)))), + col("ts").eq(null()).or(col("ts").eq(lit(ts_us(-1000)))), + col("ts") + .eq(lit(ts_us(500))) + .and(col("ts").gt(lit(ts_us(0)))), + col("ts") + .not_eq(lit(ts_us(500))) + .or(col("ts").lt(lit(ts_us(0)))), + ]; + let native: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![ + -2000, -1000, 0, 1000, 2000, + ])); + let old: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![-2, -1, 0, 1, 2])); + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + native.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![native]).unwrap(); + for expr in predicates { + let physical = create_physical_expr( + &expr, + &DFSchema::try_from(schema.as_ref().clone()).unwrap(), + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let result = physical + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let expected = boolean_array_to_scan_mask(as_boolean_array(&result).unwrap()) + .values() + .clone(); + let actual = match cast_to_ms(&expr).unwrap() { + TimestampUnitCast::Matched => BooleanBuffer::new_set(old.len()), + TimestampUnitCast::Pruned => BooleanBuffer::new_unset(old.len()), + TimestampUnitCast::Filter(filter) => filter.evaluate_array(&old).unwrap(), + }; + assert_eq!(actual, expected, "{expr}"); + } + } + + fn ts_us(v: i64) -> ScalarValue { + ScalarValue::TimestampMicrosecond(Some(v), None) + } + + fn cast_to_ms(expr: &Expr) -> Option { + SimpleFilterEvaluator::try_new(expr) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::timestamp_millisecond_datatype()) + } + + /// Evaluates a cast `Filter` outcome against millisecond values and + /// returns the mask. + fn eval_ms_mask(cast: Option, values: &[i64]) -> Vec { + let filter = match cast.expect("cast must apply") { + TimestampUnitCast::Filter(filter) => filter, + other => panic!("expected Filter outcome, got {other:?}"), + }; + let array = Arc::new(TimestampMillisecondArray::from(values.to_vec())) as ArrayRef; + filter.evaluate_array(&array).unwrap().iter().collect() + } + + #[test] + fn cast_timestamp_unit_converts_representable_literal() { + // ts = 7_000_000us evaluated against a ms column becomes ts = 7000ms. + assert_eq!( + vec![false, true, false], + eval_ms_mask( + cast_to_ms(&col("ts").eq(lit(ts_us(7_000_000)))), + &[6_999, 7_000, 7_001] + ) + ); + // != complements =. + assert_eq!( + vec![true, false, true], + eval_ms_mask( + cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_000)))), + &[6_999, 7_000, 7_001] + ) + ); + } + + #[test] + fn cast_timestamp_unit_prunes_non_representable_equality() { + // 7_000_500us is not a whole millisecond: no ms row can equal it, + // and every ms row satisfies `!=`. + assert!(matches!( + cast_to_ms(&col("ts").eq(lit(ts_us(7_000_500)))), + Some(TimestampUnitCast::Pruned) + )); + assert!(matches!( + cast_to_ms(&col("ts").not_eq(lit(ts_us(7_000_500)))), + Some(TimestampUnitCast::Matched) + )); + } + + #[test] + fn cast_timestamp_unit_strengthens_inequalities() { + // 2_500_500us is strictly between 2500ms and 2501ms: the cast must + // not round the literal to a boundary the original predicate + // excludes (>= 2500ms would wrongly match 2500ms). + let values = vec![2_499, 2_500, 2_501]; + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_500)))), &values) + ); + + // Representable boundary keeps the original operator. + let values = vec![2_499, 2_500, 2_501]; + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![false, true, true], + eval_ms_mask(cast_to_ms(&col("ts").gt_eq(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![true, false, false], + eval_ms_mask(cast_to_ms(&col("ts").lt(lit(ts_us(2_500_000)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask(cast_to_ms(&col("ts").lt_eq(lit(ts_us(2_500_000)))), &values) + ); + } + + #[test] + fn cast_timestamp_unit_handles_negative_instants() { + // -2_500_500us floors to -2501ms with remainder 500us: only rows + // with instant >= -2.5005ms may match >= / >. + let values = vec![-2_502, -2_501, -2_500]; + assert_eq!( + vec![false, false, true], + eval_ms_mask( + cast_to_ms(&col("ts").gt_eq(lit(ts_us(-2_500_500)))), + &values + ) + ); + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").gt(lit(ts_us(-2_500_500)))), &values) + ); + assert_eq!( + vec![true, true, false], + eval_ms_mask( + cast_to_ms(&col("ts").lt_eq(lit(ts_us(-2_500_500)))), + &values + ) + ); + // Exactly representable negative literal: -2_500_000us == -2500ms. + assert_eq!( + vec![false, false, true], + eval_ms_mask(cast_to_ms(&col("ts").eq(lit(ts_us(-2_500_000)))), &values) + ); + } + + #[test] + fn cast_timestamp_unit_or_chain_drops_unrepresentable_literals() { + // Only 7_000_000us is a whole millisecond; 7_000_500us cannot match. + let expr = col("ts") + .eq(lit(ts_us(7_000_500))) + .or(col("ts").eq(lit(ts_us(7_000_000)))); + assert_eq!( + vec![false, true], + eval_ms_mask(cast_to_ms(&expr), &[6_999, 7_000]) + ); + + // A chain where no literal is representable matches nothing. + let expr = col("ts") + .eq(lit(ts_us(7_000_500))) + .or(col("ts").eq(lit(ts_us(6_000_500)))); + assert!(matches!(cast_to_ms(&expr), Some(TimestampUnitCast::Pruned))); + } + + #[test] + fn cast_timestamp_unit_null_literal_prunes() { + // A NULL literal never compares true. + assert!(matches!( + cast_to_ms(&col("ts").eq(lit(ScalarValue::TimestampMicrosecond(None, None)))), + Some(TimestampUnitCast::Pruned) + )); + } + + #[test] + fn cast_timestamp_unit_same_unit_returns_filter_directly() { + // A literal already in the target unit is returned unchanged. + let expr = col("ts").eq(lit(ts_us(7_000_000))); + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + let cast = filter + .cast_timestamp_unit(&ConcreteDataType::timestamp_microsecond_datatype()) + .unwrap(); + let TimestampUnitCast::Filter(f) = cast else { + panic!("expected Filter, got {cast:?}") + }; + assert!(filter.is_eq() && f.is_eq()); + assert_eq!(filter.column_name(), f.column_name()); + assert_eq!(filter.literal_value(), f.literal_value()); + } + + #[test] + fn cast_timestamp_unit_rejects_non_timestamps() { + // Non-timestamp target type. + assert!( + SimpleFilterEvaluator::try_new(&col("ts").gt(lit(ts_us(1)))) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::int64_datatype()) + .is_none() + ); + // Non-timestamp literal against a timestamp target. + assert!(cast_to_ms(&col("ts").gt(lit(42_i64))).is_none()); + } + + #[test] + fn cast_timestamp_unit_to_finer_unit() { + // 5 seconds against a nanosecond column becomes 5_000_000_000ns. + let filter = SimpleFilterEvaluator::try_new( + &col("ts").eq(lit(ScalarValue::TimestampSecond(Some(5), None))), + ) + .unwrap() + .cast_timestamp_unit(&ConcreteDataType::timestamp_nanosecond_datatype()) + .and_then(|cast| match cast { + TimestampUnitCast::Filter(filter) => Some(filter), + _ => None, + }) + .unwrap(); + let array = Arc::new(TimestampNanosecondArray::from(vec![ + 4_999_999_999, + 5_000_000_000, + ])) as ArrayRef; + assert_eq!( + vec![false, true], + filter + .evaluate_array(&array) + .unwrap() + .iter() + .collect::>() + ); + } +} diff --git a/src/mito2/src/engine/prune_test.rs b/src/mito2/src/engine/prune_test.rs index f2171510345..71e6845e96f 100644 --- a/src/mito2/src/engine/prune_test.rs +++ b/src/mito2/src/engine/prune_test.rs @@ -288,7 +288,7 @@ async fn test_prune_memtable_complex_expr_with_format(flat_format: bool) { ) .await; - // ts filter will be ignored when pruning time series in memtable. + // Both the compound time range and the tag predicate are evaluated during the scan. let filters = vec![time_range_expr(4, 7), col("tag_0").lt(lit("6"))]; let stream = engine @@ -306,10 +306,6 @@ async fn test_prune_memtable_complex_expr_with_format(flat_format: bool) { +-------+---------+---------------------+ | tag_0 | field_0 | ts | +-------+---------+---------------------+ -| 0 | 0.0 | 1970-01-01T00:00:00 | -| 1 | 1.0 | 1970-01-01T00:00:01 | -| 2 | 2.0 | 1970-01-01T00:00:02 | -| 3 | 3.0 | 1970-01-01T00:00:03 | | 4 | 4.0 | 1970-01-01T00:00:04 | | 5 | 5.0 | 1970-01-01T00:00:05 | +-------+---------+---------------------+"; diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index e10968d3a32..39a1e66f6d8 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -73,6 +73,215 @@ use crate::test_util; use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{CreateRequestBuilder, TestEnv, reopen_region}; +#[tokio::test] +async fn test_nullable_label_filters_in_memtables_and_ssts() { + use datafusion_expr::{BinaryExpr as LogicalBinaryExpr, Cast, Expr}; + + for (flat, memtable) in [(false, "time_series"), (true, "bulk")] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", memtable) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = [None, Some(""), Some("tmpfs"), Some("ext4")] + .into_iter() + .enumerate() + .map(|(i, label)| { + let mut row = row(vec![ + ValueData::StringValue(String::new()), + ValueData::F64Value(i as f64), + ValueData::TimestampMillisecondValue(i as i64), + ]); + row.values[0].value_data = label.map(|v| ValueData::StringValue(v.into())); + row + }) + .collect(); + test_util::put_rows(&engine, region_id, Rows { schema, rows }).await; + for flushed in [false, true] { + if flushed { + test_util::flush_region(&engine, region_id, None).await; + } + let mut cases = Vec::new(); + for (op, literal, include_null, expected) in [ + (Operator::NotEq, "tmpfs", true, vec![0, 1, 3]), + ( + Operator::RegexNotMatch, + "^(?:tmpfs|vfat)$", + true, + vec![0, 1, 3], + ), + (Operator::Eq, "", true, vec![0, 1]), + (Operator::NotEq, "", false, vec![2, 3]), + (Operator::RegexMatch, "^(?:ext4)$", false, vec![3]), + ] { + for cast in [false, true] { + let column = if cast { + Expr::Cast(Cast::new(Box::new(col("tag_0")), DataType::Utf8)) + } else { + col("tag_0") + }; + let comparison = Expr::BinaryExpr(LogicalBinaryExpr::new( + Box::new(column), + op, + Box::new(lit(literal)), + )); + let filter = if include_null { + col("tag_0").is_null().or(comparison) + } else { + comparison + }; + cases.push((filter, expected.clone())); + } + } + let column = col("tag_0"); + cases.extend([ + (column.clone().is_null(), vec![0]), + (column.clone().is_not_null(), vec![1, 2, 3]), + ( + column + .clone() + .is_null() + .or(column.clone().not_eq(lit("tmpfs"))) + .and(column.clone().is_null().or(column.clone().not_eq(lit("")))), + vec![0, 3], + ), + ( + column.clone().eq(lit("ext4")).or(column.eq(lit("tmpfs"))), + vec![2, 3], + ), + ]); + for (filter, expected) in cases { + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + filters: vec![filter.clone()], + distribution: flat.then_some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut actual = Vec::new(); + for batch in batches.iter() { + let ts = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + actual.extend(ts.values().iter().copied()); + } + actual.sort_unstable(); + assert_eq!(actual, expected, "flat={flat}, flushed={flushed}, {filter}"); + } + } + } +} + +#[tokio::test] +async fn test_compound_time_filters_before_last_row() { + for (flat, memtable) in [ + (false, "time_series"), + (true, "time_series"), + (true, "bulk"), + ] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", memtable) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = (0..5) + .map(|i| { + row(vec![ + ValueData::StringValue("series".into()), + ValueData::F64Value(i as f64), + ValueData::TimestampMillisecondValue(i), + ]) + }) + .collect(); + test_util::put_rows(&engine, region_id, Rows { schema, rows }).await; + let ts = |value| lit(ScalarValue::TimestampMillisecond(Some(value), None)); + let cases = [ + (col("ts").eq(ts(1)).or(col("ts").eq(ts(3))), vec![1, 3]), + (col("ts").gt_eq(ts(1)).and(col("ts").lt(ts(3))), vec![1, 2]), + (col("ts").is_null().or(col("ts").lt(ts(2))), vec![0, 1]), + ]; + for flushed in [false, true] { + if flushed { + test_util::flush_region(&engine, region_id, None).await; + } + for (filter, expected) in &cases { + for selector in [ + None, + Some(TimeSeriesRowSelector::LastRow { after_merge: false }), + Some(TimeSeriesRowSelector::LastRow { after_merge: true }), + ] { + let stream = engine + .scan_to_stream( + region_id, + ScanRequest { + filters: vec![ + col("tag_0").is_null().or(col("tag_0").eq(lit("series"))), + filter.clone(), + ], + distribution: flat.then_some(TimeSeriesDistribution::PerSeries), + series_row_selector: selector, + ..Default::default() + }, + ) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut actual = Vec::new(); + for batch in batches.iter() { + actual.extend( + batch + .column_by_name("ts") + .unwrap() + .as_primitive::() + .values() + .iter() + .copied(), + ); + } + actual.sort_unstable(); + let expected = if selector.is_some() { + expected.last().copied().into_iter().collect::>() + } else { + expected.clone() + }; + assert_eq!( + actual, expected, + "flat={flat}, memtable={memtable}, flushed={flushed}, selector={selector:?}, {filter}" + ); + } + } + } + } +} + #[tokio::test] async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResult<()> { // Create a region with a JSON2 field whose physical Parquet representation is a nested struct. diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index 4e720df88df..fad746927d4 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -641,7 +641,13 @@ impl Iter { predicate .exprs() .iter() - .filter_map(SimpleFilterEvaluator::try_new) + .filter_map(|expr| { + SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + metadata + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + }) + }) .collect::>() }) .unwrap_or_default(); diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index dc68835a11f..5cfd6a29844 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -29,10 +29,12 @@ use common_telemetry::tracing::Instrument; use common_telemetry::{debug, error, tracing, warn}; use common_time::range::TimestampRange; use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_common::pruning::PruningStatistics; -use datafusion_common::{Column, ScalarValue}; +use datafusion_common::{Column, ScalarValue, ToDFSchema}; use datafusion_expr::Expr; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::extension::json::is_json2_extension_type; @@ -57,8 +59,8 @@ use crate::access_layer::AccessLayerRef; use crate::cache::CacheStrategy; use crate::config::DEFAULT_MAX_CONCURRENT_SCAN_FILES; use crate::error::{ - InvalidPartitionExprSnafu, InvalidRequestSnafu, RegionSequenceDomainBrokenSnafu, Result, - SequenceRangeUnsupportedSnafu, + EvalPartitionFilterSnafu, InvalidPartitionExprSnafu, InvalidRequestSnafu, + RegionSequenceDomainBrokenSnafu, Result, SequenceRangeUnsupportedSnafu, }; #[cfg(feature = "enterprise")] use crate::extension::{BoxedExtensionRange, BoxedExtensionRangeProvider}; @@ -2197,6 +2199,22 @@ impl PredicateGroup { expr: expr_json.clone(), })?; + // Persisted partition bounds retain their original types after schema changes. + let schema = metadata + .schema + .arrow_schema() + .clone() + .to_dfschema_ref() + .context(EvalPartitionFilterSnafu)?; + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_schema(schema.clone()) + .build(), + ); + let logical_expr = simplifier + .coerce(logical_expr, &schema) + .and_then(|expr| simplifier.simplify(expr)) + .context(EvalPartitionFilterSnafu)?; combined_exprs.push(logical_expr); region_partition_expr = Some(expr); } @@ -2339,6 +2357,75 @@ mod tests { .with_files(vec![file]) } + #[test] + fn test_partition_time_filter_after_unit_widening() { + use datatypes::arrow::array::TimestampMicrosecondArray; + + let bound = Value::Timestamp(Timestamp::new_millisecond(1)); + for (expr, expected) in [ + ( + partition_col("ts").lt(bound.clone()), + vec![true, false, false], + ), + (partition_col("ts").gt_eq(bound), vec![false, true, true]), + ] { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }); + builder.primary_key(vec![]); + builder.partition_expr_json(Some(expr.as_json_str().unwrap())); + let metadata = builder.build().unwrap(); + let group = PredicateGroup::new(&metadata, &[]).unwrap(); + let filters = group.time_filters().unwrap(); + let input: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![999, 1000, 1001])); + assert_eq!(filters.len(), 1); + assert_eq!( + filters[0].evaluate_array(&input).unwrap(), + datafusion_common::arrow::buffer::BooleanBuffer::from(expected), + ); + } + } + + #[test] + fn test_partition_time_filter_overflow_after_unit_widening() { + // Year 3000 fits milliseconds but overflows nanoseconds. + let bound = Value::Timestamp(Timestamp::new_millisecond(32_503_680_000_000)); + let expr = partition_col("ts").lt(bound); + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 0, + }); + builder.primary_key(vec![]); + builder.partition_expr_json(Some(expr.as_json_str().unwrap())); + let metadata = builder.build().unwrap(); + + let error = PredicateGroup::new(&metadata, &[]) + .err() + .expect("an overflowing partition bound must fail scan construction"); + let crate::error::Error::EvalPartitionFilter { error, .. } = error else { + panic!("unexpected error: {error:?}"); + }; + assert!( + error + .to_string() + .contains("converted value exceeds the representable i64 range"), + "{error}" + ); + } + #[tokio::test] async fn test_total_rows_is_exact_after_partition_filter() { let expr = partition_col("k0").gt_eq(Value::String("foo".into())); diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index 4362440e415..d5401f4a6de 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -341,8 +341,7 @@ pub(crate) fn build_bulk_filter_plan( predicate: Option<&Predicate>, ) -> BulkFilterPlan { let metadata = read_format.metadata(); - // Bulk memtable only needs simple binary filters here. Any filter that - // cannot be reduced to a SimpleFilterContext stays out of this fast path. + // Filters not recognized by SimpleFilterContext stay out of this fast path. let simple_filters: Vec = predicate .into_iter() .flat_map(|predicate| { diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 90c2728e35d..a3cbdef90f1 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -2014,7 +2014,12 @@ impl SimpleFilterContext { expected_meta: Option<&RegionMetadata>, expr: &Expr, ) -> Option { - let filter = SimpleFilterEvaluator::try_new(expr)?; + let filter = SimpleFilterEvaluator::try_new_with_column_type(expr, &|name| { + expected_meta + .unwrap_or(sst_meta) + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + })?; let expr_str = format!("{expr:?}"); let (column_metadata, maybe_filter) = match expected_meta { Some(meta) => { @@ -3159,6 +3164,15 @@ mod tests { )); } + #[test] + fn test_nullable_label_filter_uses_missing_sst_column_default() { + let metadata = Arc::new(sst_region_metadata()); + let expected = expected_metadata_with_reused_tag_name(&metadata); + let expr = col("tag_0").is_null().or(col("tag_0").not_eq(lit("lo"))); + let ctx = SimpleFilterContext::new_opt(&metadata, Some(&expected), &expr).unwrap(); + assert!(matches!(ctx.filter(), MaybeFilter::Matched)); + } + #[test] fn test_simple_filter_context_drops_mismatched_field_filter() { let (sst_metadata, latest_metadata) = mock_metadata(); diff --git a/src/query/src/dummy_catalog.rs b/src/query/src/dummy_catalog.rs index 5644194f1f9..8f0a2c7ac60 100644 --- a/src/query/src/dummy_catalog.rs +++ b/src/query/src/dummy_catalog.rs @@ -31,6 +31,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion_common::DataFusionError; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; use datatypes::arrow::datatypes::SchemaRef; +use datatypes::data_type::DataType; use datatypes::types::json_type::JsonNativeType; use futures::stream::BoxStream; use session::context::{QueryContext, QueryContextRef}; @@ -209,7 +210,13 @@ impl TableProvider for DummyTableProvider { .iter() .map(|e| { // Simple filter on primary key columns are precisely evaluated. - if let Some(simple_filter) = SimpleFilterEvaluator::try_new(e) { + if let Some(simple_filter) = + SimpleFilterEvaluator::try_new_with_column_type(e, &|name| { + self.metadata + .column_by_name(name) + .map(|column| column.column_schema.data_type.as_arrow_type()) + }) + { if self .metadata .column_by_name(simple_filter.column_name()) @@ -646,6 +653,37 @@ mod tests { FLOW_RETURN_REGION_SEQ, FLOW_SINK_TABLE_ID, }; + #[test] + fn test_nullable_label_filter_pushdown() { + use datafusion_expr::{Cast, col, lit}; + use datatypes::arrow::datatypes::DataType as ArrowDataType; + + let provider = crate::optimizer::test_util::mock_table_provider(RegionId::new(1, 1)); + let cast_label = Expr::Cast(Cast::new(Box::new(col("k0")), ArrowDataType::Utf8)); + let cast_field = Expr::Cast(Cast::new(Box::new(col("v0")), ArrowDataType::Utf8)); + let filters = [ + col("k0").is_null().or(col("k0").not_eq(lit("lo"))), + col("k0").is_null().or(cast_label.not_eq(lit("lo"))), + col("k0").is_null().or(col("v0").not_eq(lit(1.0))), + cast_field.eq(lit("1")), + col("k0").in_list(vec![lit(""), lit("lo")], false), + col("k0").in_list(vec![lit(""), lit("lo")], true), + ]; + assert_eq!( + provider + .supports_filters_pushdown(&filters.iter().collect::>()) + .unwrap(), + vec![ + TableProviderFilterPushDown::Exact, + TableProviderFilterPushDown::Exact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + ] + ); + } + fn test_region_id() -> RegionId { RegionId::new(1024, 1) } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 88464bc9ee8..07faf3b53dd 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -2557,28 +2557,21 @@ impl PromPlanner { // for it, or the column exists but is NULL on that row — the latter // is the norm for logical metrics sharing a physical table, which // holds the union of their label columns. + let mut null_label = None; let col = if let Some(column_name) = column_name { let column = DfExpr::Column(Column::from_name(&column_name)); let field = table_schema .index_of_column_by_name(None, &column_name) .map(|index| table_schema.field(index)); if accepts_empty - && let Some(data_type) = field - .filter(|field| { - field.is_nullable() - && Self::string_value_data_type(field.data_type()).is_some() - }) - .map(|field| field.data_type()) - { - let empty = Self::string_scalar_value(data_type, Some(String::new())) - .expect("nullable label has a string type"); - DfExpr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![column, DfExpr::Literal(empty, None)], + && field.is_some_and(|field| { + field.is_nullable() + && Self::string_value_data_type(field.data_type()).is_some() }) - } else { - column + { + null_label = Some(column.clone().is_null()); } + column } else { DfExpr::Literal(ScalarValue::Utf8(Some(String::new())), None) }; @@ -2631,6 +2624,10 @@ impl PromPlanner { } } }; + let expr = match null_label { + Some(null_label) => null_label.or(expr), + None => expr, + }; exprs.push(expr); } diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index f0c918351d1..9ca8a8b646c 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -535,6 +535,7 @@ async fn matchers_read_absent_labels_as_empty() { schema.clone(), vec![Arc::new(StringArray::from(vec![ None, + Some(""), Some(GREPTIME_TEMPORALITY_DELTA), ]))], ) @@ -542,10 +543,19 @@ async fn matchers_read_absent_labels_as_empty() { let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); for (matcher, expected) in [ - (r#"="""#, vec![None]), - (r#"!="delta""#, vec![None]), - (r#"=~".*""#, vec![None, Some(GREPTIME_TEMPORALITY_DELTA)]), - (r#"!~"delta""#, vec![None]), + (r#"="""#, vec![None, Some("")]), + (r#"!="delta""#, vec![None, Some("")]), + ( + r#"=~".*""#, + vec![None, Some(""), Some(GREPTIME_TEMPORALITY_DELTA)], + ), + (r#"!~"delta""#, vec![None, Some("")]), + (r#"!="""#, vec![Some(GREPTIME_TEMPORALITY_DELTA)]), + ( + r#"=~"delta|""#, + vec![None, Some(""), Some(GREPTIME_TEMPORALITY_DELTA)], + ), + (r#"!~"delta|""#, vec![]), (r#"="delta""#, vec![Some(GREPTIME_TEMPORALITY_DELTA)]), ] { let query = format!(r#"metric{{{marker}{matcher}}}"#); @@ -558,10 +568,10 @@ async fn matchers_read_absent_labels_as_empty() { }; let expressions = PromPlanner::matchers_to_expr(selector.matchers, scan.schema()).unwrap(); let display = expressions.iter().map(ToString::to_string).join(" AND "); - if matcher == r#"="delta""# { - assert!(!display.contains("coalesce"), "{display}"); + if [r#"="delta""#, r#"!="""#, r#"!~"delta|""#].contains(&matcher) { + assert!(!display.contains("IS NULL"), "{display}"); } else if !expressions.is_empty() { - assert!(display.contains("coalesce"), "{display}"); + assert!(display.contains("IS NULL"), "{display}"); } let plan = if let Some(filter) = conjunction(expressions) { LogicalPlanBuilder::from(scan) @@ -599,7 +609,7 @@ async fn matchers_read_absent_labels_as_empty() { // The rule is about NULL, not about the marker: any nullable label column // can be NULL where the series does not carry the label. A non-nullable one // has nothing to normalize. - for (nullable, wants_coalesce) in [(true, true), (false, false)] { + for (nullable, wants_null_filter) in [(true, true), (false, false)] { let ordinary_schema = Arc::new(ArrowSchema::new(vec![Field::new( "label", ArrowDataType::Utf8, @@ -626,8 +636,8 @@ async fn matchers_read_absent_labels_as_empty() { .map(ToString::to_string) .join(" AND "); assert_eq!( - wants_coalesce, - expressions.contains("coalesce"), + wants_null_filter, + expressions.contains("IS NULL"), "nullable={nullable}: {expressions}" ); } diff --git a/tests/cases/standalone/common/promql/regex.result b/tests/cases/standalone/common/promql/regex.result index c76b743eafa..c2a10626ff8 100644 --- a/tests/cases/standalone/common/promql/regex.result +++ b/tests/cases/standalone/common/promql/regex.result @@ -146,13 +146,93 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[1000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["host"] REDACTED -|_|_|_FilterExec: CASE WHEN host@1 IS NOT NULL THEN host@1 ELSE_END = REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["CASE WHEN host IS NOT NULL THEN host ELSE Dictionary(UInt32, Utf8(\"\")) END = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "host", "val"], "filters": ["host IS NULL OR host = Dictionary(UInt32, Utf8(\"\"))", "ts >= TimestampMillisecond(-299999, None)", "ts <= TimestampMillisecond(0, None)"], "REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ +INSERT INTO test VALUES (0, NULL, 2), (0, '', 3), (0, 'ext4', 4), (0, 'tmpfs', 5); + +Affected Rows: 4 + +TQL EVAL (0, 0, '1s') test{host!="tmpfs"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host!~"tmpfs|vfat"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host=""}; + ++---------------------+------+-----+ +| ts | host | val | ++---------------------+------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | ++---------------------+------+-----+ + +TQL EVAL (0, 0, '1s') test{host!=""}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{host=~"ext4|"}; + ++---------------------+------+-----+ +| ts | host | val | ++---------------------+------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | ext4 | 4 | ++---------------------+------+-----+ + +TQL EVAL (0, 0, '1s') test{host!~"ext4|"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + +TQL EVAL (0, 0, '1s') test{missing!="lo"}; + ++---------------------+-------------------+-----+ +| ts | host | val | ++---------------------+-------------------+-----+ +| 1970-01-01T00:00:00 | | 2 | +| 1970-01-01T00:00:00 | | 3 | +| 1970-01-01T00:00:00 | 10.0.160.237:8080 | 1 | +| 1970-01-01T00:00:00 | 10.0.160.237:8081 | 1 | +| 1970-01-01T00:00:00 | ext4 | 4 | +| 1970-01-01T00:00:00 | tmpfs | 5 | ++---------------------+-------------------+-----+ + DROP TABLE test; Affected Rows: 0 diff --git a/tests/cases/standalone/common/promql/regex.sql b/tests/cases/standalone/common/promql/regex.sql index 9804e098e64..80a8f3ed381 100644 --- a/tests/cases/standalone/common/promql/regex.sql +++ b/tests/cases/standalone/common/promql/regex.sql @@ -59,4 +59,14 @@ TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".*"}; -- SQLNESS REPLACE (flat_format.*) REDACTED TQL ANALYZE VERBOSE (0, 0, '1s') test{host!~".+"}; +INSERT INTO test VALUES (0, NULL, 2), (0, '', 3), (0, 'ext4', 4), (0, 'tmpfs', 5); + +TQL EVAL (0, 0, '1s') test{host!="tmpfs"}; +TQL EVAL (0, 0, '1s') test{host!~"tmpfs|vfat"}; +TQL EVAL (0, 0, '1s') test{host=""}; +TQL EVAL (0, 0, '1s') test{host!=""}; +TQL EVAL (0, 0, '1s') test{host=~"ext4|"}; +TQL EVAL (0, 0, '1s') test{host!~"ext4|"}; +TQL EVAL (0, 0, '1s') test{missing!="lo"}; + DROP TABLE test; From 7f0f592b6b8be81acdaf3ac057b3a741bc7ebe97 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 8 Oct 2026 09:00:10 +0000 Subject: [PATCH 102/133] chore: update rust toolchain to 1.98.1 (#9429) --- flake.lock | 18 ++++++++--------- flake.nix | 2 +- rust-toolchain.toml | 2 +- .../function/src/scalars/vector/impl_conv.rs | 6 ++++-- src/common/meta/src/instruction.rs | 4 ++-- src/common/meta/src/kv_backend/chroot.rs | 2 +- src/common/meta/src/rpc/router.rs | 2 +- src/common/meta/src/rpc/store.rs | 6 +++--- src/common/meta/src/sequence.rs | 4 ++-- src/common/query/src/stream.rs | 2 +- src/datatypes/src/types/vector_type.rs | 6 ++++-- src/flow/src/batching_mode/utils/test.rs | 1 - src/index/src/bloom_filter/creator.rs | 6 ++++-- src/log-store/src/kafka/consumer.rs | 11 ++++------ .../src/handler/collect_stats_handler.rs | 2 +- .../src/procedure/repartition/utils.rs | 2 +- src/meta-srv/src/service/heartbeat.rs | 1 - .../src/row_converter/sparse/checked.rs | 2 +- src/mito2/benches/memtable_bench.rs | 2 +- src/mito2/src/compaction/scheduler_test.rs | 2 +- src/mito2/src/gc/worker_test.rs | 10 +++------- src/mito2/src/memtable/bulk.rs | 12 +++++------ src/mito2/src/sst/file_ref.rs | 8 ++------ .../src/etl/transform/transformer/greptime.rs | 2 +- .../src/extension_plan/histogram_fold.rs | 2 +- src/query/src/dist_plan/analyzer.rs | 2 +- src/servers/src/http/jaeger.rs | 2 +- src/servers/src/prom_remote_write/v2.rs | 2 +- src/sql/src/parsers/create_parser.rs | 4 ++-- src/sql/src/statements/copy.rs | 10 +++++----- src/sql/src/statements/create.rs | 20 +++++++++---------- src/sql/src/statements/option_map.rs | 4 ++-- src/sql/src/statements/show.rs | 8 ++++---- src/table/src/metadata.rs | 2 +- src/table/src/predicate.rs | 11 ++++------ src/table/src/test_util/memtable.rs | 1 - tests-integration/src/cluster.rs | 2 +- tests-integration/src/standalone.rs | 2 +- tests-integration/src/tests/instance_test.rs | 8 ++++---- 39 files changed, 92 insertions(+), 103 deletions(-) diff --git a/flake.lock b/flake.lock index 27c1797f10b..621f1e2f3a0 100644 --- a/flake.lock +++ b/flake.lock @@ -8,11 +8,11 @@ "rust-analyzer-src": "rust-analyzer-src" }, "locked": { - "lastModified": 1785314864, - "narHash": "sha256-imz9J5iMNersDeQ1ymenFg35N6C2R3MCaGg2i0XCIrE=", + "lastModified": 1790642863, + "narHash": "sha256-JIwE6EnebvxW7QFFEitWHuZDXwN6Qcf6ZC8DIvYW3y8=", "owner": "nix-community", "repo": "fenix", - "rev": "594418b2c9ea0731bee6988236ef5bfd97e81dca", + "rev": "a8f386c482e78e9a48be5abecdc17a9b098a95a1", "type": "github" }, "original": { @@ -41,11 +41,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1785133411, - "narHash": "sha256-Yjv0WEg39KRYS0rBdTbu6Fc/or/ihAKk13W9sQ6VWd0=", + "lastModified": 1790533781, + "narHash": "sha256-F+b8u43BVF5DlM0pz7JQSbtbwKy5PrHIT5zQUekXsD0=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "2f5a153c270b70cb0f8c11f46d96d6d3bc39f4e3", + "rev": "cf5e76507c6e23b59f7e0ffcc7baa2a39ddd8442", "type": "github" }, "original": { @@ -65,11 +65,11 @@ "rust-analyzer-src": { "flake": false, "locked": { - "lastModified": 1785261141, - "narHash": "sha256-sA+DHPejWD68mLA7gtTiHGd5T41F6W+NKf0g+ibFCW4=", + "lastModified": 1790532271, + "narHash": "sha256-j8y/AO3AMDF9p+11aIDhnA21JnNj5PjvcejTI43vI8Q=", "owner": "rust-lang", "repo": "rust-analyzer", - "rev": "bec66814323579659ffd77c909b3d963af118ece", + "rev": "03fcb77246f2568adb0e9b2fa60d19c6cc1686f4", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index 6e4151fb3e8..df71403b24f 100644 --- a/flake.nix +++ b/flake.nix @@ -22,7 +22,7 @@ lib = nixpkgs.lib; rustToolchain = fenix.packages.${system}.fromToolchainName { name = (lib.importTOML ./rust-toolchain.toml).toolchain.channel; - sha256 = "sha256-h+t2xTBz5yt2YIO+1VMIIGlCU7gyp2LYOFvaV1nwOXU="; + sha256 = "sha256-p8h3Sl/YRByZfZTAKXdsvF6xEenXKrXSVvpphmZENH4="; }; in { diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 2d45363a5be..642220914d8 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "1.96.1" +channel = "1.98.1" diff --git a/src/common/function/src/scalars/vector/impl_conv.rs b/src/common/function/src/scalars/vector/impl_conv.rs index 219361a302c..00d23cedf7f 100644 --- a/src/common/function/src/scalars/vector/impl_conv.rs +++ b/src/common/function/src/scalars/vector/impl_conv.rs @@ -53,8 +53,10 @@ pub fn binlit_as_veclit(bytes: &[u8]) -> Result> { }) } else { let v = bytes - .chunks_exact(std::mem::size_of::()) - .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|chunk| f32::from_le_bytes(*chunk)) .collect::>(); Ok(Cow::Owned(v)) } diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index 60ebe14e37c..6ae074b0ac9 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -732,12 +732,12 @@ impl PackedRegionFileRefs { } let mut refs = HashSet::new(); - for bytes in files.chunks_exact(16) { + for bytes in files.as_chunks::<16>().0 { let mut id = [0; 16]; id.copy_from_slice(bytes); refs.insert(FileRef::new(region, FileId::from_bytes(id), None)); } - for bytes in indexed.chunks_exact(24) { + for bytes in indexed.as_chunks::<24>().0 { let mut id = [0; 16]; id.copy_from_slice(&bytes[..16]); let mut version = [0; 8]; diff --git a/src/common/meta/src/kv_backend/chroot.rs b/src/common/meta/src/kv_backend/chroot.rs index 1a7e7c615ac..4d2d16d0f2f 100644 --- a/src/common/meta/src/kv_backend/chroot.rs +++ b/src/common/meta/src/kv_backend/chroot.rs @@ -82,7 +82,7 @@ impl KvBackend for ChrootKvBackend { async fn batch_put(&self, mut req: BatchPutRequest) -> Result { for kv in req.kvs.iter_mut() { - kv.key = self.key_prepend_root(kv.key.drain(..).collect()); + kv.key = self.key_prepend_root(std::mem::take(&mut kv.key)); } let mut res = self.inner.batch_put(req).await?; res.prev_kvs = res diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 929f57515ed..20fcd955abb 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -48,7 +48,7 @@ pub fn region_distribution(region_routes: &[RegionRoute]) -> RegionDistribution .add_follower_region(region_number); } } - for (_, region_role_set) in regions_id_map.iter_mut() { + for region_role_set in regions_id_map.values_mut() { // Sort the regions in ascending order. region_role_set.sort() } diff --git a/src/common/meta/src/rpc/store.rs b/src/common/meta/src/rpc/store.rs index 25536993972..a9c36199761 100644 --- a/src/common/meta/src/rpc/store.rs +++ b/src/common/meta/src/rpc/store.rs @@ -199,7 +199,7 @@ impl RangeResponse { } pub fn take_kvs(&mut self) -> Vec { - self.kvs.drain(..).collect() + std::mem::take(&mut self.kvs) } } @@ -451,7 +451,7 @@ impl BatchPutResponse { } pub fn take_prev_kvs(&mut self) -> Vec { - self.prev_kvs.drain(..).collect() + std::mem::take(&mut self.prev_kvs) } } @@ -780,7 +780,7 @@ impl DeleteRangeResponse { } pub fn take_prev_kvs(&mut self) -> Vec { - self.prev_kvs.drain(..).collect() + std::mem::take(&mut self.prev_kvs) } } diff --git a/src/common/meta/src/sequence.rs b/src/common/meta/src/sequence.rs index cd0a8ebf881..9d56170833c 100644 --- a/src/common/meta/src/sequence.rs +++ b/src/common/meta/src/sequence.rs @@ -188,7 +188,7 @@ impl Inner { } error::NextSequenceSnafu { - err_msg: format!("{}.next()", &self.name), + err_msg: format!("{}.next()", self.name), } .fail() } @@ -264,7 +264,7 @@ impl Inner { } error::NextSequenceSnafu { - err_msg: format!("{}.next_range()", &self.name), + err_msg: format!("{}.next_range()", self.name), } .fail() } diff --git a/src/common/query/src/stream.rs b/src/common/query/src/stream.rs index 0a46164d9b4..d6f1735e06c 100644 --- a/src/common/query/src/stream.rs +++ b/src/common/query/src/stream.rs @@ -95,7 +95,7 @@ impl DisplayAs for StreamScanAdapter { f, "StreamScanAdapter: [], schema: [" )?; - write!(f, "{:?}", &self.arrow_schema)?; + write!(f, "{:?}", self.arrow_schema)?; write!(f, "]") } } diff --git a/src/datatypes/src/types/vector_type.rs b/src/datatypes/src/types/vector_type.rs index d01949c4737..ab756be90eb 100644 --- a/src/datatypes/src/types/vector_type.rs +++ b/src/datatypes/src/types/vector_type.rs @@ -86,8 +86,10 @@ pub fn vector_type_value_to_string(val: &[u8], dim: u32) -> Result { } let elements = val - .chunks_exact(std::mem::size_of::()) - .map(|e| f32::from_le_bytes(e.try_into().unwrap())); + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|e| f32::from_le_bytes(*e)); let mut s = String::from("["); for (i, e) in elements.enumerate() { diff --git a/src/flow/src/batching_mode/utils/test.rs b/src/flow/src/batching_mode/utils/test.rs index 1f390d1594c..1af8880a7d4 100644 --- a/src/flow/src/batching_mode/utils/test.rs +++ b/src/flow/src/batching_mode/utils/test.rs @@ -22,7 +22,6 @@ use common_recordbatch::recordbatch::merge_record_batches; use common_recordbatch::{RecordBatch, util}; use common_time::Timestamp; use datafusion_common::ScalarValue; -use datafusion_common::tree_node::TreeNode as _; use datafusion_expr::{GroupingSet, TableScanBuilder}; use datatypes::arrow::array::{Array, AsArray}; use datatypes::arrow::datatypes::{Float64Type, Int64Type, UInt64Type}; diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index 83c0ae0c697..dc1dc54d289 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/src/index/src/bloom_filter/creator.rs @@ -270,8 +270,10 @@ mod tests { /// Converts a slice of bytes to a vector of `u64`. pub fn u64_vec_from_bytes(bytes: &[u8]) -> Vec { bytes - .chunks_exact(std::mem::size_of::()) - .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap())) + .as_chunks::<{ std::mem::size_of::() }>() + .0 + .iter() + .map(|chunk| u64::from_le_bytes(*chunk)) .collect() } diff --git a/src/log-store/src/kafka/consumer.rs b/src/log-store/src/kafka/consumer.rs index 8a7cba5a4b6..d064a9bc9ac 100644 --- a/src/log-store/src/kafka/consumer.rs +++ b/src/log-store/src/kafka/consumer.rs @@ -158,13 +158,10 @@ impl RecordsBuffer { impl RecordsBuffer { fn pop_front(&mut self) -> Option { while let Some(index) = self.index.peek() { - if let Some(record_and_offset) = self.buffer.pop_front() { - if index == record_and_offset.offset as u64 { - self.index.next(); - return Some(record_and_offset); - } - } else { - return None; + let record_and_offset = self.buffer.pop_front()?; + if index == record_and_offset.offset as u64 { + self.index.next(); + return Some(record_and_offset); } } diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index e3ae7ebe915..cc653f9982c 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -40,7 +40,7 @@ struct EpochStats { impl EpochStats { #[inline] fn drain_all(&mut self) -> Vec { - self.stats.drain(..).collect() + std::mem::take(&mut self.stats) } #[inline] diff --git a/src/meta-srv/src/procedure/repartition/utils.rs b/src/meta-srv/src/procedure/repartition/utils.rs index 50dbc57f9a8..cf6ea726f34 100644 --- a/src/meta-srv/src/procedure/repartition/utils.rs +++ b/src/meta-srv/src/procedure/repartition/utils.rs @@ -82,7 +82,7 @@ pub fn merge_and_validate_region_wal_options( table_id: TableId, ) -> Result { // Doesn't allow overwriting existing WAL options. - for (region_number, _) in new_region_wal_options.iter() { + for region_number in new_region_wal_options.keys() { if region_wal_options.contains_key(region_number) { return error::UnexpectedSnafu { violated: format!( diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs index 5cdc5b57894..782daa289de 100644 --- a/src/meta-srv/src/service/heartbeat.rs +++ b/src/meta-srv/src/service/heartbeat.rs @@ -419,7 +419,6 @@ mod tests { use std::sync::Arc; use api::v1::meta::heartbeat_server::Heartbeat; - use api::v1::meta::*; use common_meta::kv_backend::memory::MemoryKvBackend; use common_telemetry::tracing_context::W3cTrace; use servers::grpc::GrpcOptions; diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs index 0cc9ed8211a..fec89b74e0d 100644 --- a/src/mito-codec/src/row_converter/sparse/checked.rs +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -116,7 +116,7 @@ impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { // encoded_value checked the markers and every chunk length. Reserve once // before unchunking so the buffer does not grow for each 8-byte chunk. buffer.reserve(encoded.len() - 2); - for chunk in encoded[2..].chunks_exact(9) { + for chunk in encoded[2..].as_chunks::<9>().0 { let len = usize::from(chunk[8]).min(8); buffer.extend_from_slice(&chunk[..len]); } diff --git a/src/mito2/benches/memtable_bench.rs b/src/mito2/benches/memtable_bench.rs index 80331941865..e4d71c9883c 100644 --- a/src/mito2/benches/memtable_bench.rs +++ b/src/mito2/benches/memtable_bench.rs @@ -268,7 +268,7 @@ fn filter_1_host(c: &mut Criterion) { ) .unwrap(); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { let iter = range.build_record_batch_iter(None, None).unwrap(); for batch in iter { let _batch = batch.unwrap(); diff --git a/src/mito2/src/compaction/scheduler_test.rs b/src/mito2/src/compaction/scheduler_test.rs index 5afcc9c37fd..42c1df23f07 100644 --- a/src/mito2/src/compaction/scheduler_test.rs +++ b/src/mito2/src/compaction/scheduler_test.rs @@ -546,7 +546,7 @@ async fn test_planning_panic_notifies_and_clears_status() { let finished = recv_compaction_pick_finished(&mut rx).await; let CompactionPlanningResult::Error(err) = &finished.result else { - panic!("expected planning error, got {:?}", &finished.result); + panic!("expected planning error, got {:?}", finished.result); }; assert!(err.to_string().contains("planning boom")); let pending_ddls = scheduler diff --git a/src/mito2/src/gc/worker_test.rs b/src/mito2/src/gc/worker_test.rs index 7bcf511f1a0..2a66294f59b 100644 --- a/src/mito2/src/gc/worker_test.rs +++ b/src/mito2/src/gc/worker_test.rs @@ -518,13 +518,9 @@ async fn write_and_list_entry(store: &ObjectStore, path: &str) -> Entry { .await .unwrap(); // List the parent directory to get the entry. - let parent = std::path::Path::new(path).parent().and_then(|p| { - if p.as_os_str().is_empty() { - None - } else { - Some(p) - } - }); + let parent = std::path::Path::new(path) + .parent() + .filter(|p| !p.as_os_str().is_empty()); let prefix = match parent { Some(p) => format!("{}/", p.to_str().unwrap()), None => String::new(), // root diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 94ccac472df..afede8904d8 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -1872,7 +1872,7 @@ mod tests { let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum(); assert_eq!(5, total_rows); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.num_rows() > 0); assert!(range.is_record_batch()); @@ -2468,7 +2468,7 @@ mod tests { let total_rows: usize = ranges.ranges.values().map(|r| r.stats().num_rows()).sum(); assert_eq!(10, total_rows); - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.num_rows() > 0); assert!(range.is_record_batch()); @@ -2554,7 +2554,7 @@ mod tests { // Read all data and verify let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); @@ -2642,7 +2642,7 @@ mod tests { assert_eq!(13, total_rows); let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); for batch_result in record_batch_iter { let batch = batch_result.unwrap(); @@ -3063,7 +3063,7 @@ mod tests { // Read all data let mut total_rows_read = 0; - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); @@ -3120,7 +3120,7 @@ mod tests { // Should return ranges but each range should produce an empty iterator // instead of an error. - for (_range_id, range) in ranges.ranges.iter() { + for range in ranges.ranges.values() { assert!(range.is_record_batch()); let record_batch_iter = range.build_record_batch_iter(None, None).unwrap(); let total_rows: usize = record_batch_iter.map(|r| r.unwrap().num_rows()).sum(); diff --git a/src/mito2/src/sst/file_ref.rs b/src/mito2/src/sst/file_ref.rs index 2f067cf3c19..28ff116cf25 100644 --- a/src/mito2/src/sst/file_ref.rs +++ b/src/mito2/src/sst/file_ref.rs @@ -77,12 +77,8 @@ impl FileReferenceManager { } fn ref_file_set(&self, region_id: RegionId) -> Option> { - let file_refs = if let Some(file_refs) = self.files_per_region.get(®ion_id) { - file_refs.clone() - } else { - // region id not found. - return None; - }; + // region id not found. + let file_refs = self.files_per_region.get(®ion_id)?.clone(); if file_refs.files.is_empty() { // still return an empty manifest to indicate no files are referenced. diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index eb25071c91c..16c9e3399b6 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -911,7 +911,7 @@ fn identity_pipeline_inner( } let column_count = schema_info.schema.len(); - for (_, row) in opt_map.iter_mut() { + for row in opt_map.values_mut() { for row in row.iter_mut() { assert!( column_count >= row.values.len(), diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 3dc51e27bcd..709117df294 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -1218,7 +1218,7 @@ impl HistogramFoldStream { fn flush_remaining(&mut self) -> DataFusionResult<()> { if self.mode == FoldMode::Optimistic && self.input_buffered_rows > 0 { - let buffered_batches: Vec<_> = self.input_buffer.drain(..).collect(); + let buffered_batches: Vec<_> = std::mem::take(&mut self.input_buffer); if !buffered_batches.is_empty() { let batch = concat_batches(&self.input_schema, buffered_batches.as_slice())?; self.switch_to_safe_mode(batch)?; diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 2c1d066054a..da26d88b433 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -653,7 +653,7 @@ impl PlanRewriter { )) })?; - for (_col_name, alias_set) in part_cols.iter_mut() { + for alias_set in part_cols.values_mut() { let aliased_cols = aliased_columns_for( &alias_set.clone().into_iter().collect(), plan, diff --git a/src/servers/src/http/jaeger.rs b/src/servers/src/http/jaeger.rs index 32ef6a61682..16dbaf76e53 100644 --- a/src/servers/src/http/jaeger.rs +++ b/src/servers/src/http/jaeger.rs @@ -328,7 +328,7 @@ impl QueryTraceParams { } .build() })?; - for (_, v) in tags_map.iter_mut() { + for v in tags_map.values_mut() { if let Some(number) = convert_string_to_number(v) { *v = number; } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index b9e34dc7f24..9178748bc77 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -658,7 +658,7 @@ fn resolve_series_labels<'a>( let mut tags = Vec::with_capacity(labels_refs.len() / 2); label_names.clear(); - for pair in labels_refs.chunks_exact(2) { + for pair in labels_refs.as_chunks::<2>().0 { let name = symbol_ref(symbols, pair[0], "label name")?; let value = symbol_ref(symbols, pair[1], "label value")?; validate_label(name)?; diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index d701415f3d2..32606679a1c 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -768,7 +768,7 @@ impl<'a> ParserContext<'a> { { return Err(ParserError::ParserError(format!( "Cannot use keyword '{}' as column name. Hint: add quotes to the name.", - &name.value + name.value ))); } @@ -786,7 +786,7 @@ impl<'a> ParserContext<'a> { InvalidSqlSnafu { msg: format!( "Cannot use keyword '{}' as column name. Hint: add quotes to the name.", - &name.value + name.value ), } ); diff --git a/src/sql/src/statements/copy.rs b/src/sql/src/statements/copy.rs index 7a59175538c..f47f26ec61a 100644 --- a/src/sql/src/statements/copy.rs +++ b/src/sql/src/statements/copy.rs @@ -46,7 +46,7 @@ pub struct CopyQueryTo { impl Display for CopyQueryTo { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "COPY ({}) TO {}", &self.query, &self.arg.location)?; + write!(f, "COPY ({}) TO {}", self.query, self.arg.location)?; if !self.arg.with.is_empty() { let options = self.arg.with.kv_pairs(); write!(f, " WITH ({})", options.join(", "))?; @@ -70,11 +70,11 @@ impl Display for CopyTable { write!(f, "COPY ")?; let (with, connection) = match self { CopyTable::To(args) => { - write!(f, "{} TO {}", &args.table_name, &args.location)?; + write!(f, "{} TO {}", args.table_name, args.location)?; (&args.with, &args.connection) } CopyTable::From(args) => { - write!(f, "{} FROM {}", &args.table_name, &args.location)?; + write!(f, "{} FROM {}", args.table_name, args.location)?; (&args.with, &args.connection) } }; @@ -101,11 +101,11 @@ impl Display for CopyDatabase { write!(f, "COPY DATABASE ")?; let (with, connection) = match self { CopyDatabase::To(args) => { - write!(f, "{} TO {}", &args.database_name, &args.location)?; + write!(f, "{} TO {}", args.database_name, args.location)?; (&args.with, &args.connection) } CopyDatabase::From(args) => { - write!(f, "{} FROM {}", &args.database_name, &args.location)?; + write!(f, "{} FROM {}", args.database_name, args.location)?; (&args.with, &args.connection) } }; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 725743ee679..9bd6c341a65 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -456,7 +456,7 @@ impl CreateTable { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{} (", &self.name)?; + writeln!(f, "{} (", self.name)?; for (index, column) in self.columns.iter().enumerate() { if index > 0 { f.write_str(LINE_SEP)?; @@ -471,7 +471,7 @@ impl CreateTable { if let Some(partitions) = &self.partitions { writeln!(f, "{partitions}")?; } - writeln!(f, "ENGINE={}", &self.engine)?; + writeln!(f, "ENGINE={}", self.engine)?; if !self.options.is_empty() { f.write_str("WITH(\n")?; self.options.write_sql(f, LINE_SEP, " ", complete)?; @@ -506,7 +506,7 @@ impl Display for CreateDatabase { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - write!(f, "{}", &self.name)?; + write!(f, "{}", self.name)?; if !self.options.is_empty() { let options = self.options.kv_pairs(); write!(f, "\nWITH(\n{}\n)", format_list_indent!(options))?; @@ -533,11 +533,11 @@ impl Display for CreateExternalTable { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{} (", &self.name)?; + writeln!(f, "{} (", self.name)?; writeln!(f, "{},", format_list_indent!(self.columns))?; writeln!(f, "{}", format_table_constraint(&self.constraints))?; writeln!(f, ")")?; - writeln!(f, "ENGINE={}", &self.engine)?; + writeln!(f, "ENGINE={}", self.engine)?; if !self.options.is_empty() { let options = self.options.kv_pairs(); write!(f, "WITH(\n{}\n)", format_list_indent!(options))?; @@ -635,8 +635,8 @@ impl Display for CreateFlow { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - writeln!(f, "{}", &self.flow_name)?; - writeln!(f, "SINK TO {}", &self.sink_table_name)?; + writeln!(f, "{}", self.flow_name)?; + writeln!(f, "SINK TO {}", self.sink_table_name)?; if let Some(expire_after) = &self.expire_after { writeln!(f, "EXPIRE AFTER '{} s'", expire_after)?; } @@ -657,7 +657,7 @@ impl Display for CreateFlow { let options = self.flow_options.kv_pairs(); writeln!(f, "WITH ({})", format_list_comma!(options))?; } - write!(f, "AS {}", &self.query) + write!(f, "AS {}", self.query) } } @@ -687,11 +687,11 @@ impl Display for CreateView { if self.if_not_exists { write!(f, "IF NOT EXISTS ")?; } - write!(f, "{} ", &self.name)?; + write!(f, "{} ", self.name)?; if !self.columns.is_empty() { write!(f, "({}) ", format_list_comma!(self.columns))?; } - write!(f, "AS {}", &self.query) + write!(f, "AS {}", self.query) } } diff --git a/src/sql/src/statements/option_map.rs b/src/sql/src/statements/option_map.rs index 5077088b67a..29808ae602d 100644 --- a/src/sql/src/statements/option_map.rs +++ b/src/sql/src/statements/option_map.rs @@ -258,10 +258,10 @@ impl Visit for OptionMap { impl VisitMut for OptionMap { fn visit(&mut self, visitor: &mut V) -> ControlFlow { - for (_, v) in self.options.iter_mut() { + for v in self.options.values_mut() { v.visit(visitor)?; } - for (_, v) in self.secrets.iter_mut() { + for v in self.secrets.values_mut() { v.expose_secret_mut().visit(visitor)?; } ControlFlow::Continue(()) diff --git a/src/sql/src/statements/show.rs b/src/sql/src/statements/show.rs index 67624192b82..d1531de360d 100644 --- a/src/sql/src/statements/show.rs +++ b/src/sql/src/statements/show.rs @@ -41,7 +41,7 @@ impl Display for ShowKind { macro_rules! format_kind { ($self: expr, $f: expr) => { if $self.kind != ShowKind::All { - write!($f, " {}", &$self.kind)?; + write!($f, " {}", $self.kind)?; } }; } @@ -71,7 +71,7 @@ impl Display for ShowColumns { if self.full { write!(f, "FULL ")?; } - write!(f, "COLUMNS IN {}", &self.table)?; + write!(f, "COLUMNS IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } @@ -90,7 +90,7 @@ pub struct ShowIndex { impl Display for ShowIndex { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "SHOW INDEX IN {}", &self.table)?; + write!(f, "SHOW INDEX IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } @@ -110,7 +110,7 @@ pub struct ShowRegion { impl Display for ShowRegion { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "SHOW REGION IN {}", &self.table)?; + write!(f, "SHOW REGION IN {}", self.table)?; if let Some(database) = &self.database { write!(f, " IN {database}")?; } diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index d91771b9c5f..c8941bb1a82 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -1143,7 +1143,7 @@ impl TableMeta { .with_context(|_| error::CastDefaultValueSnafu { reason: format!( "Failed to cast default value from {:?} to type {:?}", - default_value, &change_column.target_type + default_value, change_column.target_type ), })?, ) diff --git a/src/table/src/predicate.rs b/src/table/src/predicate.rs index 7290c35c70b..8dc85b610f4 100644 --- a/src/table/src/predicate.rs +++ b/src/table/src/predicate.rs @@ -505,13 +505,10 @@ fn extract_from_in_list_expr( for expr in list { if let Expr::Literal(scalar, _) = expr { return_none_if_utf8!(scalar); - if let Some(timestamp) = scalar_value_to_timestamp(scalar, None) { - init_range = init_range.or(&TimestampRange::single(timestamp)) - } else { - // TODO(hl): maybe we should raise an error here since cannot parse - // timestamp value from in list expr - return None; - } + // TODO(hl): maybe we should raise an error here since cannot parse + // timestamp value from in list expr + let timestamp = scalar_value_to_timestamp(scalar, None)?; + init_range = init_range.or(&TimestampRange::single(timestamp)); } } Some(init_range) diff --git a/src/table/src/test_util/memtable.rs b/src/table/src/test_util/memtable.rs index 0ae751d35f4..a0824facd0d 100644 --- a/src/table/src/test_util/memtable.rs +++ b/src/table/src/test_util/memtable.rs @@ -188,7 +188,6 @@ impl Stream for MemtableStream { #[cfg(test)] mod test { use common_recordbatch::util; - use datatypes::prelude::*; use datatypes::schema::ColumnSchema; use datatypes::vectors::{Helper, Int32Vector, StringVector}; diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 4f762ba709b..21a0d5eba2d 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -380,7 +380,7 @@ impl GreptimeDbClusterBuilder { let home_dir = if let Some(home_dir) = &self.shared_home_dir { home_dir.path().to_str().unwrap().to_string() } else { - let home_tmp_dir = create_temp_dir(&format!("gt_home_{}", &self.cluster_name)); + let home_tmp_dir = create_temp_dir(&format!("gt_home_{}", self.cluster_name)); let home_dir = home_tmp_dir.path().to_str().unwrap().to_string(); guards.push(TestGuard { home_guard: FileDirGuard::new(home_tmp_dir), diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 5a36d495df8..98eb336d9cc 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -385,7 +385,7 @@ impl GreptimeDbStandaloneBuilder { let procedure_config = ProcedureConfig::default(); let kv_backend = standalone::build_metadata_kvbackend( - format!("{}/kv", &opts.storage.data_home), + format!("{}/kv", opts.storage.data_home), kv_backend_config, ) .unwrap(); diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 7cc02cfc456..a1ccac27328 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -2648,7 +2648,7 @@ async fn test_execute_copy_from_orc_with_cast(instance: Arc) { let output = execute_sql( &instance, - &format!("copy demo from '{}' WITH(FORMAT='orc');", &filepath), + &format!("copy demo from '{}' WITH(FORMAT='orc');", filepath), ) .await .data; @@ -2687,7 +2687,7 @@ async fn test_execute_copy_from_orc(instance: Arc) { let output = execute_sql( &instance, - &format!("copy demo from '{}' WITH(FORMAT='orc');", &filepath), + &format!("copy demo from '{}' WITH(FORMAT='orc');", filepath), ) .await .data; @@ -2724,7 +2724,7 @@ async fn test_cast_type_issue_1594(instance: Arc) { let output = execute_sql( &instance, - &format!("copy tsbs_cpu from '{}' WITH(FORMAT='csv');", &filepath), + &format!("copy tsbs_cpu from '{}' WITH(FORMAT='csv');", filepath), ) .await; @@ -2762,7 +2762,7 @@ async fn test_copy_from_csv_skip_bad_records(instance: Arc) { &instance, &format!( "copy csv_skip_bad_records from '{}' WITH(FORMAT='csv', skip_bad_records='true');", - &filepath + filepath ), ) .await From c46fbde9acd6f87bfcb7d0fa12d79e80d871b17b Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:23:25 +0000 Subject: [PATCH 103/133] fix(json2): preserve metadata when adding columns via alter table (#8696) * fix(json2): support adding columns via alter table Signed-off-by: fys * fix(json2): scope missing-column null handling to JSON2 --------- Signed-off-by: fys --- Cargo.lock | 1 + src/mito2/src/read/compat.rs | 44 +++++++----- src/operator/src/expr_helper.rs | 24 +++++++ src/sql/Cargo.toml | 1 + src/sql/src/statements.rs | 72 ++++++++++++++----- .../standalone/common/types/json/json2.result | 62 ++++++++++++++++ .../standalone/common/types/json/json2.sql | 30 ++++++++ .../common/types/json/json2_limit.result | 21 ++++++ .../common/types/json/json2_limit.sql | 13 ++++ 9 files changed, 234 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6ec2a2d2764..79e325b19b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14005,6 +14005,7 @@ version = "1.4.0-alpha.0" dependencies = [ "api", "arrow-buffer 59.2.0", + "arrow-schema 59.2.0", "chrono", "common-base", "common-catalog", diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index 0e566b7e078..acc05f0e9d2 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use api::v1::SemanticType; use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, BinaryBuilder, DictionaryArray, UInt32Array, + Array, ArrayRef, BinaryArray, BinaryBuilder, DictionaryArray, UInt32Array, new_null_array, }; use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; @@ -27,7 +27,7 @@ use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::DataType; use datatypes::value::Value; -use datatypes::vectors::VectorRef; +use datatypes::vectors::{Helper, VectorRef}; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -192,20 +192,32 @@ impl FlatCompatBatch { }); } else { // Create a default vector with 1 element for that column. - let default_vector = expect_column - .column_schema - .create_default_vector(1) - .context(CreateDefaultSnafu { - region_id: expect_metadata.region_id, - column: &expect_column.column_schema.name, - })? - .with_context(|| CompatReaderSnafu { - region_id: expect_metadata.region_id, - reason: format!( - "column {} does not have a default value to read", - expect_column.column_schema.name - ), - })?; + let default_vector = if expect_data_type.is_json2() { + // JSON2 columns do not support defaults, so missing columns use NULL. + // `expect_data_type` may be a query-concretized JSON2 type, + // while region metadata still carries its unconcretized type. + // The default vector must match the target Arrow field above. + Helper::try_into_vector(new_null_array(&expect_data_type.as_arrow_type(), 1)) + .context(CreateDefaultSnafu { + region_id: expect_metadata.region_id, + column: &expect_column.column_schema.name, + })? + } else { + expect_column + .column_schema + .create_default_vector(1) + .context(CreateDefaultSnafu { + region_id: expect_metadata.region_id, + column: &expect_column.column_schema.name, + })? + .with_context(|| CompatReaderSnafu { + region_id: expect_metadata.region_id, + reason: format!( + "column {} does not have a default value to read", + expect_column.column_schema.name + ), + })? + }; index_or_defaults.push(IndexOrDefault::DefaultValue { default_vector, semantic_type: expect_column.semantic_type, diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 1d06c24ff7f..02c1af46706 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -1192,6 +1192,7 @@ mod tests { use std::collections::HashMap; use api::v1::{SetDatabaseOptions, UnsetDatabaseOptions}; + use datatypes::extension::json::is_json2_extension_type; use datatypes::value::Value; use session::context::{QueryContext, QueryContextBuilder}; use sql::dialect::GreptimeDbDialect; @@ -1201,6 +1202,29 @@ mod tests { use super::*; + #[test] + fn test_to_alter_add_json2_column_preserves_extension() { + let sql = "ALTER TABLE monitor ADD COLUMN payload JSON2;"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = expr.kind.unwrap() else { + unreachable!() + }; + let column_def = add_columns[0].column_def.as_ref().unwrap(); + let column_schema = try_as_column_schema(column_def).unwrap(); + let schema = Schema::new(vec![column_schema]); + + assert!(is_json2_extension_type(&schema.arrow_schema().fields()[0])); + } + #[test] fn test_create_flow_tql_expr() { let sql = r#" diff --git a/src/sql/Cargo.toml b/src/sql/Cargo.toml index 685d96336b1..de972bec30e 100644 --- a/src/sql/Cargo.toml +++ b/src/sql/Cargo.toml @@ -13,6 +13,7 @@ enterprise = [] [dependencies] api.workspace = true arrow-buffer.workspace = true +arrow-schema.workspace = true chrono.workspace = true common-base.workspace = true common-catalog.workspace = true diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 2ea716678c7..8a001510045 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -36,7 +36,10 @@ pub mod truncate; use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; -use api::v1::SemanticType; +use api::v1::{ColumnOptions, SemanticType}; +use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, +}; use common_sql::default_constraint::parse_column_default_constraint; use common_time::timezone::Timezone; use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; @@ -148,25 +151,26 @@ pub fn column_to_schema( column_schema.set_inverted_index(column.extensions.inverted_index_options.is_some()); - let is_json2_column = if let SqlDataType::Custom(object_name, _) = column.data_type() { - object_name - .0 - .first() - .map(|x| x.to_string_unquoted().eq_ignore_ascii_case(JSON2_TYPE_NAME)) - .unwrap_or_default() - } else { - false - }; - if is_json2_column { - let settings = column - .extensions - .build_json_settings()? - .unwrap_or_else(JsonSettings::new_v2); - let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(settings))); + set_json2_extension(&mut column_schema, column)?; + + Ok(column_schema) +} + +/// Attaches extension metadata and settings to JSON2 columns. +fn set_json2_extension(column_schema: &mut ColumnSchema, column: &Column) -> Result<()> { + if !column_schema.data_type.is_json2() { + return Ok(()); + } + + let settings = column + .extensions + .build_json_settings()? + .unwrap_or_else(JsonSettings::new_v2); + if let Some(extension) = json2_extension(column.data_type(), settings) { column_schema.with_extension_type(&extension); } - Ok(column_schema) + Ok(()) } /// Convert `ColumnDef` in sqlparser to `ColumnDef` in gRPC proto. @@ -204,6 +208,22 @@ pub fn sql_column_def_to_grpc_column_def( SemanticType::Field }; + // TODO(fys): Extend the ALTER TABLE ADD COLUMN parser to support JSON2 + // type hints and pass the parsed JsonSettings through this conversion. + let options = json2_extension(&col.data_type, JsonSettings::new_v2()).map(|extension| { + let mut options = ColumnOptions::default(); + options.options.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + Json2ExtensionType::NAME.to_string(), + ); + if let Some(metadata) = extension.serialize_metadata() { + options + .options + .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata); + } + options + }); + Ok(api::v1::ColumnDef { name, data_type: datatype as i32, @@ -212,10 +232,26 @@ pub fn sql_column_def_to_grpc_column_def( semantic_type: semantic_type as _, comment: String::new(), datatype_extension: datatype_ext, - options: None, + options, }) } +fn json2_extension(data_type: &SqlDataType, settings: JsonSettings) -> Option { + let SqlDataType::Custom(name, _) = data_type else { + return None; + }; + if !name.0.first().is_some_and(|name| { + name.to_string_unquoted() + .eq_ignore_ascii_case(JSON2_TYPE_NAME) + }) { + return None; + } + + Some(Json2ExtensionType::new(Arc::new(JsonMetadata::new( + settings, + )))) +} + pub fn sql_data_type_to_concrete_data_type(data_type: &SqlDataType) -> Result { match data_type { SqlDataType::BigInt(_) | SqlDataType::Int64 => Ok(ConcreteDataType::int64_datatype()), diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index ce0c9748ce1..7114d9bafa1 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -513,3 +513,65 @@ drop table json2_finite_paths; Affected Rows: 0 +-- A newly added JSON2 column must not read data from a dropped column with the +-- same name in an older SST. Column identity is determined by column id, not +-- by the parquet root name. +create table json2_drop_add_same_name ( + ts timestamp time index, + b json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', +); + +Affected Rows: 0 + +insert into json2_drop_add_same_name +values (1, '{"a": "old"}'); + +Affected Rows: 1 + +admin flush_table('json2_drop_add_same_name'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_drop_add_same_name') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +alter table json2_drop_add_same_name drop column b; + +Affected Rows: 0 + +alter table json2_drop_add_same_name add column b json2; + +Affected Rows: 0 + +insert into json2_drop_add_same_name +values (2, '{"a": "new"}'); + +Affected Rows: 1 + +admin flush_table('json2_drop_add_same_name'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_drop_add_same_name') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +select ts, b.a +from json2_drop_add_same_name +order by ts; + ++-------------------------+--------------------------------------------------+ +| ts | json_get(json2_drop_add_same_name.b,Utf8("$.a")) | ++-------------------------+--------------------------------------------------+ +| 1970-01-01T00:00:00.001 | | +| 1970-01-01T00:00:00.002 | new | ++-------------------------+--------------------------------------------------+ + +drop table json2_drop_add_same_name; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2.sql b/tests/cases/standalone/common/types/json/json2.sql index 17248741c35..55642b7530c 100644 --- a/tests/cases/standalone/common/types/json/json2.sql +++ b/tests/cases/standalone/common/types/json/json2.sql @@ -157,3 +157,33 @@ from json2_finite_paths order by ts; drop table json2_finite_paths; + +-- A newly added JSON2 column must not read data from a dropped column with the +-- same name in an older SST. Column identity is determined by column id, not +-- by the parquet root name. +create table json2_drop_add_same_name ( + ts timestamp time index, + b json2 +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', +); + +insert into json2_drop_add_same_name +values (1, '{"a": "old"}'); + +admin flush_table('json2_drop_add_same_name'); + +alter table json2_drop_add_same_name drop column b; +alter table json2_drop_add_same_name add column b json2; + +insert into json2_drop_add_same_name +values (2, '{"a": "new"}'); + +admin flush_table('json2_drop_add_same_name'); + +select ts, b.a +from json2_drop_add_same_name +order by ts; + +drop table json2_drop_add_same_name; diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 555e76f1e5e..32bbfbaf0a7 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -95,6 +95,27 @@ drop table json2_alter_non_append; Affected Rows: 0 +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +-- ALTER TABLE ADD COLUMN does not yet support JSON2 settings or type hints. +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +Error: 2000(InvalidSyntax), Invalid SQL syntax: sql parser error: Expected: type modifiers, found: = at Line: 1, Column: 81 + +alter table json2_alter_add_settings add column k json2(service string); + +Error: 2000(InvalidSyntax), SQL data type not supported yet: Custom(ObjectName([Identifier(Ident { value: "json2", quote_style: None, span: Span(Location(1,51)..Location(1,56)) })]), ["service", "string"]) + +drop table json2_alter_add_settings; + +Affected Rows: 0 + create table json2_set_append_mode_false ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 653fdcec12c..1bbadfd5adc 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -58,6 +58,19 @@ alter table json2_alter_non_append add column j json2; drop table json2_alter_non_append; +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +-- ALTER TABLE ADD COLUMN does not yet support JSON2 settings or type hints. +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +alter table json2_alter_add_settings add column k json2(service string); + +drop table json2_alter_add_settings; + create table json2_set_append_mode_false ( ts timestamp time index, j json2 From 5affb50bb5d2fdf90963f08dfd54a80b423f5fd6 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:09:08 +0000 Subject: [PATCH 104/133] ci: update dev-builder image tag (#9469) Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 9144432b66d..7f92098262f 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 1.96.1-822042c7-20260924091315 +DEV_BUILDER_IMAGE_TAG ?= 1.98.1-7f0f592b-20261008091949 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu From 65b9a356921fc5157fc74ec7119b4c0988a128c7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 8 Oct 2026 10:29:26 +0000 Subject: [PATCH 105/133] fix(metasrv): retry failed heartbeat address updates (#9437) * fix(metasrv): retry failed heartbeat address updates Signed-off-by: WenyXu * refactor(metasrv): move shared heartbeat helpers to utils Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/AGENTS.md | 6 +- src/meta-srv/src/handler.rs | 1 + .../handler/collect_cluster_info_handler.rs | 247 +++++++++------ .../src/handler/collect_stats_handler.rs | 209 +++++++++---- .../src/handler/remap_flow_peer_handler.rs | 136 +++++---- src/meta-srv/src/handler/test_utils.rs | 37 ++- src/meta-srv/src/handler/utils.rs | 281 ++++++++++++++++++ 7 files changed, 714 insertions(+), 203 deletions(-) create mode 100644 src/meta-srv/src/handler/utils.rs diff --git a/src/meta-srv/AGENTS.md b/src/meta-srv/AGENTS.md index e1e23581525..5f9f573978c 100644 --- a/src/meta-srv/AGENTS.md +++ b/src/meta-srv/AGENTS.md @@ -38,7 +38,11 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`; - **Heartbeat**: `service/heartbeat.rs` receives datanode/frontend heartbeats and runs the handler chain in `handler/`; responses carry mailbox messages (DDL - results, cache invalidation). + results, cache invalidation). Node information is refreshed in memory on each + heartbeat. Datanode statistics collection and flownode remapping share + `handler/utils.rs` for per-node address persistence. Datanode address + updates retain the statistics handler's per-node ordering. Address epochs are + confirmed only after a successful save; cache invalidation remains best effort. - **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to `common-meta`'s `DdlManager`, executed via `common-procedure`. - **Region migration**: `procedure/region_migration/manager.rs` plus per-step diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 320f26fae7c..6bf71674a82 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -81,6 +81,7 @@ pub mod publish_heartbeat_handler; pub mod region_lease_handler; pub mod remap_flow_peer_handler; pub mod response_header_handler; +mod utils; #[cfg(test)] pub mod test_utils; diff --git a/src/meta-srv/src/handler/collect_cluster_info_handler.rs b/src/meta-srv/src/handler/collect_cluster_info_handler.rs index 6bee86ad711..12fee395136 100644 --- a/src/meta-srv/src/handler/collect_cluster_info_handler.rs +++ b/src/meta-srv/src/handler/collect_cluster_info_handler.rs @@ -46,30 +46,21 @@ impl HeartbeatHandler for CollectFrontendClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let frontend_workloads = get_frontend_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Frontend(FrontendStatus { - workloads: frontend_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Frontend(FrontendStatus { + workloads: frontend_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -89,29 +80,20 @@ impl HeartbeatHandler for CollectFlownodeClusterInfoHandler { ctx: &mut Context, _acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; let flownode_workloads = get_flownode_workloads(req.node_workloads.as_ref()); - let value = NodeInfo { - peer, - last_activity_ts: common_time::util::current_time_millis(), - status: NodeStatus::Flownode(FlownodeStatus { - workloads: flownode_workloads, - }), - version: info.version, - git_commit: info.git_commit, - start_time_ms: info.start_time_ms, - total_cpu_millicores: info.total_cpu_millicores, - total_memory_bytes: info.total_memory_bytes, - cpu_usage_millicores: info.cpu_usage_millicores, - memory_usage_bytes: info.memory_usage_bytes, - hostname: info.hostname, - env_vars, - }; - - put_into_memory_store(ctx, key, value).await?; + update + .save_to_memory( + ctx, + common_time::util::current_time_millis(), + NodeStatus::Flownode(FlownodeStatus { + workloads: flownode_workloads, + }), + ) + .await?; Ok(HandleControl::Continue) } @@ -132,7 +114,7 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { ctx: &mut Context, acc: &mut HeartbeatAccumulator, ) -> Result { - let Some((key, peer, info, env_vars)) = extract_base_info(req) else { + let Some(update) = NodeInfoUpdate::from_request(req) else { return Ok(HandleControl::Continue); }; @@ -147,16 +129,64 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { .count(); let follower_regions = stat.region_stats.len() - leader_regions; - let value = NodeInfo { + update + .save_to_memory( + ctx, + stat.timestamp_millis, + NodeStatus::Datanode(DatanodeStatus { + rcus: stat.rcus, + wcus: stat.wcus, + leader_regions, + follower_regions, + workloads: stat.datanode_workloads.clone(), + }), + ) + .await?; + + Ok(HandleControl::Continue) + } +} + +/// Common fields collected for one heartbeat's in-memory node information +/// update. Persistent routing addresses have a separate lifecycle. +struct NodeInfoUpdate { + key: NodeInfoKey, + peer: Peer, + info: PbNodeInfo, + env_vars: HashMap, +} + +impl NodeInfoUpdate { + fn from_request(request: &HeartbeatRequest) -> Option { + let key = NodeInfoKey::new(request)?; + let peer = request.peer.clone()?; + let info = request.info.clone()?; + let env_vars = EnvVars::from_extensions(&request.extensions) + .inspect_err(|e| { + warn!(e; "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer); + }) + .unwrap_or_default() + .map(|e| e.vars) + .unwrap_or_default(); + Some(Self { + key, peer, - last_activity_ts: stat.timestamp_millis, - status: NodeStatus::Datanode(DatanodeStatus { - rcus: stat.rcus, - wcus: stat.wcus, - leader_regions, - follower_regions, - workloads: stat.datanode_workloads.clone(), - }), + info, + env_vars, + }) + } + + async fn save_to_memory( + self, + ctx: &Context, + last_activity_ts: i64, + status: NodeStatus, + ) -> Result<()> { + let info = self.info; + let value = NodeInfo { + peer: self.peer, + last_activity_ts, + status, version: info.version, git_commit: info.git_commit, start_time_ms: info.start_time_ms, @@ -165,53 +195,88 @@ impl HeartbeatHandler for CollectDatanodeClusterInfoHandler { cpu_usage_millicores: info.cpu_usage_millicores, memory_usage_bytes: info.memory_usage_bytes, hostname: info.hostname, - env_vars, + env_vars: self.env_vars, }; - - put_into_memory_store(ctx, key, value).await?; - - Ok(HandleControl::Continue) + let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; + ctx.in_memory + .put(PutRequest { + key: (&self.key).into(), + value, + ..Default::default() + }) + .await + .context(SaveClusterInfoSnafu)?; + Ok(()) } } -fn extract_base_info( - request: &HeartbeatRequest, -) -> Option<(NodeInfoKey, Peer, PbNodeInfo, HashMap)> { - let HeartbeatRequest { peer, info, .. } = request; - let key = NodeInfoKey::new(request)?; - let Some(peer) = &peer else { - return None; - }; - let Some(info) = &info else { - return None; - }; +#[cfg(test)] +mod tests { + use api::v1::meta::RequestHeader; + use common_meta::datanode::Stat; - let env_vars = EnvVars::from_extensions(&request.extensions) - .inspect_err(|e| { - warn!(e; - "Failed to deserialize __env_vars from heartbeat extensions, peer: {}", peer + use super::*; + use crate::handler::test_utils::TestEnv; + + #[tokio::test] + async fn test_node_info_refreshes_without_epoch_change() { + let mut ctx = TestEnv::new().ctx(); + let handlers: Vec<(Role, Box)> = vec![ + (Role::Frontend, Box::new(CollectFrontendClusterInfoHandler)), + (Role::Datanode, Box::new(CollectDatanodeClusterInfoHandler)), + (Role::Flownode, Box::new(CollectFlownodeClusterInfoHandler)), + ]; + for (role, handler) in handlers { + let mut req = HeartbeatRequest { + header: Some(RequestHeader { + role: role as i32, + ..Default::default() + }), + peer: Some(Peer { + id: 7, + addr: "node:3001".into(), + }), + info: Some(PbNodeInfo { + version: "v1".into(), + hostname: "host".into(), + ..Default::default() + }), + node_epoch: 10, + ..Default::default() + }; + let mut acc = HeartbeatAccumulator { + stat: Some(Stat { + timestamp_millis: 123, + rcus: 5, + wcus: 6, + ..Default::default() + }), + ..Default::default() + }; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let key: Vec = (&NodeInfoKey::new(&req).unwrap()).into(); + let first = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(first.peer, req.peer.clone().unwrap()); + assert_eq!(first.version, "v1"); + assert_eq!(first.hostname, "host"); + req.info.as_mut().unwrap().version = "v2".into(); + req.info.as_mut().unwrap().memory_usage_bytes = 42; + acc.stat.as_mut().unwrap().timestamp_millis = 456; + handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); + let second = + NodeInfo::try_from(ctx.in_memory.get(&key).await.unwrap().unwrap().value).unwrap(); + assert_eq!(second.version, "v2"); + assert_eq!(second.memory_usage_bytes, 42); + assert_eq!( + serde_json::to_value(&second.status).unwrap(), + serde_json::to_value(&first.status).unwrap() ); - }) - .unwrap_or_default() - .map(|e| e.vars) - .unwrap_or_default(); - - Some((key, peer.clone(), info.clone(), env_vars)) -} - -async fn put_into_memory_store(ctx: &mut Context, key: NodeInfoKey, value: NodeInfo) -> Result<()> { - let key = (&key).into(); - let value = value.try_into().context(InvalidClusterInfoFormatSnafu)?; - let put_req = PutRequest { - key, - value, - ..Default::default() - }; - - ctx.in_memory - .put(put_req) - .await - .context(SaveClusterInfoSnafu)?; - - Ok(()) + if role == Role::Datanode { + assert_eq!(second.last_activity_ts, 456); + } + // Node information must not be persisted to the configured backend. + assert!(ctx.kv_backend.get(&key).await.unwrap().is_none()); + } + } } diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index cc653f9982c..4b58beaf0fb 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -15,19 +15,19 @@ use std::cmp::Ordering; use std::sync::Arc; -use api::v1::meta::{HeartbeatRequest, Role}; +use api::v1::meta::{HeartbeatRequest, Peer, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::peer::Peer; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; +use common_telemetry::{error, warn}; use dashmap::DashMap; use snafu::ResultExt; use tokio::sync::Mutex; use crate::error::{self, Result}; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; @@ -74,6 +74,7 @@ const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { stats_cache: DashMap>>, flush_stats_factor: usize, + address_updater: NodeAddressUpdater, } impl Default for CollectStatsHandler { @@ -87,6 +88,7 @@ impl CollectStatsHandler { Self { flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR), stats_cache: DashMap::default(), + address_updater: NodeAddressUpdater::default(), } } } @@ -119,6 +121,7 @@ impl HeartbeatHandler for CollectStatsHandler { let key: Vec = key.into(); + let current_epoch = current_stat.node_epoch; let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { Ordering::Greater => { @@ -146,11 +149,13 @@ impl HeartbeatHandler for CollectStatsHandler { true }; - // Need to refresh the [datanode -> address] mapping - if refresh { - // Safety: `epoch_stats.stats` is not empty - let last = epoch_stats.stats.last().unwrap(); - rewrite_node_address(ctx, last).await; + // Keep address updates and statistics in the original per-node order. + // Address acknowledgement is independent of the statistics epoch so a + // failed save can be retried by the next heartbeat of the same epoch. + if epoch_stats.epoch() == Some(current_epoch) + && let Some(stat) = epoch_stats.stats.last() + { + update_datanode_address(&self.address_updater, ctx, stat).await; } if !refresh && epoch_stats.len() < self.flush_stats_factor { @@ -178,48 +183,26 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &Context, stat: &Stat) { - let peer = Peer { - id: stat.id, - addr: stat.addr.clone(), - }; - let key = NodeAddressKey::with_datanode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; - - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!( - "Successfully updated datanode `NodeAddressValue`: {:?}", - peer - ); - // broadcast invalidating cache - let cache_idents = stat - .table_ids() - .into_iter() - .map(CacheIdent::TableId) - .collect::>(); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update datanode `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!( - "Failed to serialize datanode `NodeAddressValue`: {:?}", - peer - ); +/// Applies datanode-specific address and table-cache updates. The caller keeps +/// its statistics lock so notifications and statistics retain their ordering. +async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) { + let updated = updater + .update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async { + let peer = Peer { id: stat.id, addr: stat.addr.clone() }; + let key = NodeAddressKey::with_datanode(stat.id).to_bytes(); + utils::save_node_address(ctx, key, peer).await + }) + .await + .inspect_err(|e| { + error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr); + }); + if let Ok(true) = updated { + let cache_idents = stat + .table_ids() + .into_iter() + .map(CacheIdent::TableId) + .collect::>(); + utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await; } } @@ -553,4 +536,126 @@ mod tests { handler.handle(&req, &mut ctx, &mut acc).await.unwrap(); } } + + #[tokio::test] + async fn test_persistence_failure_retries_but_notification_failure_does_not() { + use crate::handler::collect_stats_handler::CollectStatsHandler; + use crate::handler::test_utils::FailingCacheInvalidator; + use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + use api::v1::meta::HeartbeatRequest; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; + use common_meta::rpc::store::PutResponse; + let mut ctx = TestEnv::new().ctx(); + let puts = Arc::new(std::sync::Mutex::new(Vec::::new())); + let backend = MockKvBackendBuilder::default() + .put_fn(Arc::new({ + let puts = puts.clone(); + move |req: PutRequest| { + let mut puts = puts.lock().unwrap(); + puts.push(req); + if puts.len() == 1 { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected persistence failure", + } + .fail(); + } + Ok(PutResponse::default()) + } + }) as _) + .build() + .unwrap(); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend))); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = CollectStatsHandler::default(); + let mut stat = Stat { + id: 101, + addr: "dn-original".into(), + node_epoch: 7, + ..Default::default() + }; + let req = HeartbeatRequest::default(); + + // A backend error must not interrupt subsequent heartbeat handlers. + assert!(matches!( + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + } + ) + .await + .unwrap(), + HandleControl::Continue + )); + assert_eq!(puts.lock().unwrap().len(), 1); + assert_eq!(invalidator.attempts(), 0); + + let stats_key: Vec = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into(); + assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some()); + + // Same epoch retries persistence using the current heartbeat's address. + stat.addr = "dn-retry".into(); + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + { + let puts = puts.lock().unwrap(); + assert_eq!(puts.len(), 2); + assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes()); + assert_eq!( + NodeAddressValue::try_from_raw_value(&puts[1].value) + .unwrap() + .peer + .addr, + "dn-retry" + ); + } + assert_eq!(invalidator.attempts(), 1); + + // Notification remains best-effort: its failure neither rewrites nor retries. + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 2); + assert_eq!(invalidator.attempts(), 1); + + stat.node_epoch = 8; + handler + .handle( + &req, + &mut ctx, + &mut HeartbeatAccumulator { + stat: Some(stat.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(puts.lock().unwrap().len(), 3); + assert_eq!(invalidator.attempts(), 2); + } } diff --git a/src/meta-srv/src/handler/remap_flow_peer_handler.rs b/src/meta-srv/src/handler/remap_flow_peer_handler.rs index f6b3c67a9e5..186cefa19e8 100644 --- a/src/meta-srv/src/handler/remap_flow_peer_handler.rs +++ b/src/meta-srv/src/handler/remap_flow_peer_handler.rs @@ -12,22 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -use api::v1::meta::{HeartbeatRequest, Peer, Role}; +use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::instruction::CacheIdent; -use common_meta::key::node_address::{NodeAddressKey, NodeAddressValue}; -use common_meta::key::{MetadataKey, MetadataValue}; -use common_meta::rpc::store::PutRequest; -use common_telemetry::{error, info, warn}; -use dashmap::DashMap; +use common_meta::key::MetadataKey; +use common_meta::key::node_address::NodeAddressKey; +use common_telemetry::error; use crate::Result; +use crate::handler::utils::{self, NodeAddressUpdater}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::Context; #[derive(Debug, Default)] pub struct RemapFlowPeerHandler { - /// flow_node_id -> epoch - epoch_cache: DashMap, + address_updater: NodeAddressUpdater, } #[async_trait::async_trait] @@ -45,61 +43,83 @@ impl HeartbeatHandler for RemapFlowPeerHandler { let Some(peer) = req.peer.as_ref() else { return Ok(HandleControl::Continue); }; - - let current_epoch = req.node_epoch; - let flow_node_id = peer.id; - - let refresh = if let Some(mut epoch) = self.epoch_cache.get_mut(&flow_node_id) { - if current_epoch > *epoch.value() { - *epoch.value_mut() = current_epoch; - true - } else { - false - } - } else { - self.epoch_cache.insert(flow_node_id, current_epoch); - true - }; - - if refresh { - rewrite_node_address(ctx, peer).await; + let updated = self + .address_updater + .update_if_needed((Role::Flownode, peer.id), req.node_epoch, || { + let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); + utils::save_node_address(ctx, key, peer.clone()) + }) + .await + .inspect_err(|e| { + // Address errors must not prevent later heartbeat handlers from running. + error!(e; "Failed to update flownode address, peer: {:?}", peer); + }); + if let Ok(true) = updated { + let cache_idents = [CacheIdent::FlowNodeAddressChange(peer.id)]; + utils::invalidate_address_caches(ctx, peer.id, &cache_idents).await; } - Ok(HandleControl::Continue) } } -async fn rewrite_node_address(ctx: &mut Context, peer: &Peer) { - let key = NodeAddressKey::with_flownode(peer.id).to_bytes(); - if let Ok(value) = NodeAddressValue::new(peer.clone()).try_as_raw_value() { - let put = PutRequest { - key, - value, - prev_kv: false, - }; +#[cfg(test)] +mod tests { + use super::*; + use crate::handler::test_utils::TestEnv; + use api::v1::meta::Peer; + use std::sync::Arc; + #[tokio::test] + async fn test_notification_failure_keeps_persisted_epoch() { + use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; + use api::v1::meta::HeartbeatRequest; - match ctx.leader_cached_kv_backend.put(put).await { - Ok(_) => { - info!("Successfully updated flow `NodeAddressValue`: {:?}", peer); - // broadcast invalidating cache to all frontends - let cache_idents = vec![CacheIdent::FlowNodeAddressChange(peer.id)]; - info!( - "Invalidate flow node cache for new address with cache idents: {:?}", - cache_idents - ); - if let Err(e) = ctx - .cache_invalidator - .invalidate(&Default::default(), &cache_idents) - .await - { - error!(e; "Failed to invalidate {} `NodeAddressKey` cache, peer: {:?}", cache_idents.len(), peer); - } - } - Err(e) => { - error!(e; "Failed to update flow `NodeAddressValue`: {:?}", peer); - } - } - } else { - warn!("Failed to serialize flow `NodeAddressValue`: {:?}", peer); + use crate::handler::test_utils::FailingCacheInvalidator; + use common_meta::key::MetadataValue; + use common_meta::key::node_address::NodeAddressValue; + let mut ctx = TestEnv::new().ctx(); + let invalidator = Arc::new(FailingCacheInvalidator::default()); + ctx.cache_invalidator = invalidator.clone(); + let handler = RemapFlowPeerHandler::default(); + let mut req = HeartbeatRequest { + peer: Some(Peer { + id: 1, + addr: "first".into(), + }), + node_epoch: 1, + ..Default::default() + }; + let key = NodeAddressKey::with_flownode(1).to_bytes(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + req.peer.as_mut().unwrap().addr = "same-epoch".into(); + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 1); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "first" + ); + req.node_epoch = 2; + handler + .handle(&req, &mut ctx, &mut HeartbeatAccumulator::default()) + .await + .unwrap(); + assert_eq!(invalidator.attempts(), 2); + let kv = ctx.kv_backend.get(&key).await.unwrap().unwrap(); + assert_eq!( + NodeAddressValue::try_from_raw_value(&kv.value) + .unwrap() + .peer + .addr, + "same-epoch" + ); } } diff --git a/src/meta-srv/src/handler/test_utils.rs b/src/meta-srv/src/handler/test_utils.rs index d8241c785ea..c2afb3cb14a 100644 --- a/src/meta-srv/src/handler/test_utils.rs +++ b/src/meta-srv/src/handler/test_utils.rs @@ -13,9 +13,13 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; -use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; +use common_meta::cache_invalidator::{ + CacheInvalidator, CacheInvalidatorRef, Context as InvalidationContext, DummyCacheInvalidator, +}; use common_meta::distributed_time_constants::BASE_HEARTBEAT_INTERVAL; +use common_meta::instruction::CacheIdent; use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; @@ -97,3 +101,34 @@ impl TestEnv { } } } + +/// Counts failed notification attempts without requiring heartbeat connections. +#[derive(Default)] +pub(crate) struct FailingCacheInvalidator { + attempts: AtomicUsize, +} + +impl FailingCacheInvalidator { + pub(crate) fn attempts(&self) -> usize { + self.attempts.load(Ordering::Relaxed) + } +} + +#[async_trait::async_trait] +impl CacheInvalidator for FailingCacheInvalidator { + async fn invalidate( + &self, + _: &InvalidationContext, + _: &[CacheIdent], + ) -> common_meta::error::Result<()> { + self.attempts.fetch_add(1, Ordering::Relaxed); + common_meta::error::UnexpectedSnafu { + err_msg: "injected notification failure", + } + .fail() + } + + fn invalidate_all(&self) -> common_meta::error::Result<()> { + Ok(()) + } +} diff --git a/src/meta-srv/src/handler/utils.rs b/src/meta-srv/src/handler/utils.rs new file mode 100644 index 00000000000..96624d87306 --- /dev/null +++ b/src/meta-srv/src/handler/utils.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::sync::Arc; + +use api::v1::meta::{Peer, Role}; +use common_meta::instruction::CacheIdent; +use common_meta::key::MetadataValue; +use common_meta::key::node_address::NodeAddressValue; +use common_meta::rpc::store::PutRequest; +use common_telemetry::{error, info}; +use dashmap::DashMap; +use snafu::ResultExt; +use tokio::sync::Mutex; + +use crate::Result; +use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu}; +use crate::metasrv::Context; + +/// Serializes address changes per node and acknowledges only successful updates. +/// Node information refreshed on every heartbeat has a separate lifecycle. +#[derive(Debug, Default)] +pub(crate) struct NodeAddressUpdater { + states: DashMap<(Role, u64), Arc>>, +} + +#[derive(Debug, Default)] +struct NodeAddressEpochs { + /// Rejects stale heartbeats even after a newer epoch failed to save. + seen_epoch: Option, + /// Suppresses retries only after persistence succeeds. + saved_epoch: Option, +} + +impl NodeAddressUpdater { + /// Saves a new or previously failed epoch. Returns true only when this call + /// saved the address, or false when the heartbeat is stale or already saved. + /// Cache invalidation is separate and does not affect acknowledgement. + pub(crate) async fn update_if_needed( + &self, + node: (Role, u64), + epoch: u64, + save: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // Never retain a DashMap guard across await, even while waiting for a + // different node: its entry may live in the same shard. + let state = { + let entry = self.states.entry(node).or_default(); + Arc::clone(entry.value()) + }; + let mut state = state.lock().await; + if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) { + return Ok(false); + } + + // Seeing an epoch rejects stale heartbeats, but only a successful save + // suppresses retries of that same epoch. Cancellation also leaves it retryable. + state.seen_epoch = Some(epoch); + save().await?; + state.saved_epoch = Some(epoch); + Ok(true) + } +} + +/// Persists an address through the leader cache. +pub(crate) async fn save_node_address(ctx: &Context, key: Vec, peer: Peer) -> Result<()> { + let address = NodeAddressValue::new(peer); + let value = address + .try_as_raw_value() + .context(InvalidNodeInfoFormatSnafu)?; + ctx.leader_cached_kv_backend + .put(PutRequest { + key, + value, + prev_kv: false, + }) + .await + .context(KvBackendSnafu)?; + info!("Successfully updated node address: {:?}", address.peer); + Ok(()) +} + +/// Preserves the existing best-effort notification behavior. Notification +/// failures do not undo a persisted address or cause same-epoch rewrites. +pub(crate) async fn invalidate_address_caches( + ctx: &Context, + node_id: u64, + cache_idents: &[CacheIdent], +) { + if let Err(e) = ctx + .cache_invalidator + .invalidate(&Default::default(), cache_idents) + .await + { + error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::sync::Notify; + use tokio::time::timeout; + + use super::*; + + #[tokio::test] + async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + assert!( + updater + .update_if_needed(key, 2, || async { + crate::error::UnexpectedSnafu { + violated: "injected save failure", + } + .fail() + }) + .await + .is_err() + ); + assert!( + !updater + .update_if_needed(key, 1, || async { + panic!("stale heartbeat must not save"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 2, || async { Ok(()) }) + .await + .unwrap() + ); + assert!( + !updater + .update_if_needed(key, 2, || async { + panic!("successful epoch must not save again"); + }) + .await + .unwrap() + ); + assert!( + updater + .update_if_needed(key, 3, || async { Ok(()) }) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn test_cancelled_save_can_be_retried() { + let updater = NodeAddressUpdater::default(); + let key = (Role::Datanode, 1); + let entered = Notify::new(); + { + let first = updater.update_if_needed(key, 0, || async { + entered.notify_one(); + std::future::pending::>().await + }); + tokio::pin!(first); + tokio::select! { + _ = entered.notified() => {}, + _ = &mut first => panic!("save must remain pending"), + } + } + assert!( + timeout( + Duration::from_secs(2), + updater.update_if_needed(key, 0, || async { Ok(()) }) + ) + .await + .unwrap() + .unwrap() + ); + } + + #[test] + fn test_same_node_serialized_other_nodes_progress() { + // A watchdog outside the runtime detects synchronous DashMap deadlocks; + // a Tokio timeout alone cannot fire if the only worker is blocked. + let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1); + let release = Arc::new(Notify::new()); + let finish = Arc::new(Notify::new()); + let worker = std::thread::spawn({ + let release = release.clone(); + let finish = finish.clone(); + move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let updater = Arc::new(NodeAddressUpdater::default()); + let entered = Arc::new(Notify::new()); + + let first = tokio::spawn({ + let updater = updater.clone(); + let entered = entered.clone(); + let release = release.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 1, || async { + entered.notify_one(); + release.notified().await; + Ok(()) + }) + .await + } + }); + entered.notified().await; + let second = tokio::spawn({ + let updater = updater.clone(); + async move { + updater + .update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) }) + .await + } + }); + tokio::task::yield_now().await; + assert!(!second.is_finished()); + // Exercise shared DashMap shards on a single-thread runtime while the + // first node is suspended in external I/O. + timeout(Duration::from_secs(2), async { + for id in 0_u64..128 { + assert!( + updater + .update_if_needed((Role::Datanode, id + 2), 1, || async { + Ok(()) + }) + .await + .unwrap() + ); + } + }) + .await + .unwrap(); + progress_tx.send(()).unwrap(); + finish.notified().await; + timeout(Duration::from_secs(2), async { + assert!(first.await.unwrap().unwrap()); + assert!(second.await.unwrap().unwrap()); + }) + .await + .unwrap(); + assert!( + !updater + .update_if_needed((Role::Datanode, 1), 1, || async { + panic!("old epoch must not overwrite a newer one"); + }) + .await + .unwrap() + ); + }); + } + }); + let progress = progress_rx.recv_timeout(Duration::from_secs(5)); + release.notify_one(); + finish.notify_one(); + progress.expect("other node updates must progress while one node save is blocked"); + worker.join().unwrap(); + } +} From 5126623e6356bd6849aeb69aa3f2ec3386effab5 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:34:20 +0000 Subject: [PATCH 106/133] fix: align TWCS assignment with half-open time windows (#9464) Signed-off-by: Lei, HUANG --- src/mito2/src/compaction/twcs.rs | 304 ++++++++++++++++++++++--------- 1 file changed, 216 insertions(+), 88 deletions(-) diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index dd407b30382..c170845be3d 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -969,6 +969,7 @@ struct Window { start: Timestamp, end: Timestamp, files: Vec, + /// Start of the half-open window, in seconds. time_window: i64, primary_key_range: Option<(bytes::Bytes, bytes::Bytes)>, } @@ -1022,12 +1023,7 @@ fn assign_to_windows<'a>( continue; } let (_, end) = f.time_range(); - let time_window = end - .convert_to(TimeUnit::Second) - .unwrap() - .value() - .align_to_ceil_by_bucket(time_window_size) - .unwrap_or(i64::MIN); + let time_window = time_window_start(end, time_window_size); match windows.entry(time_window) { Entry::Occupied(mut e) => { @@ -1043,37 +1039,41 @@ fn assign_to_windows<'a>( windows.into_iter().collect() } +/// Uses the same half-open windows as memtables, including timestamps in the first +/// second after a boundary. An unrepresentable window start is clamped to `i64::MIN`. +fn time_window_start(timestamp: Timestamp, time_window_size: i64) -> i64 { + // Conversion to seconds cannot overflow. + timestamp + .convert_to(TimeUnit::Second) + .and_then(|ts| ts.value().align_by_bucket(time_window_size)) + .unwrap_or(i64::MIN) +} + fn time_window_intersects_range( - window_end: i64, + window_start: i64, time_window_size: i64, time_range: &TimestampRange, ) -> bool { - let first_window = match time_range.start() { - None => i64::MIN, - Some(start) => { - let Some(first_window) = start - .convert_to(TimeUnit::Second) - .and_then(|timestamp| timestamp.value().align_to_ceil_by_bucket(time_window_size)) - else { - return false; - }; - first_window - } - }; + if time_range.is_empty() { + return false; + } + let first_window = time_range + .start() + .map_or(i64::MIN, |start| time_window_start(start, time_window_size)); let last_window = match time_range.end() { None => i64::MAX, Some(end) => { - let Some(last_window) = end + // The range end is exclusive; keep the last second that intersects it. + let Some(last_second) = end .convert_to_ceil(TimeUnit::Second) .and_then(|timestamp| timestamp.value().checked_sub(1)) - .and_then(|timestamp| timestamp.align_to_ceil_by_bucket(time_window_size)) else { return false; }; - last_window + time_window_start(Timestamp::new_second(last_second), time_window_size) } }; - (first_window..=last_window).contains(&window_end) + (first_window..=last_window).contains(&window_start) } fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { @@ -1091,7 +1091,7 @@ fn window_has_overlap(this: &Window, windows: &BTreeMap) -> bool { } /// Finds the latest active writing window among all files. -/// Returns `None` when there are no files or all files are corrupted. +/// Returns `None` when there are no files. fn find_latest_window_in_seconds<'a>( files: impl Iterator, time_window_size: i64, @@ -1107,9 +1107,7 @@ fn find_latest_window_in_seconds<'a>( latest_timestamp = Some(end); } } - latest_timestamp - .and_then(|ts| ts.convert_to_ceil(TimeUnit::Second)) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) + latest_timestamp.map(|ts| time_window_start(ts, time_window_size)) } /// Finds the active window from the file with the highest sequence number. @@ -1119,9 +1117,7 @@ fn find_latest_window_in_seconds<'a>( /// survives the transient state where level 0 is empty right after its files /// were compacted away. Returns `None` when no file carries a sequence. /// -/// The window key follows the same convention as [`assign_to_windows`] -/// (truncate to seconds, then align up), since it is compared against the -/// window keys produced there. +/// The window key is its inclusive start, as in [`assign_to_windows`]. fn find_active_window_by_sequence<'a>( files: impl Iterator, time_window_size: i64, @@ -1129,12 +1125,7 @@ fn find_active_window_by_sequence<'a>( files .filter(|f| f.meta_ref().sequence.is_some()) .max_by_key(|f| f.meta_ref().sequence) - .and_then(|f| { - f.time_range() - .1 - .convert_to(TimeUnit::Second) - .and_then(|ts| ts.value().align_to_ceil_by_bucket(time_window_size)) - }) + .map(|f| time_window_start(f.time_range().1, time_window_size)) } #[cfg(test)] @@ -1533,7 +1524,7 @@ mod tests { new_file_handle_with_sequence(FileId::random(), 0, 999, 1, 0), new_file_handle_with_sequence(FileId::random(), 2000, 2999, 0, 0), ]; - assert_eq!(Some(3), active_window_of(&files)); + assert_eq!(Some(2), active_window_of(&files)); // No sequence and no L0: None, exactly as before. let files = [new_file_handle_with_sequence( @@ -1549,7 +1540,7 @@ mod tests { #[test] fn test_get_latest_window_in_seconds() { assert_eq!( - Some(1), + Some(0), find_latest_window_in_seconds([new_file_handle(FileId::random(), 0, 999, 0)].iter(), 1) ); assert_eq!( @@ -1561,7 +1552,7 @@ mod tests { ); assert_eq!( - Some(-9223372036854000), + Some(-9223372036857600), find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MIN + 1, 0)].iter(), 3600, @@ -1569,7 +1560,7 @@ mod tests { ); assert_eq!( - (i64::MAX / 10000000 + 1) * 10000, + (i64::MAX / 10000000) * 10000, find_latest_window_in_seconds( [new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0)].iter(), 10000, @@ -1578,7 +1569,7 @@ mod tests { ); assert_eq!( - Some((i64::MAX / 3600000 + 1) * 3600), + Some((i64::MAX / 3600000) * 3600), find_latest_window_in_seconds( [ new_file_handle(FileId::random(), i64::MIN, i64::MAX, 0), @@ -1590,6 +1581,114 @@ mod tests { ); } + #[test] + fn test_window_assignment_and_activity_at_boundaries() { + // Before, at, and after positive/negative boundaries, including the first + // fractional second: all picker paths must agree with half-open partitions. + for (unit, ticks_per_second) in [ + (TimeUnit::Second, 1), + (TimeUnit::Millisecond, 1_000), + (TimeUnit::Microsecond, 1_000_000), + (TimeUnit::Nanosecond, 1_000_000_000), + ] { + for window_size in [1, 3, 86_400] { + for boundary in [-window_size, 0, window_size] { + for (offset, expected_window) in [ + (-1, boundary - window_size), + (0, boundary), + (ticks_per_second - 1, boundary), + (ticks_per_second, boundary + i64::from(window_size == 1)), + ] { + let timestamp = Timestamp::new(boundary * ticks_per_second + offset, unit); + let file = FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + ); + let files = [file]; + let windows = assign_to_windows(files.iter(), window_size); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>(), + "{timestamp:?}, window size {window_size}" + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), window_size) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), window_size) + ); + } + } + } + } + } + + #[test] + fn test_window_assignment_at_timestamp_limits() { + for (seconds, expected_window) in [(i64::MIN, i64::MIN), (i64::MAX, i64::MAX - 1)] { + let timestamp = Timestamp::new_second(seconds); + let files = [FileHandle::new( + FileMeta { + time_range: (timestamp, timestamp), + sequence: NonZeroU64::new(1), + ..Default::default() + }, + crate::test_util::new_noop_file_purger(), + )]; + let windows = assign_to_windows(files.iter(), 3); + assert_eq!( + vec![expected_window], + windows.keys().copied().collect::>() + ); + assert_eq!( + Some(expected_window), + find_active_window_by_sequence(files.iter(), 3) + ); + assert_eq!( + Some(expected_window), + find_latest_window_in_seconds(files.iter(), 3) + ); + assert!(time_window_intersects_range( + expected_window, + 3, + &TimestampRange::single(timestamp) + )); + } + } + + #[tokio::test] + async fn test_midnight_file_is_not_compacted_with_previous_day() { + // Regression for #9441: flush already separates the midnight-only SST. + let midnight = 86_400_000; + CompactionPickerTestCase { + window_size: 86_400, + input_files: [ + (midnight - 240_000, midnight - 180_000), + (midnight - 180_000, midnight - 90_000), + (midnight - 90_000, midnight - 30_000), + (midnight, midnight), + ] + .into_iter() + .enumerate() + .map(|(index, (start, end))| { + new_file_handle_with_sequence(FileId::random(), start, end, 0, index as u64 + 1) + }) + .collect(), + expected_outputs: vec![ExpectedOutput { + input_files: vec![0, 1, 2], + output_level: 1, + }], + } + .check() + .await; + } + #[test] fn test_assign_to_windows() { let windows = assign_to_windows( @@ -1619,7 +1718,7 @@ mod tests { assert_eq!( files[0], windows - .get(&0) + .get(&-3) .unwrap() .files() .next() @@ -1630,7 +1729,7 @@ mod tests { assert_eq!( files[1], windows - .get(&3) + .get(&0) .unwrap() .files() .next() @@ -1641,7 +1740,7 @@ mod tests { assert_eq!( files[2], windows - .get(&12) + .get(&9) .unwrap() .files() .next() @@ -1732,6 +1831,7 @@ mod tests { .collect::>(); let windows = assign_to_windows(file_handles.iter(), time_window); + assert_eq!(expected_files.len(), windows.len()); for (expected_window, overlapping, window_files) in expected_files { let actual_window = windows.get(expected_window).unwrap(); @@ -1756,8 +1856,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999)]), ], ); @@ -1774,9 +1874,8 @@ mod tests { &[(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)], 2, &[ - (0, false, vec![(0, 999)]), - (2, false, vec![(1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (3000, 3999)]), ], ); @@ -1790,9 +1889,8 @@ mod tests { ], 2, &[ - (0, true, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(0, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 3999), (2000, 2999), (3000, 3999)]), ], ); @@ -1806,43 +1904,39 @@ mod tests { ], 2, &[ - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(1999, 3999), (3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(1999, 3999), (2000, 2999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (2999, 3999), // window 4 + (3000, 3999), // window 2 + (2999, 3999), // window 2 ], 2, &[ - // window 2 overlaps with window 4 - (0, false, vec![(0, 999)]), - (2, true, vec![(1000, 1999), (2000, 2999)]), - (4, true, vec![(2999, 3999), (3000, 3999)]), + // Overlap within a window does not imply cross-window overlap. + (0, false, vec![(0, 999), (1000, 1999)]), + (2, false, vec![(2000, 2999), (2999, 3999), (3000, 3999)]), ], ); check_assign_to_windows_with_overlapping( &[ (0, 999), // window 0 - (1000, 1999), // window 2 + (1000, 1999), // window 0 (2000, 2999), // window 2 - (3000, 3999), // window 4 - (0, 1000), // // window 2 + (3000, 3999), // window 2 + (0, 2000), // window 2 ], 2, &[ - // only window 0 overlaps with window 2. - (0, true, vec![(0, 999)]), - (2, true, vec![(0, 1000), (1000, 1999), (2000, 2999)]), - (4, false, vec![(3000, 3999)]), + (0, true, vec![(0, 999), (1000, 1999)]), + (2, true, vec![(0, 2000), (2000, 2999), (3000, 3999)]), ], ); } @@ -1862,7 +1956,7 @@ mod tests { new_file_handle_with_size_sequence_and_primary_key_range( FileId::random(), 500, - 1999, + 2999, 0, 2, 10, @@ -1872,7 +1966,8 @@ mod tests { let windows = assign_to_windows(files.iter(), 2); - let overlapping = window_has_overlap(windows.get(&2).unwrap(), &windows); + assert_eq!(2, windows.len()); + let overlapping = window_has_overlap(windows.get(&0).unwrap(), &windows); assert!(!overlapping); } @@ -2592,7 +2687,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), None) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), None) .await .unwrap(); @@ -2609,27 +2704,60 @@ mod tests { #[test] fn test_filter_time_windows_by_time_range() { - let time_range = TimestampRange::new( - Timestamp::new_millisecond(1_200), - Timestamp::new_millisecond(1_800), - ) - .unwrap(); + // Interior, exclusive boundary ends, fractional boundary ends, negative + // timestamps, and empty ranges must all use the same half-open windows. + for (start, end, expected_windows) in [ + (1_200, 1_800, vec![0]), + (0, 3_000, vec![0]), + (3_000, 6_000, vec![3]), + (2_999, 3_000, vec![0]), + (2_999, 3_001, vec![0, 3]), + (3_000, 3_001, vec![3]), + (3_001, 3_002, vec![3]), + (-3_000, 0, vec![-3]), + (-3_001, -2_999, vec![-6, -3]), + (0, 0, vec![]), + (1_500, 1_500, vec![]), + ] { + let time_range = TimestampRange::with_unit(start, end, TimeUnit::Millisecond).unwrap(); + let selected = [-6, -3, 0, 3, 6] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &time_range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{time_range}"); + } + } - assert!(time_window_intersects_range(3, 3, &time_range)); - assert!(!time_window_intersects_range(9, 3, &time_range)); + #[test] + fn test_filter_time_windows_with_unbounded_and_extreme_ranges() { + for (range, expected_windows) in [ + (TimestampRange::min_to_max(), vec![-3, 0, 3]), + ( + TimestampRange::from_start(Timestamp::new_second(0)), + vec![0, 3], + ), + ( + TimestampRange::until_end(Timestamp::new_second(0), false), + vec![-3], + ), + ] { + let selected = [-3, 0, 3] + .into_iter() + .filter(|window| time_window_intersects_range(*window, 3, &range)) + .collect::>(); + assert_eq!(expected_windows, selected, "{range}"); + } - let boundary_range = - TimestampRange::new(Timestamp::new_second(0), Timestamp::new_second(3)).unwrap(); - assert!(time_window_intersects_range(0, 3, &boundary_range)); - assert!(time_window_intersects_range(3, 3, &boundary_range)); - assert!(!time_window_intersects_range(6, 3, &boundary_range)); - - let overflowing_range = TimestampRange::new( + let upper_range = TimestampRange::new( Timestamp::new_second(i64::MAX - 1), Timestamp::new_second(i64::MAX), ) .unwrap(); - assert!(!time_window_intersects_range(0, 4, &overflowing_range)); + assert!(!time_window_intersects_range(0, 4, &upper_range)); + assert!(time_window_intersects_range(i64::MAX - 3, 4, &upper_range)); + + let lower_range = TimestampRange::until_end(Timestamp::new_second(i64::MIN), false); + assert!(!time_window_intersects_range(i64::MIN, 3, &lower_range)); } #[tokio::test] @@ -2659,7 +2787,7 @@ mod tests { }; let output = picker - .build_output_with_time_range(RegionId::from_u64(123), windows, Some(9), Some(3)) + .build_output_with_time_range(RegionId::from_u64(123), windows, Some(6), Some(3)) .await .unwrap(); From 6265c31e4afa4e064fd0f266f6756cae4853a1db Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 8 Oct 2026 13:14:06 +0000 Subject: [PATCH 107/133] feat: batch logical table DDL during import-v2 (#9467) * feat: batch logical table DDL during import-v2 Signed-off-by: jeremyhi * fix: preserve one-shot batch DDL transport and context Signed-off-by: jeremyhi * test: check schema-only import state at its own path Signed-off-by: jeremyhi * fix: preserve oversized import DDL compatibility Signed-off-by: jeremyhi * refactor: reuse import DDL encoding and metric constants Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/cli/src/data/import_v2/command.rs | 166 +++---- src/cli/src/data/import_v2/executor.rs | 406 ++++++++++++++++-- src/cli/src/data/import_v2/packed.rs | 85 ++-- src/cli/src/database.rs | 308 ++++++++++++- .../tests/export_logical_tables.rs | 129 +++++- 5 files changed, 904 insertions(+), 190 deletions(-) diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 33cf9cb1ffb..62b9de08816 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -222,12 +222,19 @@ impl Import { manifest.validate_layout().map_err(|reason| { crate::data::import_v2::error::InvalidPackedSnapshotSnafu { reason }.build() })?; - if manifest.is_packed() { - self.database_client - .require_packed_import() + let capabilities = if manifest.is_packed() { + let capabilities = self + .database_client + .capabilities() .await .context(crate::data::import_v2::error::DatabaseSnafu)?; - } + capabilities + .require("metric_packed_import") + .context(crate::data::import_v2::error::DatabaseSnafu)?; + Some(capabilities) + } else { + None + }; info!("Snapshot contains {} schema(s)", manifest.schemas.len()); @@ -279,6 +286,15 @@ impl Import { return Ok(()); } + let capabilities = match capabilities { + Some(capabilities) => capabilities, + None => self + .database_client + .capabilities() + .await + .context(crate::data::import_v2::error::DatabaseSnafu)?, + }; + let mut resume_session = if !data_tasks.is_empty() { let state_path = resolve_state_path( self.state_path.as_deref(), @@ -316,7 +332,9 @@ impl Import { false } else { let executor = DdlExecutor::new(&self.database_client); - executor.execute_strict(&ddl_statements).await?; + executor + .execute_strict(&ddl_statements, capabilities.supports("metric_batch_ddl")) + .await?; if let Some(session) = resume_session.as_mut() { session.mark_ddl_completed().await?; } @@ -356,9 +374,7 @@ impl Import { .await .context(SnapshotStorageSnafu)?; statements.extend( - parse_ddl_statements(&content) - .into_iter() - .map(|sql| ddl_statement_for_schema(schema, sql)), + iter_ddl_statements(&content).map(|sql| ddl_statement_for_schema(schema, sql)), ); } @@ -402,90 +418,92 @@ impl ImportTaskExecutor for CopyDatabaseImportTaskExecutor<'_> { } } +#[cfg(test)] fn parse_ddl_statements(content: &str) -> Vec { - let mut statements = Vec::new(); - let mut current = String::new(); + iter_ddl_statements(content).collect() +} + +pub(crate) fn iter_ddl_statements(content: &str) -> impl Iterator + '_ { let mut chars = content.chars().peekable(); let mut in_single_quote = false; let mut in_double_quote = false; let mut in_line_comment = false; let mut in_block_comment = false; - while let Some(ch) = chars.next() { - if in_line_comment { - if ch == '\n' { - in_line_comment = false; - current.push('\n'); - } - continue; - } - - if in_block_comment { - if ch == '*' && chars.peek() == Some(&'/') { - chars.next(); - in_block_comment = false; - } - continue; - } - - if in_single_quote { - current.push(ch); - if ch == '\'' { - if chars.peek() == Some(&'\'') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_single_quote = false; + std::iter::from_fn(move || { + let mut current = String::new(); + while let Some(ch) = chars.next() { + if in_line_comment { + if ch == '\n' { + in_line_comment = false; + current.push('\n'); } + continue; } - continue; - } - if in_double_quote { - current.push(ch); - if ch == '"' { - if chars.peek() == Some(&'"') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_double_quote = false; + if in_block_comment { + if ch == '*' && chars.peek() == Some(&'/') { + chars.next(); + in_block_comment = false; } + continue; } - continue; - } - match ch { - '-' if chars.peek() == Some(&'-') => { - chars.next(); - in_line_comment = true; - } - '/' if chars.peek() == Some(&'*') => { - chars.next(); - in_block_comment = true; - } - '\'' => { - in_single_quote = true; + if in_single_quote { current.push(ch); - } - '"' => { - in_double_quote = true; - current.push(ch); - } - ';' => { - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); + if ch == '\'' { + if chars.peek() == Some(&'\'') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_single_quote = false; + } } - current.clear(); + continue; + } + + if in_double_quote { + current.push(ch); + if ch == '"' { + if chars.peek() == Some(&'"') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_double_quote = false; + } + } + continue; + } + + match ch { + '-' if chars.peek() == Some(&'-') => { + chars.next(); + in_line_comment = true; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + in_block_comment = true; + } + '\'' => { + in_single_quote = true; + current.push(ch); + } + '"' => { + in_double_quote = true; + current.push(ch); + } + ';' => { + let statement = current.trim(); + if !statement.is_empty() { + return Some(statement.to_string()); + } + current.clear(); + } + _ => current.push(ch), } - _ => current.push(ch), } - } - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); - } - - statements + let statement = current.trim(); + (!statement.is_empty()).then(|| statement.to_string()) + }) } fn ddl_statement_for_schema(schema: &str, sql: String) -> DdlStatement { diff --git a/src/cli/src/data/import_v2/executor.rs b/src/cli/src/data/import_v2/executor.rs index 3f2bf66ae68..6f152bc5c6d 100644 --- a/src/cli/src/data/import_v2/executor.rs +++ b/src/cli/src/data/import_v2/executor.rs @@ -14,8 +14,18 @@ //! DDL execution for import. +use std::collections::HashSet; + +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_telemetry::info; +use servers::query_handler::sql::{MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS}; use snafu::ResultExt; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::statement::Statement; +use store_api::metric_engine_consts::{ + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; use crate::data::import_v2::error::{DatabaseSnafu, Result}; use crate::database::DatabaseClient; @@ -54,69 +64,379 @@ impl<'a> DdlExecutor<'a> { Self { client } } - /// Executes a list of DDL statements, stopping on first error. - pub async fn execute_strict(&self, statements: &[DdlStatement]) -> Result<()> { - let total = statements.len(); - + /// Executes DDL in order, stopping at the first failed request. + pub async fn execute_strict( + &self, + statements: &[DdlStatement], + batch_supported: bool, + ) -> Result<()> { + let mut batch = LogicalBatch::default(); for (i, stmt) in statements.iter().enumerate() { - let preview = preview_sql(&stmt.sql); - - info!("Executing DDL ({}/{}): {}", i + 1, total, preview); - - if let Some(schema) = stmt.execution_schema.as_deref() { + info!("Executing DDL ({}/{})", i + 1, statements.len()); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let candidate = if batch_supported { + classify(stmt, self.client.catalog()).context(DatabaseSnafu)? + } else { + None + }; + if let Some(candidate) = candidate { + if !batch.accepts(&candidate, &stmt.sql) { + batch.flush(self.client).await?; + } + batch.push(candidate, &stmt.sql); + } else { + batch.flush(self.client).await?; self.client .sql(&stmt.sql, schema) .await .context(DatabaseSnafu)?; - } else { - self.client - .sql_in_public(&stmt.sql) - .await - .context(DatabaseSnafu)?; } } + batch.flush(self.client).await + } +} +struct LogicalTable { + catalog: String, + execution_schema: String, + schema: String, + physical: String, + name: String, +} + +fn classify(stmt: &DdlStatement, catalog: &str) -> crate::error::Result> { + if stmt.sql.len() > MAX_LOGICAL_TABLE_DDL_BYTES { + return Ok(None); + } + let parsed = parse_ddl(&stmt.sql)?; + let [Statement::CreateTable(create)] = parsed.as_slice() else { + return Ok(None); + }; + let Some(physical) = create.options.get(LOGICAL_TABLE_METADATA_KEY) else { + return Ok(None); + }; + if create.engine != METRIC_ENGINE_NAME + || create.options.value(PHYSICAL_TABLE_METADATA_KEY).is_some() + { + return Ok(None); + } + let parts: Vec<_> = create + .name + .0 + .iter() + .filter_map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let (catalog, schema, name) = match parts.as_slice() { + [name] => (catalog, schema, *name), + [schema, name] => (catalog, *schema, *name), + [catalog, schema, name] => (*catalog, *schema, *name), + _ => return Ok(None), + }; + Ok(Some(LogicalTable { + catalog: catalog.into(), + execution_schema: stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME) + .into(), + schema: schema.into(), + physical: physical.into(), + name: name.into(), + })) +} + +pub(crate) fn parse_ddl(sql: &str) -> crate::error::Result> { + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()).map_err( + |_| { + crate::error::InvalidArgumentsSnafu { + msg: "cannot parse import DDL", + } + .build() + }, + ) +} + +#[derive(Default)] +struct LogicalBatch { + group: Option, + names: HashSet, + sql: String, +} + +impl LogicalBatch { + fn accepts(&self, table: &LogicalTable, sql: &str) -> bool { + self.group.as_ref().is_none_or(|g| { + g.catalog == table.catalog + && g.schema == table.schema + && g.physical == table.physical + && g.execution_schema == table.execution_schema + }) && !self.names.contains(&table.name) + && self.names.len() < MAX_LOGICAL_TABLE_DDL_STATEMENTS + && self.sql.len() + (if self.sql.is_empty() { 0 } else { 2 }) + sql.len() + <= MAX_LOGICAL_TABLE_DDL_BYTES + } + + fn push(&mut self, table: LogicalTable, sql: &str) { + if !self.sql.is_empty() { + self.sql.push_str(";\n"); + } + self.sql.push_str(sql); + self.names.insert(table.name.clone()); + self.group = Some(table); + } + + async fn flush(&mut self, client: &DatabaseClient) -> Result<()> { + if let Some(group) = self.group.as_ref() { + client + .logical_tables(&self.sql, &group.execution_schema, self.names.len()) + .await + .context(DatabaseSnafu)?; + self.sql.clear(); + self.names.clear(); + self.group = None; + } Ok(()) } } -fn preview_sql(sql: &str) -> String { - let mut chars = sql.chars(); - let preview: String = chars.by_ref().take(80).collect(); - if chars.next().is_some() { - format!("{preview}...") - } else { - preview - } -} #[cfg(test)] mod tests { use super::*; - #[test] - fn test_statement_without_execution_schema_uses_public() { - let stmt = DdlStatement::new("CREATE DATABASE IF NOT EXISTS test_db".to_string()); - assert_eq!(stmt.execution_schema, None); + fn logical(name: &str, physical: &str) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='{physical}')" + ) } #[test] - fn test_statement_with_execution_schema_preserves_context() { - let stmt = DdlStatement::with_execution_schema( - r#"CREATE TABLE IF NOT EXISTS "my""schema"."metrics" (ts TIMESTAMP TIME INDEX)"# - .to_string(), - r#"my"schema"#.to_string(), - ); - assert_eq!(stmt.execution_schema.as_deref(), Some(r#"my"schema"#)); + fn classification_and_batch_boundaries() { + let mut batch = LogicalBatch::default(); + for (name, context, physical, expected) in [ + ("A", "public", "p.q", true), + ("public.b", "other", "p.q", false), + ("greptime.public.c", "public", "p.q", true), + ("a", "public", "p.q", false), + ("\"A\"", "public", "p.q", true), + ("d", "public", "p", false), + ("d", "other", "p.q", false), + ("other.public.d", "public", "p.q", false), + ] { + let sql = logical(name, physical); + let stmt = DdlStatement::with_execution_schema(sql.clone(), context.into()); + let table = classify(&stmt, "greptime").unwrap().unwrap(); + assert_eq!(batch.accepts(&table, &sql), expected, "{name}"); + if expected { + batch.push(table, &sql); + } + } + for sql in [ + "CREATE DATABASE d", + "CREATE TABLE t (ts TIMESTAMP TIME INDEX)", + "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table='')", + "CREATE VIEW v AS SELECT * FROM a", + ] { + assert!( + classify(&DdlStatement::new(sql.into()), "greptime") + .unwrap() + .is_none() + ); + } + let mut batch = LogicalBatch::default(); + for i in 0..128 { + let sql = logical(&format!("t{i}"), "p"); + let table = classify(&DdlStatement::new(sql.clone()), "greptime") + .unwrap() + .unwrap(); + assert!(batch.accepts(&table, &sql)); + batch.push(table, &sql); + } + let table = classify(&DdlStatement::new(logical("next", "p")), "greptime") + .unwrap() + .unwrap(); + assert!(!batch.accepts(&table, &logical("next", "p"))); + batch.names.clear(); + batch.sql = "测".repeat(MAX_LOGICAL_TABLE_DDL_BYTES / 3 - 1); + assert!(batch.accepts(&table, "ab")); + assert!(!batch.accepts(&table, "测")); + let oversized = + DdlStatement::new(logical("large", &"x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES))); + assert!(classify(&oversized, "greptime").unwrap().is_none()); } - #[test] - fn test_preview_sql_truncates_at_char_boundary() { - let sql = format!( - "CREATE TABLE {} (ts TIMESTAMP TIME INDEX)", - "测".repeat(100) - ); - let preview = preview_sql(&sql); - assert!(preview.ends_with("...")); - assert!(preview.is_char_boundary(preview.len())); + #[tokio::test] + async fn oversized_ddl_uses_ordinary_sql_in_order() { + let large = "x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES); + let statements = [ + DdlStatement::new(logical("before", "p")), + DdlStatement::new(format!( + "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}')" + )), + DdlStatement::new(logical("between", "p")), + DdlStatement::new(format!("CREATE VIEW v AS SELECT '{large}' AS payload")), + DdlStatement::new(logical("oversized", &large)), + DdlStatement::new(logical("after", "p")), + ]; + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), statements.len()); + for (i, (request, statement)) in requests.iter().zip(&statements).enumerate() { + let endpoint = if [0, 2, 5].contains(&i) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{endpoint} "))); + let body = request.split_once("\r\n\r\n").unwrap().1; + let sql = url::form_urlencoded::parse(body.as_bytes()) + .find(|(key, _)| key == "sql") + .unwrap() + .1; + assert_eq!(sql, statement.sql); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn qualified_targets_preserve_request_context() { + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + let statements = [ + DdlStatement::new(logical("other.metrics.a", "p")), + DdlStatement::with_execution_schema( + logical("other.metrics.b", "p"), + "existing_schema".into(), + ), + ]; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + for (request, schema) in requests.iter().zip(["public", "existing_schema"]) { + assert!(request.starts_with("POST /v1/ddl/logical-tables ")); + let body = request.split_once("\r\n\r\n").unwrap().1; + let db = url::form_urlencoded::parse(body.as_bytes()) + .find(|(k, _)| k == "db") + .unwrap() + .1; + assert_eq!(db, format!("greptime-{schema}")); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn transport_preserves_order_and_stops_without_fallback() { + use crate::database::tests::test_server; + let statements = [ + DdlStatement::new("CREATE DATABASE d".into()), + DdlStatement::new(logical("a", "p").replace( + "on_physical_table='p'", + "on_physical_table='p', access_key_id='original-secret'", + )), + DdlStatement::new(logical("a", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for supported in [true, false] { + let (client, requests, server) = test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, supported) + .await + .unwrap(); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 4); + for (i, request) in requests.iter().enumerate() { + let path = if supported && (i == 1 || i == 2) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{path} "))); + assert!( + request + .to_lowercase() + .contains("authorization: basic dxnlcjpwyxnzd29yza==") + ); + let body = request.split_once("\r\n\r\n").unwrap().1; + let form: std::collections::HashMap<_, _> = + url::form_urlencoded::parse(body.as_bytes()) + .into_owned() + .collect(); + assert_eq!(form["sql"], statements[i].sql); + assert_eq!(form["db"], "greptime-public"); + } + server.abort(); + } + let statements = [ + DdlStatement::new(logical("a", "p")), + DdlStatement::new(logical("b", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for (status, body, success) in [ + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":0}]}"#, + true, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + false, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":1}]}"#, + false, + ), + (200, r#"{"execution_time_ms":0,"output":[{},{}]}"#, false), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"error":"secret"}]}"#, + false, + ), + (500, r#"{"error":"secret"}"#, false), + (200, "invalid", false), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await; + assert_eq!(result.is_ok(), success); + assert_eq!(requests.lock().unwrap().len(), if success { 2 } else { 1 }); + if let Err(error) = result { + assert!(!format!("{error:?}").contains("secret")); + } + server.abort(); + let _ = server.await; + } } } diff --git a/src/cli/src/data/import_v2/packed.rs b/src/cli/src/data/import_v2/packed.rs index 303c866ba69..c2f1ec35889 100644 --- a/src/cli/src/data/import_v2/packed.rs +++ b/src/cli/src/data/import_v2/packed.rs @@ -16,12 +16,11 @@ use std::collections::HashSet; use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; use snafu::ResultExt; -use sql::dialect::GreptimeDbDialect; -use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use crate::data::export_v2::manifest::{ChunkStatus, Manifest}; use crate::data::import_v2::error::{InvalidPackedSnapshotSnafu, Result, SnapshotStorageSnafu}; +use crate::data::import_v2::{command, executor}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; use crate::data::snapshot_storage::SnapshotStorage; @@ -51,49 +50,46 @@ pub(crate) async fn validate_snapshot( .read_text(&ddl_path_for_schema(schema)) .await .context(SnapshotStorageSnafu)?; - let statements = ParserContext::create_with_dialect( - &ddl, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .map_err(|e| { - InvalidPackedSnapshotSnafu { - reason: e.to_string(), - } - .build() - })?; let mut names = HashSet::new(); - for statement in statements { - if let Statement::CreateTable(create) = statement { - if create.engine.eq_ignore_ascii_case("metric") - && create.options.get("on_physical_table").is_none() - { - continue; + for sql in command::iter_ddl_statements(&ddl) { + let statements = executor::parse_ddl(&sql).map_err(|_| { + InvalidPackedSnapshotSnafu { + reason: "invalid snapshot DDL", } - let parts: Vec<_> = create - .name - .0 - .iter() - .map(|p| p.as_ident().map(|i| i.value.as_str())) - .collect(); - let allowed = match parts.as_slice() { - [Some(_)] => true, - [Some(s), Some(_)] => *s == schema, - [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, - _ => false, - }; - if !allowed { - return InvalidPackedSnapshotSnafu { - reason: "snapshot table DDL escapes its catalog/schema", + .build() + })?; + for statement in statements { + if let Statement::CreateTable(create) = statement { + if create.engine.eq_ignore_ascii_case("metric") + && create.options.get("on_physical_table").is_none() + { + continue; } - .fail(); - } - let name = parts.last().and_then(|p| *p).unwrap_or_default(); - if !names.insert(name.to_string()) { - return InvalidPackedSnapshotSnafu { - reason: "duplicate snapshot table DDL", + let parts: Vec<_> = create + .name + .0 + .iter() + .map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let allowed = match parts.as_slice() { + [Some(_)] => true, + [Some(s), Some(_)] => *s == schema, + [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, + _ => false, + }; + if !allowed { + return InvalidPackedSnapshotSnafu { + reason: "snapshot table DDL escapes its catalog/schema", + } + .fail(); + } + let name = parts.last().and_then(|p| *p).unwrap_or_default(); + if !names.insert(name.to_string()) { + return InvalidPackedSnapshotSnafu { + reason: "duplicate snapshot table DDL", + } + .fail(); } - .fail(); } } } @@ -180,13 +176,14 @@ mod tests { .unwrap(); let uri = url::Url::from_file_path(dir.path()).unwrap(); let storage = OpenDalStorage::from_uri(uri.as_str(), &Default::default()).unwrap(); - storage.write_text("schema/ddl/public.sql", "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\";").await.unwrap(); + let large = "x".repeat(servers::query_handler::sql::MAX_LOGICAL_TABLE_DDL_BYTES); + storage.write_text("schema/ddl/public.sql", &format!("CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\"; CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}'); CREATE VIEW large_view AS SELECT '{large}' AS payload;")).await.unwrap(); let path = "data/public/1/pack-index.json"; - let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":12}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0}]}); + let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":24}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0},{"table_name":"ordinary","object":"pack-0.bin","offset":12,"length":12,"row_count":0}]}); storage.write_text(path, &index.to_string()).await.unwrap(); assert!(dir.path().join(path).is_file()); storage - .write_text("data/public/1/pack-0.bin", "abcdefghijkl") + .write_text("data/public/1/pack-0.bin", "abcdefghijklmnopqrstuvwx") .await .unwrap(); let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index da277e166be..1efe6d7d0a6 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -26,9 +26,7 @@ use servers::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use snafu::ResultExt; -use crate::error::{ - BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result, SerdeJsonSnafu, -}; +use crate::error::{BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result}; #[derive(Debug, Clone)] pub struct DatabaseClient { @@ -100,49 +98,84 @@ impl DatabaseClient { self.sql(sql, DEFAULT_SCHEMA_NAME).await } - /// Requires the explicit packed-import protocol before any restore mutation. - pub async fn require_packed_import(&self) -> Result<()> { - self.require_packed_capability("metric_packed_import").await + pub(crate) fn catalog(&self) -> &str { + &self.catalog } /// Requires packed export support before creating snapshot artifacts. pub async fn require_packed_export(&self) -> Result<()> { - self.require_packed_capability("metric_packed_export").await + self.capabilities().await?.require("metric_packed_export") } - async fn require_packed_capability(&self, capability: &str) -> Result<()> { + /// Reads target capabilities once; older targets advertise no extensions. + pub(crate) async fn capabilities(&self) -> Result { let url = format!("http://{}/v1/capabilities", self.addr); let mut request = self.http_client().await?.get(&url).timeout(self.timeout); if let Some(auth) = &self.auth_header { request = request.header("Authorization", auth); } let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { - reason: "packed snapshot capability request failed", + reason: "capability request failed", })?; if response.status() == reqwest::StatusCode::NOT_FOUND { - return crate::error::InvalidArgumentsSnafu { - msg: format!("server does not support {capability}"), - } - .fail(); + return Ok(Capabilities(serde_json::json!({}))); } let response = response .error_for_status() .with_context(|_| HttpQuerySqlSnafu { - reason: "packed snapshot capability request rejected", + reason: "capability request rejected", })?; let body = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot read capability response", })?; - let value: Value = serde_json::from_str(&body).context(SerdeJsonSnafu)?; - if !value.is_object() { + let value: Value = serde_json::from_str(&body).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: "invalid capability response", + } + .build() + })?; + if !value.is_object() + || value.get("error").is_some() + || value.get("code").is_some_and(|v| v.as_u64() != Some(0)) + { return crate::error::UnexpectedSnafu { msg: "invalid capability response: expected JSON object", } .fail(); } - if value.get(capability).and_then(Value::as_u64) != Some(1) { + for key in [ + "metric_batch_ddl", + "metric_packed_import", + "metric_packed_export", + ] { + if value.get(key).is_some_and(|v| v.as_u64().is_none()) { + return crate::error::UnexpectedSnafu { + msg: "invalid capability version", + } + .fail(); + } + } + Ok(Capabilities(value)) + } + + /// Submits a batch once. An ambiguous failure must not be replayed here. + pub(crate) async fn logical_tables(&self, sql: &str, schema: &str, count: usize) -> Result<()> { + let db = format!("{}-{}", self.catalog, schema); + if 8 + form_encoded_len(&db) + form_encoded_len(sql) > 4 * 1024 * 1024 { return crate::error::InvalidArgumentsSnafu { - msg: format!("server does not support {capability} version 1"), + msg: "logical-table batch exceeds 4 MiB form limit", + } + .fail(); + } + let response = self.ddl_response(sql, schema, "ddl/logical-tables").await?; + if response.output().len() != count + || response + .output() + .iter() + .any(|output| !matches!(output, GreptimeQueryOutput::AffectedRows(0))) + { + return crate::error::UnexpectedSnafu { + msg: "invalid logical-table batch output", } .fail(); } @@ -163,7 +196,16 @@ impl DatabaseClient { sql: &str, schema: &str, ) -> Result { - let url = format!("http://{}/v1/sql", self.addr); + self.ddl_response(sql, schema, "sql").await + } + + async fn ddl_response( + &self, + sql: &str, + schema: &str, + endpoint: &str, + ) -> Result { + let url = format!("http://{}/v1/{endpoint}", self.addr); let params = [ ("db", format!("{}-{}", self.catalog, schema)), ("sql", sql.to_string()), @@ -172,8 +214,21 @@ impl DatabaseClient { .http_client() .await? .post(&url) - .form(¶ms) .header("Content-Type", "application/x-www-form-urlencoded"); + if endpoint == "ddl/logical-tables" { + let body = url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(¶ms) + .finish(); + // A streaming body cannot be cloned for reqwest's protocol retries. + request = request + .timeout(self.timeout) + .header(reqwest::header::CONTENT_LENGTH, body.len()) + .body(reqwest::Body::wrap_stream(futures::stream::once( + std::future::ready(Ok::<_, std::io::Error>(body)), + ))); + } else { + request = request.form(¶ms); + } if let Some(ref auth) = self.auth_header { request = request.header("Authorization", auth); } @@ -239,7 +294,19 @@ impl DatabaseClient { async fn http_client(&self) -> Result<&reqwest::Client> { self.client .get_or_try_init(|| async { - let mut builder = reqwest::Client::builder(); + let mut builder = reqwest::Client::builder().redirect( + reqwest::redirect::Policy::custom(|attempt| { + if attempt + .previous() + .first() + .is_some_and(|url| url.path() == "/v1/ddl/logical-tables") + { + attempt.stop() + } else { + reqwest::redirect::Policy::default().redirect(attempt) + } + }), + ); if let Some(proxy) = self.proxy.clone() { builder = builder.proxy(proxy); } @@ -252,6 +319,30 @@ impl DatabaseClient { } } +fn form_encoded_len(value: &str) -> usize { + url::form_urlencoded::byte_serialize(value.as_bytes()) + .map(str::len) + .sum() +} + +pub(crate) struct Capabilities(Value); + +impl Capabilities { + pub(crate) fn supports(&self, capability: &str) -> bool { + self.0.get(capability).and_then(Value::as_u64) == Some(1) + } + + pub(crate) fn require(&self, capability: &str) -> Result<()> { + if !self.supports(capability) { + return crate::error::InvalidArgumentsSnafu { + msg: format!("server does not support {capability} version 1"), + } + .fail(); + } + Ok(()) + } +} + /// Split at `-`. pub(crate) fn split_database(database: &str) -> Result<(String, Option)> { let (catalog, schema) = match database.split_once('-') { @@ -267,7 +358,7 @@ pub(crate) fn split_database(database: &str) -> Result<(String, Option)> } #[cfg(test)] -mod tests { +pub(crate) mod tests { #[tokio::test] async fn packed_capability_probe_checks_auth_and_protocol_version() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -312,7 +403,11 @@ mod tests { if export { client.require_packed_export().await.is_ok() } else { - client.require_packed_import().await.is_ok() + client + .capabilities() + .await + .and_then(|c| c.require("metric_packed_import")) + .is_ok() }, supported, "{status}: {body}" @@ -324,6 +419,173 @@ mod tests { use super::*; + #[test] + fn form_size_counts_utf8_and_escaping() { + for (db, sql) in [ + ("greptime-public", "SELECT 1"), + ("测-试", "+&='秘密';\n"), + ("a_b.*", "a~b"), + ] { + let encoded = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("db", db) + .append_pair("sql", sql) + .finish(); + assert_eq!( + 8 + form_encoded_len(db) + form_encoded_len(sql), + encoded.len() + ); + } + } + + #[tokio::test] + async fn batch_capabilities_distinguish_absence_from_errors() { + for (status, body, expected) in [ + (200, r#"{"metric_batch_ddl":1}"#, Some(true)), + (200, r#"{"metric_batch_ddl":2}"#, Some(false)), + (200, "{}", Some(false)), + (404, "", Some(false)), + (401, "{}", None), + (403, "{}", None), + (500, "{}", None), + (200, "[]", None), + (200, "bad", None), + (200, r#"{"metric_batch_ddl":"1"}"#, None), + (200, r#"{"error":"secret"}"#, None), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = client.capabilities().await; + assert_eq!( + result.as_ref().ok().map(|c| c.supports("metric_batch_ddl")), + expected + ); + assert_eq!(requests.lock().unwrap().len(), 1); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_redirects_are_not_replayed() { + for status in [303, 307, 308] { + let (client, requests, server) = test_server( + status, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + assert_eq!(requests.lock().unwrap().len(), 1); + assert!(requests.lock().unwrap()[0].starts_with("POST /v1/ddl/logical-tables ")); + // Ordinary SQL retains the shared client's existing redirect behavior. + assert!(client.sql("SELECT 1", "public").await.is_err()); + assert!(requests.lock().unwrap().len() > 2); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_disconnect_and_timeout_are_errors() { + use tokio::io::AsyncReadExt; + for disconnect in [true, false] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + None, + Duration::from_millis(100), + None, + true, + ); + let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + if !disconnect { + let _ = stopped.await; + } + }); + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + let _ = stop.send(()); + server.await.unwrap(); + client.catalog = "测".repeat(500_000); + let error = client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("4 MiB form limit")); + } + } + + pub(crate) async fn test_server( + status: u16, + body: &str, + ) -> ( + DatabaseClient, + Arc>>, + tokio::task::JoinHandle<()>, + ) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + Some("user:password".into()), + Duration::from_secs(5), + None, + true, + ); + let requests = Arc::new(std::sync::Mutex::new(Vec::new())); + let captured = requests.clone(); + let body = body.to_string(); + let server = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + loop { + let mut buf = [0; 4096]; + let n = socket.read(&mut buf).await.unwrap(); + if n == 0 { + break; + } + bytes.extend_from_slice(&buf[..n]); + if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&bytes[..end]).to_lowercase(); + let length = headers + .lines() + .find_map(|l| l.strip_prefix("content-length: ")) + .map(|n| n.parse::().unwrap()) + .unwrap_or(0); + if bytes.len() >= end + 4 + length { + break; + } + } + } + captured + .lock() + .unwrap() + .push(String::from_utf8(bytes).unwrap()); + let location = if (300..400).contains(&status) { + "Location: /v1/sql\r\n" + } else { + "" + }; + socket.write_all(format!("HTTP/1.1 {status} Response\r\n{location}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + } + }); + (client, requests, server) + } + #[test] fn test_split_database() { let result = split_database("catalog-schema").unwrap(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index da1232a243b..8e3057bb7aa 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1416,6 +1416,59 @@ async fn export_http(instance: Arc) -> (String, tokio::task::JoinHandl (addr, task) } +#[derive(Default)] +struct ImportRequests { + batches: std::sync::atomic::AtomicUsize, + fail_batch: std::sync::atomic::AtomicBool, + capabilities: std::sync::atomic::AtomicUsize, +} + +async fn import_http( + instance: Arc, + older: bool, +) -> (String, tokio::task::JoinHandle<()>, Arc) { + use axum::response::IntoResponse; + use std::sync::atomic::Ordering; + let requests = Arc::new(ImportRequests::default()); + requests.fail_batch.store(!older, Ordering::SeqCst); + let observed = requests.clone(); + let server = servers::http::HttpServerBuilder::new(Default::default()) + .with_sql_handler(instance) + .build(); + let app = server + .build(server.make_app()) + .unwrap() + .layer(axum::middleware::from_fn( + move |request: axum::extract::Request, next: axum::middleware::Next| { + let observed = observed.clone(); + async move { + let batch = request.uri().path() == "/v1/ddl/logical-tables"; + if request.uri().path() == "/v1/capabilities" { + observed.capabilities.fetch_add(1, Ordering::SeqCst); + if older { + return axum::http::StatusCode::NOT_FOUND.into_response(); + } + } + if batch { + observed.batches.fetch_add(1, Ordering::SeqCst); + } + let response = next.run(request).await; + if batch && observed.fail_batch.swap(false, Ordering::SeqCst) { + assert!(response.status().is_success()); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + response + } + }, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (addr, task, requests) +} + struct FailSecondChunk { fail: std::sync::atomic::AtomicBool, copies: std::sync::atomic::AtomicUsize, @@ -1451,7 +1504,12 @@ impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { #[tokio::test] async fn metric_export_v2_cli_resume_roundtrip() { - metric_export_v2_cli_roundtrip(false, false).await; + metric_export_v2_cli_roundtrip(false, false, false).await; +} + +#[tokio::test] +async fn metric_import_v2_older_target_roundtrip() { + metric_export_v2_cli_roundtrip(false, false, true).await; } #[tokio::test] @@ -1461,13 +1519,13 @@ async fn metric_export_v2_cli_s3_resume_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) { - metric_export_v2_cli_roundtrip(true, false).await; + metric_export_v2_cli_roundtrip(true, false, false).await; } } #[tokio::test(flavor = "multi_thread")] async fn packed_export_v2_cli_local_roundtrip() { - metric_export_v2_cli_roundtrip(false, true).await; + metric_export_v2_cli_roundtrip(false, true, false).await; } #[tokio::test(flavor = "multi_thread")] @@ -1475,7 +1533,7 @@ async fn packed_export_v2_cli_s3_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) { - metric_export_v2_cli_roundtrip(true, true).await; + metric_export_v2_cli_roundtrip(true, true, false).await; } } @@ -1508,7 +1566,7 @@ async fn exported_table_path( ) } -async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { +async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { use servers::interceptor::SqlQueryInterceptorRef; let plugins = common_base::Plugins::new(); let faults = Arc::new(FailSecondChunk { @@ -1864,7 +1922,8 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { ) .await; } - let (target_addr, target_server) = export_http(target.fe_instance().clone()).await; + let (target_addr, target_server, requests) = + import_http(target.fe_instance().clone(), older).await; let state = destination.path().join("restore-state.json"); let mut import = vec![ "import-v2", @@ -1879,7 +1938,29 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { "never", ]; import.extend(storage_args.iter().map(String::as_str)); + use std::sync::atomic::Ordering; + if !older { + assert!(run_data_cli(&import).await.is_err()); + assert_eq!(requests.batches.load(Ordering::SeqCst), 1); + assert_eq!(requests.capabilities.load(Ordering::SeqCst), 1); + let failed: serde_json::Value = + serde_json::from_slice(&std::fs::read(&state).unwrap()).unwrap(); + assert_eq!(failed["ddl_completed"], false); + assert!( + failed["tasks"] + .as_array() + .unwrap() + .iter() + .all(|t| t["status"] == "pending") + ); + } run_data_cli(&import).await.unwrap(); + assert_eq!( + requests.capabilities.load(Ordering::SeqCst), + if older { 1 } else { 2 } + ); + assert_eq!(requests.batches.load(Ordering::SeqCst) > 1, !older); + assert!(!state.exists()); for name in names { let source_table = table(instance, &name).await; let target_table = table(target.fe_instance(), &name).await; @@ -1954,6 +2035,42 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { assert!(schema_manifest.schema_only && schema_manifest.chunks.is_empty()); assert_eq!(schema_manifest.version, 1); assert!(schema_manifest.data_layout.is_none()); + let schema_target = GreptimeDbStandaloneBuilder::new("schema_only_target") + .build() + .await; + let (schema_addr, schema_server, schema_requests) = + import_http(schema_target.fe_instance().clone(), false).await; + schema_requests.fail_batch.store(false, Ordering::SeqCst); + let schema_state = destination.path().join("schema-only-state.json"); + let mut schema_import = vec![ + "import-v2", + "--addr", + &schema_addr, + "--from", + &schema_uri, + "--state-path", + schema_state.to_str().unwrap(), + "--no-proxy", + "--progress", + "never", + ]; + schema_import.extend(storage_args.iter().map(String::as_str)); + let probes = schema_requests.capabilities.load(Ordering::SeqCst); + schema_import.push("--dry-run"); + run_data_cli(&schema_import).await.unwrap(); + assert_eq!(schema_requests.capabilities.load(Ordering::SeqCst), probes); + schema_import.pop(); + let batches = schema_requests.batches.load(Ordering::SeqCst); + run_data_cli(&schema_import).await.unwrap(); + assert!(schema_requests.batches.load(Ordering::SeqCst) > batches); + assert!(!schema_state.exists()); + assert!( + values(schema_target.fe_instance(), "SELECT * FROM dashboard") + .await + .is_empty() + ); + schema_server.abort(); + let _ = schema_server.await; } server.abort(); target_server.abort(); From e4872ec1a2d0d040d9f05c370f238cd7f411328e Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 8 Oct 2026 13:16:14 +0000 Subject: [PATCH 108/133] feat(log-store): collect object store WAL objects below a persisted reclaim boundary (#9471) * feat(log-store): collect object store WAL objects below a persisted reclaim boundary Every object header carries a reclaim boundary: every object below it was fully flushed when the object was sealed. The writer computes the boundary from the catalog and the obsolete watermarks after every obsolete and every indexed object, and stamps it on every sealed object and on the start object. The boundary of the last indexed object permits deletion: every present object below it, indexed, listed by recovery off the chain, or left by a failed batch, is deleted in the background, at most four at a time. Recovery judges every candidate with its own boundary: a link to a predecessor below it holds whether or not the predecessor is present, and the chain stops there. A read that listed an object before it was collected skips it instead of failing. The header grows from 46 to 54 bytes; the format version stays 1. Signed-off-by: jeremyhi * perf(log-store): scan the reclaim boundary from the previous boundary The boundary was recomputed by scanning every region's indexed objects from the first one, so the flushed objects of busy regions were scanned again on every indexed object while another region held the boundary. Scan the objects in sequence order from the previous boundary instead and stop at the first one holding a live segment. Watermarks never move down and a region first appears above the boundary, so the result is the same and an object is not scanned again once the boundary passed it. Signed-off-by: jeremyhi * test(log-store): name the parked operation helpers after both their uses Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/metrics.rs | 14 + src/log-store/src/object_store_wal.rs | 24 +- src/log-store/src/object_store_wal/catalog.rs | 172 ++- src/log-store/src/object_store_wal/format.rs | 70 +- src/log-store/src/object_store_wal/io.rs | 15 + src/log-store/src/object_store_wal/store.rs | 1276 +++++++++++++++-- .../engine/object_store_wal_recovery_test.rs | 108 +- tests-integration/tests/object_store_wal.rs | 67 +- 8 files changed, 1536 insertions(+), 210 deletions(-) diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index 824058d70af..badb5a80b70 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -97,4 +97,18 @@ lazy_static! { "object store logstore skipped corrupted segments total", ) .unwrap(); + /// Counter of objects the object store logstore deleted because no + /// region needs them any more. + pub static ref METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_deleted_objects_total", + "object store logstore deleted objects total", + ) + .unwrap(); + /// Counter of object deletes of the object store logstore that failed and + /// are attempted again at the next collection. + pub static ref METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_failed_deletes_total", + "object store logstore failed object deletes total", + ) + .unwrap(); } diff --git a/src/log-store/src/object_store_wal.rs b/src/log-store/src/object_store_wal.rs index d4bfc096393..6a25afb3e77 100644 --- a/src/log-store/src/object_store_wal.rs +++ b/src/log-store/src/object_store_wal.rs @@ -24,9 +24,9 @@ //! //! The header carries the magic `GTWALOBJ`, the format version, the object //! sequence, the epoch of the instance that wrote the object, the link to the -//! object it extends, and a CRC32 of the header. Each segment holds the -//! entries of exactly one region, ordered by entry id, and segments are ordered -//! by region id. The footer indexes every segment with its region id, entry id +//! object it extends, the reclaim boundary, and a CRC32 of the header. Each +//! segment holds the entries of exactly one region, ordered by entry id, and +//! segments are ordered by region id. The footer indexes every segment with its region id, entry id //! range, byte range and CRC32. The fixed-size trailer points at the footer and //! carries the CRC32 of the footer and of the whole object, so a reader locates //! the footer by reading the fixed-length trailer at the end of the object. @@ -40,12 +40,18 @@ //! object it extends by sequence and epoch, and recovery replays only the chain //! that ends at the complete object with the largest epoch and sequence. An //! object is complete when every link on its chain names a present object of -//! the recorded epoch, back to an object that starts the chain. Objects off the -//! chain are orphans, which are never replayed but keep their sequences. Each -//! open writes an object without segments above every present object before -//! it accepts writes, and its epoch is one above the sequence of that object, -//! so no two instances share an epoch and a late object of an earlier -//! instance never ends the chain. +//! the recorded epoch, back to an object that starts the chain or to a link +//! below the reclaim boundary of the object. Objects off the chain are +//! orphans, which are never replayed but keep their sequences. Each open +//! writes an object without segments above every present object before it +//! accepts writes, and its epoch is one above the sequence of that object, so +//! no two instances share an epoch and a late object of an earlier instance +//! never ends the chain. +//! +//! The reclaim boundary of an object states that every object with a lower +//! sequence was fully flushed when the object was sealed, so the chain need +//! not reach below it. Every present object below the boundary of the last +//! indexed object is deleted, on the chain or not. //! //! Recovery lists the objects, reads and verifies only the header, trailer and //! footer of each, and indexes the footers of the chain in sequence order to diff --git a/src/log-store/src/object_store_wal/catalog.rs b/src/log-store/src/object_store_wal/catalog.rs index 4866da359cf..70f7aa5b329 100644 --- a/src/log-store/src/object_store_wal/catalog.rs +++ b/src/log-store/src/object_store_wal/catalog.rs @@ -14,10 +14,11 @@ //! In-memory index over the footers of the objects of one WAL prefix. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::ops::Bound::{Excluded, Unbounded}; use snafu::{OptionExt, ensure}; +use store_api::logstore::EntryId; use store_api::storage::RegionId; use crate::error::{ @@ -32,6 +33,9 @@ use crate::object_store_wal::format::FooterEntry; pub(crate) struct ObjectCatalog { objects: BTreeMap>, regions: BTreeMap>, + /// Largest entry id ever indexed per region, kept after the object that + /// held it was removed. + max_entry_ids: HashMap, } impl ObjectCatalog { @@ -123,11 +127,71 @@ impl ObjectCatalog { .entry(entry.region_id) .or_default() .insert(object_seq, entry.clone()); + self.max_entry_ids + .entry(entry.region_id) + .and_modify(|max| *max = (*max).max(entry.max_entry_id)) + .or_insert(entry.max_entry_id); } self.objects.insert(object_seq, footer); Ok(()) } + /// Removes the object `object_seq` from the index once it was deleted and + /// returns the footer it held, or `None` if it was not indexed. The + /// largest entry id of every region it held is kept. + pub(crate) fn remove_object(&mut self, object_seq: u64) -> Option> { + let footer = self.objects.remove(&object_seq)?; + for entry in &footer { + if let Some(objects) = self.regions.get_mut(&entry.region_id) { + objects.remove(&object_seq); + if objects.is_empty() { + self.regions.remove(&entry.region_id); + } + } + } + Some(footer) + } + + /// Returns true while the object `object_seq` is indexed. + pub(crate) fn contains_object(&self, object_seq: u64) -> bool { + self.objects.contains_key(&object_seq) + } + + /// Returns the reclaim boundary the indexed objects and the obsolete + /// watermarks allow, given that every object below `from` holds no live + /// segment: the first indexed object at or above `from` that holds a + /// live segment, or the sequence after the last indexed object when none + /// does, and never less than `from`. + /// + /// A segment is live when its maximum entry id is above the watermark of + /// its region, or its region has no watermark; a region whose watermark + /// is [`EntryId::MAX`] holds no live segment. Watermarks never move down + /// and a region first appears in an object above the boundary, so the + /// scan starts at the previous boundary and never reads an object twice + /// once the boundary has passed it. + pub(crate) fn reclaim_boundary( + &self, + from: u64, + obsolete_entry_ids: &HashMap, + ) -> u64 { + let live = self.objects.range(from..).find(|(_, footer)| { + footer.iter().any(|entry| { + obsolete_entry_ids + .get(&entry.region_id) + .is_none_or(|obsolete| entry.max_entry_id > *obsolete) + }) + }); + match live { + Some((&object_seq, _)) => object_seq, + None => self + .objects + .last_key_value() + .map_or(from, |(&object_seq, _)| { + object_seq.saturating_add(1).max(from) + }), + } + } + /// Returns the objects that hold entries of `region_id` overlapping /// `start_entry_id..=end_entry_id`, ordered by object sequence. pub(crate) fn objects_for_entry_range( @@ -157,12 +221,10 @@ impl ObjectCatalog { .collect()) } - /// Returns the largest entry id indexed for `region_id`. + /// Returns the largest entry id ever indexed for `region_id`, which is + /// kept after the object holding it was removed. pub(crate) fn region_max_entry_id(&self, region_id: RegionId) -> Option { - self.regions - .get(®ion_id)? - .last_key_value() - .map(|(_, entry)| entry.max_entry_id) + self.max_entry_ids.get(®ion_id).copied() } /// Returns the sequence to assign to the next object written after recovery. @@ -183,9 +245,9 @@ impl ObjectCatalog { .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, }; let floor = self - .regions - .keys() - .filter_map(|region_id| self.region_max_entry_id(*region_id)) + .max_entry_ids + .values() + .copied() .map(sequence_floor) .max() .unwrap_or(0); @@ -431,6 +493,98 @@ mod tests { ); } + #[test] + fn test_catalog_reclaim_boundary_follows_the_watermarks() { + let region_one = RegionId::new(1, 1); + let region_two = RegionId::new(2, 1); + let mut catalog = ObjectCatalog::default(); + catalog.insert_object(0, Vec::new()).unwrap(); + catalog + .insert_object( + 1, + vec![ + footer_entry(region_one, entry_id(1, 1), entry_id(1, 2)), + footer_entry(region_two, entry_id(1, 1), entry_id(1, 1)), + ], + ) + .unwrap(); + catalog + .insert_object( + 2, + vec![footer_entry(region_one, entry_id(2, 1), entry_id(2, 1))], + ) + .unwrap(); + catalog + .insert_object( + 3, + vec![footer_entry(region_two, entry_id(3, 1), entry_id(3, 1))], + ) + .unwrap(); + let boundary = |catalog: &ObjectCatalog, obsolete: &[(RegionId, EntryId)]| { + catalog.reclaim_boundary(0, &obsolete.iter().copied().collect()) + }; + + // A region without a watermark holds the boundary at its first object. + assert_eq!(1, boundary(&catalog, &[])); + assert_eq!(1, boundary(&catalog, &[(region_one, entry_id(2, 1))])); + // A watermark inside an object keeps that object. + assert_eq!( + 1, + boundary( + &catalog, + &[(region_one, entry_id(1, 1)), (region_two, EntryId::MAX)] + ) + ); + assert_eq!( + 2, + boundary( + &catalog, + &[(region_one, entry_id(1, 2)), (region_two, entry_id(1, 1))] + ) + ); + assert_eq!( + 3, + boundary( + &catalog, + &[(region_one, entry_id(2, 1)), (region_two, entry_id(1, 1))] + ) + ); + // A region of `obsolete_all` contributes nothing, and when no region + // does, the boundary is past the last object. + assert_eq!( + 4, + boundary( + &catalog, + &[(region_one, entry_id(2, 1)), (region_two, EntryId::MAX)] + ) + ); + + // Objects below the previous boundary are not scanned again. + assert_eq!(3, catalog.reclaim_boundary(3, &HashMap::new())); + assert_eq!(5, catalog.reclaim_boundary(5, &HashMap::new())); + + // Removing objects keeps the largest entry id of every region, and a + // region without indexed objects contributes nothing. + for object_seq in [0, 1, 2] { + assert!(catalog.remove_object(object_seq).is_some()); + } + assert!(catalog.remove_object(2).is_none()); + assert!(!catalog.contains_object(2)); + assert!(catalog.contains_object(3)); + assert_eq!( + Some(entry_id(2, 1)), + catalog.region_max_entry_id(region_one) + ); + assert!( + catalog + .objects_for_entry_range(region_one, 0, EntryId::MAX) + .unwrap() + .is_empty() + ); + assert_eq!(3, boundary(&catalog, &[])); + assert_eq!(4, catalog.next_object_seq().unwrap()); + } + fn footer_entry(region_id: RegionId, min_entry_id: u64, max_entry_id: u64) -> FooterEntry { FooterEntry { region_id, diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index 6dd46c87242..bd7e9ae6410 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -32,8 +32,9 @@ const FORMAT_VERSION: u16 = 1; const NO_PREDECESSOR: u64 = u64::MAX; /// Length of the object header: magic, version, object sequence, epoch, -/// predecessor sequence, predecessor epoch and the CRC32 of the header. -pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 8 + 8 + 8 + 4; +/// predecessor sequence, predecessor epoch, reclaim boundary and the CRC32 of +/// the header. +pub(crate) const HEADER_LEN: usize = 8 + 2 + 8 + 8 + 8 + 8 + 8 + 4; /// Length of the fixed trailer: footer offset, footer length, footer CRC32, /// object CRC32 and magic. pub(crate) const TRAILER_LEN: usize = 8 + 8 + 4 + 4 + 8; @@ -55,6 +56,10 @@ pub(crate) struct Header { pub(crate) epoch: u64, /// The object this object extends, `None` for the first object of a chain. pub(crate) prev: Option, + /// Every object below this sequence was fully flushed when this object + /// was sealed, so the chain recovery replays need not reach below it. It + /// is at or below the sequence of the object. + pub(crate) reclaim_boundary: u64, } /// Names the object another object extends: its sequence and the epoch @@ -224,6 +229,7 @@ pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ let epoch = reader.u64("header")?; let prev_object_seq = reader.u64("header")?; let prev_epoch = reader.u64("header")?; + let reclaim_boundary = reader.u64("header")?; let header_crc32 = reader.u32("header")?; let checksum = crc32fast::hash(&bytes[..HEADER_LEN - 4]); ensure!( @@ -232,6 +238,14 @@ pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ reason: checksum_mismatch("header", header_crc32, checksum), } ); + ensure!( + reclaim_boundary <= object_seq, + CorruptedWalObjectSnafu { + reason: format!( + "reclaim boundary {reclaim_boundary} is above the object sequence {object_seq}" + ), + } + ); let prev = (prev_object_seq != NO_PREDECESSOR).then_some(ChainLink { object_seq: prev_object_seq, epoch: prev_epoch, @@ -240,6 +254,7 @@ pub(crate) fn decode_header(bytes: &[u8]) -> Result
{ object_seq, epoch, prev, + reclaim_boundary, }) } @@ -582,6 +597,7 @@ fn encode_header(header: &Header, output: &mut BytesMut) { }); output.put_u64(prev.object_seq); output.put_u64(prev.epoch); + output.put_u64(header.reclaim_boundary); output.put_u32(crc32fast::hash(&output[start..])); } @@ -730,6 +746,7 @@ mod tests { object_seq: 40, epoch: 4, }), + reclaim_boundary: 37, } } @@ -838,7 +855,7 @@ mod tests { // Every field after the version is covered by the header checksum, // which recovery verifies without reading the whole object. - for offset in [10, 18, 26, 34, HEADER_LEN - 1] { + for offset in [10, 18, 26, 34, 42, HEADER_LEN - 1] { let mut bad_field = encoded.bytes.to_vec(); bad_field[offset] ^= 1; assert_corrupted(decode_header(&bad_field), "header checksum mismatch"); @@ -855,6 +872,19 @@ mod tests { assert_eq!(first, decode_object(&encoded.bytes).unwrap().header); } + #[test] + fn test_format_rejects_a_reclaim_boundary_above_the_object() { + let above = Header { + reclaim_boundary: header().object_seq + 1, + ..header() + }; + let encoded = encode_object(above, &records()).unwrap(); + assert_corrupted( + decode_header(&encoded.bytes), + "reclaim boundary 43 is above the object sequence 42", + ); + } + #[test] fn test_format_rejects_trailer_magic_and_invalid_footer_range() { let encoded = encode_object(header(), &records()).unwrap(); @@ -971,15 +1001,16 @@ mod tests { /// the encoder and the decoder change together. const FIXTURE_V1_HEX: &str = concat!( // Header: magic, version 1, object sequence 7, epoch 3, predecessor - // sequence 6, predecessor epoch 2, header CRC32. + // sequence 6, predecessor epoch 2, reclaim boundary 5, header CRC32. "475457414c4f424a", "0001", "0000000000000007", "0000000000000003", "0000000000000006", "0000000000000002", - "1ca121b1", - // Segment of region 1 at offset 46: region id, entry count, then + "0000000000000005", + "118e773f", + // Segment of region 1 at offset 54: region id, entry count, then // (entry id, payload length, payload) per entry. "0000000100000001", "00000003", @@ -992,7 +1023,7 @@ mod tests { "0000000000000003", "00000003", "636363", - // Segment of region 2 at offset 100. + // Segment of region 2 at offset 108. "0000000200000001", "00000002", "000000000000000a", @@ -1001,28 +1032,28 @@ mod tests { "000000000000000b", "00000002", "7979", - // Footer at offset 139: entry count, then per segment region id, min and + // Footer at offset 147: entry count, then per segment region id, min and // max entry id, entry count, segment offset, length and CRC32. "00000002", "0000000100000001", "0000000000000001", "0000000000000003", "00000003", - "000000000000002e", + "0000000000000036", "0000000000000036", "25ac0486", "0000000200000001", "000000000000000a", "000000000000000b", "00000002", - "0000000000000064", + "000000000000006c", "0000000000000027", "65dc08ec", // Trailer: footer offset, footer length, footer CRC32, object CRC32, magic. - "000000000000008b", + "0000000000000093", "0000000000000064", - "6b97e3e4", - "631bb7d1", + "d80c9f07", + "65c91839", "475457414c54524c", ); @@ -1060,6 +1091,7 @@ mod tests { object_seq: 6, epoch: 2, }), + reclaim_boundary: 5, }; let footer = vec![ FooterEntry { @@ -1067,7 +1099,7 @@ mod tests { min_entry_id: 1, max_entry_id: 3, entry_count: 3, - segment_offset: 46, + segment_offset: 54, segment_len: 54, segment_crc32: 0x25ac0486, }, @@ -1076,12 +1108,12 @@ mod tests { min_entry_id: 10, max_entry_id: 11, entry_count: 2, - segment_offset: 100, + segment_offset: 108, segment_len: 39, segment_crc32: 0x65dc08ec, }, ]; - assert_eq!(271, fixture.len()); + assert_eq!(279, fixture.len()); let decoded = decode_object(&fixture).unwrap(); assert_eq!(header, decoded.header); @@ -1089,10 +1121,10 @@ mod tests { assert_eq!(fixture_records(), decoded.records); assert_eq!( FixedTrailer { - footer_offset: 139, + footer_offset: 147, footer_len: 100, - footer_crc32: 0x6b97e3e4, - object_crc32: 0x631bb7d1, + footer_crc32: 0xd80c9f07, + object_crc32: 0x65c91839, }, decode_trailer(&fixture[fixture.len() - TRAILER_LEN..]).unwrap() ); diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index 0af6ac70278..4fc7e971535 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -138,6 +138,16 @@ impl ObjectStoreIo { }) } + /// Deletes the object `object_seq`. Deleting an object that does not + /// exist succeeds. + pub(crate) async fn delete(&self, object_seq: u64) -> Result<()> { + let path = self.object_path(object_seq); + self.store.delete(&path).await.context(WalObjectStoreSnafu { + operation: "delete", + path, + }) + } + /// Lists the objects under the prefix, ordered by object sequence. Keys /// that do not follow the object layout are ignored. pub(crate) async fn list(&self) -> Result> { @@ -357,6 +367,11 @@ mod tests { vec![1, 1, 2], objects.iter().map(|object| object.size).collect::>() ); + + // A delete of an object that is already gone succeeds. + io.delete(2).await.unwrap(); + io.delete(2).await.unwrap(); + assert_eq!(vec![1, 10], object_seqs(io.list().await.unwrap())); } #[tokio::test] diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index 8750a61cc4b..dd3c3958336 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -14,6 +14,7 @@ //! Object store WAL construction, recovery, writes and region reads. +use std::cmp::Reverse; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; use std::fmt; use std::ops::Range; @@ -23,7 +24,7 @@ use std::time::{Duration, Instant}; use async_stream::try_stream; use bytes::Bytes; -use common_telemetry::{info, warn}; +use common_telemetry::{debug, info, warn}; use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; use futures::future::BoxFuture; use futures::stream::FuturesUnordered; @@ -34,9 +35,7 @@ use store_api::logstore::entry::{Entry, NaiveEntry}; use store_api::logstore::provider::{ObjectStoreProvider, Provider}; use store_api::logstore::{AppendBatchResponse, EntryId, LogStore, SendableEntryStream, WalIndex}; use store_api::storage::RegionId; -#[cfg(any(test, feature = "testing"))] -use tokio::sync::watch; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, watch}; use tokio::time::MissedTickBehavior; use crate::error::{ @@ -46,7 +45,10 @@ use crate::error::{ StaleWalObjectSnafu, UnconfirmedWalEpochStartSnafu, WalObjectSequenceExhaustedSnafu, WalObjectSequenceUnsettledSnafu, }; -use crate::metrics::METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL; +use crate::metrics::{ + METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, +}; use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ @@ -61,6 +63,7 @@ const COMMAND_BUFFER: usize = 1024; const APPEND_BUFFER: usize = 16; const MIN_FLUSH_INTERVAL: Duration = Duration::from_millis(10); const MAX_IN_FLIGHT_CREATES: usize = 4; +const MAX_IN_FLIGHT_DELETES: usize = 4; /// Delay before a create that failed transiently is attempted again in the /// `enqueued` acknowledgement mode, where no caller is left to retry it. const CREATE_RETRY_DELAY: Duration = Duration::from_millis(100); @@ -89,7 +92,8 @@ const RECOVERY_TAIL_WINDOW: usize = 64 * 1024; /// `enqueued` mode it returns on admission and the object is created in the /// background. While [`MAX_SEALED_BATCHES`] batches /// wait to become durable no append is admitted. A segment that a read cannot -/// decode is handled as `on_corrupted_segment` says. +/// decode is handled as `on_corrupted_segment` says. Objects below the reclaim +/// boundary of the last indexed object are deleted in the background. pub struct ObjectStoreLogStore { prefix: String, ack_mode: AckMode, @@ -99,6 +103,8 @@ pub struct ObjectStoreLogStore { obsolete_entry_ids: ObsoleteEntryIds, /// Segments that reads skipped, per region and by object sequence. wal_holes: WalHoles, + /// Objects with a delete in flight, which a read consults. + deleting: Arc, terminal_error: TerminalError, stopped: Arc, command_tx: mpsc::Sender, @@ -135,6 +141,72 @@ pub struct WalHole { type TerminalError = Arc>>>; +/// The delete attempts in flight, one per object. +/// +/// The actor registers an attempt before its delete starts, unindexes the +/// object if the delete succeeded, and only then settles the attempt and takes +/// it out, so an object without an attempt was either never collected or is +/// already unindexed. [`is_collected`] depends on that order. +#[derive(Debug, Default)] +struct DeletingObjects { + /// The settlement flag of the attempt in flight per object sequence. + attempts: Mutex>>, +} + +impl DeletingObjects { + fn lock(&self) -> std::sync::MutexGuard<'_, BTreeMap>> { + self.attempts.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn start(&self, object_seq: u64) { + self.lock().insert(object_seq, watch::channel(false).0); + } + + fn contains(&self, object_seq: u64) -> bool { + self.lock().contains_key(&object_seq) + } + + /// Settles the attempt of `object_seq` and takes it out, releasing the + /// reads that wait for it. + fn finish(&self, object_seq: u64) { + if let Some(settled) = self.lock().remove(&object_seq) { + settled.send_replace(true); + } + } + + /// Settles every attempt the actor will never complete because it exited: + /// their objects stay indexed, so waiting reads report their error. + fn abandon_all(&self) { + for (_, settled) in std::mem::take(&mut *self.lock()) { + settled.send_replace(true); + } + } + + /// Returns the settlement of the attempt in flight for `object_seq`, or + /// `None` when no delete of it is in flight. + fn attempt(&self, object_seq: u64) -> Option> { + self.lock().get(&object_seq).map(watch::Sender::subscribe) + } + + /// Waits until no delete is in flight. + #[cfg(any(test, feature = "testing"))] + async fn wait_until_empty(&self) { + loop { + let in_flight = self + .lock() + .values() + .map(watch::Sender::subscribe) + .collect::>(); + if in_flight.is_empty() { + return; + } + for mut settled in in_flight { + let _ = settled.wait_for(|settled| *settled).await; + } + } + } +} + impl fmt::Debug for ObjectStoreLogStore { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ObjectStoreLogStore") @@ -193,6 +265,8 @@ impl ObjectStoreLogStore { durable_entry_ids, tip, max_epoch, + reclaim_boundary, + unindexed, } = recover(io.as_ref()).await?; // Every epoch is one above the sequence of the start object its // instance created, and every object is at or above its start object, @@ -205,7 +279,8 @@ impl ObjectStoreLogStore { ), } ); - let start = start_epoch(io.as_ref(), next_object_seq, tip).await?; + let (start, moved_past) = + start_epoch(io.as_ref(), next_object_seq, tip, reclaim_boundary).await?; let epoch = start.epoch; // The epoch identifies this instance in every object it writes. info!( @@ -217,8 +292,12 @@ impl ObjectStoreLogStore { .with_context(|_| InvalidWalObjectSnafu { path: io.object_path(start.object_seq), })?; + // No watermark is recorded yet. + let writer_reclaim_boundary = catalog.reclaim_boundary(reclaim_boundary, &HashMap::new()); let catalog = Arc::new(RwLock::new(catalog)); let obsolete_entry_ids = ObsoleteEntryIds::default(); + let wal_holes = WalHoles::default(); + let deleting = Arc::new(DeletingObjects::default()); let terminal_error = TerminalError::default(); let stopped = Arc::new(AtomicBool::new(false)); let (command_tx, command_rx) = mpsc::channel(COMMAND_BUFFER); @@ -240,6 +319,7 @@ impl ObjectStoreLogStore { io: io.clone(), catalog: catalog.clone(), obsolete_entry_ids: obsolete_entry_ids.clone(), + wal_holes: wal_holes.clone(), terminal_error: terminal_error.clone(), stopped: stopped.clone(), command_rx, @@ -262,6 +342,13 @@ impl ObjectStoreLogStore { .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT), epoch, last_indexed: start, + reclaim_boundary: writer_reclaim_boundary, + durable_reclaim_boundary: reclaim_boundary, + unindexed: unindexed.into_iter().chain(moved_past).collect(), + creating: HashSet::new(), + deleting: deleting.clone(), + deletes: FuturesUnordered::new(), + failed_deletes: HashSet::new(), flush_interval: config.flush_interval, #[cfg(any(test, feature = "testing"))] admitted_appends: admitted_appends_tx, @@ -284,7 +371,8 @@ impl ObjectStoreLogStore { io, catalog, obsolete_entry_ids, - wal_holes: WalHoles::default(), + wal_holes, + deleting, terminal_error, stopped, command_tx, @@ -509,6 +597,14 @@ impl ObjectStoreLogStore { } } + /// Waits until no object delete is in flight. The deletes a collection + /// starts are in flight by the time the `obsolete`, `obsolete_all`, + /// [`seal_open_batch`](Self::seal_open_batch) or `durable` append that + /// started it returns. + pub async fn wait_for_garbage_collection(&self) { + self.deleting.wait_until_empty().await; + } + /// Sets the stopped flag without sending the stop command, which is the /// state a store is in between the two steps of [`stop`](LogStore::stop). pub fn begin_stop(&self) { @@ -595,7 +691,8 @@ impl LogStore for ObjectStoreLogStore { /// Reads the provider's region from `entry_id`, hiding obsolete entries. /// The catalog locates each segment, so a caller-supplied index is unnecessary. - /// A segment that does not decode fails the read or is skipped and recorded + /// An object collected after the read listed it is skipped. Otherwise a + /// segment that does not decode fails the read or is skipped and recorded /// as a hole, as `on_corrupted_segment` says; a failed fetch always fails /// the read. async fn read( @@ -631,6 +728,8 @@ impl LogStore for ObjectStoreLogStore { }; let io = self.io.clone(); + let catalog = self.catalog.clone(); + let deleting = self.deleting.clone(); let provider = provider.clone(); let on_corrupted_segment = self.on_corrupted_segment; let wal_holes = self.wal_holes.clone(); @@ -638,10 +737,18 @@ impl LogStore for ObjectStoreLogStore { for (object_seq, footer_entry) in objects { let records = match fetch_segment(io.as_ref(), object_seq, &footer_entry).await { Ok(records) => records, + Err(error) if is_collected(&catalog, &deleting, object_seq).await => { + debug!( + "Skipped WAL object {} (sequence {}), collected after the read listed it: {error}", + io.object_path(object_seq), + object_seq + ); + continue; + } Err(error @ Error::InvalidWalObject { .. }) if on_corrupted_segment == CorruptedSegmentAction::Skip => { - record_skipped_segment(&wal_holes, io.as_ref(), object_seq, &footer_entry, &error); + record_skipped_segment(&catalog, &wal_holes, io.as_ref(), object_seq, &footer_entry, &error); continue; } Err(error) => Err(error)?, @@ -742,11 +849,22 @@ impl LogStore for ObjectStoreLogStore { } } - /// Hides every entry of the region in memory. The caller re-establishes - /// the watermark after a restart. + /// Hides every entry of the region in memory and lets collection delete + /// its objects. The caller re-establishes the watermark after a restart. async fn obsolete_all(&self, provider: &Provider, region_id: RegionId) -> Result<()> { self.check_region(provider, region_id)?; record_obsolete(&self.obsolete_entry_ids, region_id, EntryId::MAX); + let (response_tx, response_rx) = oneshot::channel(); + if self + .command_tx + .send(Command::Collect { + response: response_tx, + }) + .await + .is_ok() + { + let _ = response_rx.await; + } Ok(()) } @@ -813,6 +931,8 @@ enum Command { watermark: EntryId, response: oneshot::Sender>, }, + /// Answered once the collection the recorded watermarks allow started. + Collect { response: oneshot::Sender<()> }, Stop { response: oneshot::Sender>, }, @@ -850,6 +970,7 @@ enum CreateState { /// created, indexed and acknowledged. struct SealedBatch { object_seq: u64, + reclaim_boundary: u64, bytes: Bytes, footer: Vec, first_admitted_at: Instant, @@ -910,10 +1031,22 @@ type CreateOutcome = (u64, Result); /// [`MAX_SEALED_BATCHES`] batches wait or an append is held back at a backlog /// threshold of the `enqueued` mode, so commands such as `stop` are handled /// while admission is held back. +/// +/// The actor also collects garbage. Every batch carries the reclaim boundary +/// current when it was sealed, see [`ObjectCatalog::reclaim_boundary`], which +/// is recomputed after every `obsolete` and every indexed object and never +/// moves down. The boundary of the last indexed object is durable: every +/// present object below it, indexed or not, is deleted in the background, up +/// to [`MAX_IN_FLIGHT_DELETES`] at a time, after an `obsolete` or an indexing +/// that moved it. An indexed object is unindexed once its delete succeeded; a +/// failed delete is attempted again at the next collection. Nothing is +/// collected once stop began or the store is poisoned, and deletes in flight +/// complete before `stop` returns. struct Actor { io: Arc, catalog: Arc>, obsolete_entry_ids: ObsoleteEntryIds, + wal_holes: WalHoles, terminal_error: TerminalError, stopped: Arc, command_rx: mpsc::Receiver, @@ -950,6 +1083,19 @@ struct Actor { /// The last indexed object, which is the start object until a batch is /// indexed. A batch sealed while no earlier batch waits extends it. last_indexed: ChainLink, + /// The reclaim boundary the next sealed batch carries. + reclaim_boundary: u64, + /// The reclaim boundary of the last indexed object. + durable_reclaim_boundary: u64, + /// Sequences of objects that may be present but are not indexed: the + /// objects recovery listed off the chain and the batches that failed. + unindexed: BTreeSet, + /// Sequences whose create is in flight. + creating: HashSet, + deleting: Arc, + deletes: FuturesUnordered)>>, + /// Deletes that failed since the last collection started. + failed_deletes: HashSet, flush_interval: Duration, #[cfg(any(test, feature = "testing"))] admitted_appends: watch::Sender, @@ -984,6 +1130,9 @@ impl Actor { Some((object_seq, result)) = self.creates.next(), if !self.creates.is_empty() => { self.on_create_completed(object_seq, result); } + Some((object_seq, result)) = self.deletes.next(), if !self.deletes.is_empty() => { + self.on_delete_completed(object_seq, result); + } Some((entries, response)) = self.append_rx.recv(), if self.stalled.is_none() && self.sealed.len() + 2 <= MAX_SEALED_BATCHES => { self.handle_append(entries, response); } @@ -994,6 +1143,11 @@ impl Actor { Some(Command::Obsolete { region_id, entry_id, watermark, response }) => { self.handle_obsolete(region_id, entry_id, watermark, response); } + Some(Command::Collect { response }) => { + self.update_reclaim_boundary(); + self.collect(); + let _ = response.send(()); + } Some(Command::Stop { response }) => { self.handle_stop(response); } @@ -1178,6 +1332,7 @@ impl Actor { epoch: self.epoch, }), ), + reclaim_boundary: self.reclaim_boundary, }; let encoded = match encode_batch(header, entries) { Ok(encoded) => encoded, @@ -1202,6 +1357,7 @@ impl Actor { } self.sealed.push_back(SealedBatch { object_seq, + reclaim_boundary: self.reclaim_boundary, bytes: encoded.bytes, footer: encoded.footer, first_admitted_at, @@ -1241,6 +1397,7 @@ impl Actor { batch.attempts += 1; let io = self.io.clone(); let object_seq = batch.object_seq; + self.creating.insert(object_seq); let bytes = batch.bytes.clone(); let epoch = self.epoch; let ack_mode = self.ack_mode; @@ -1290,14 +1447,17 @@ impl Actor { } fn on_create_completed(&mut self, object_seq: u64, result: Result) { + self.creating.remove(&object_seq); // A batch that already failed, or that the store gave up on when it - // poisoned itself: the object may exist, but it is off the chain. + // poisoned itself: the object may exist, but it is off the chain, and + // collection may delete it now. let Some(index) = self .sealed .iter() .position(|batch| batch.object_seq == object_seq) else { self.start_creates(); + self.start_deletes(); return; }; match result { @@ -1368,6 +1528,13 @@ impl Actor { object_seq: batch.object_seq, epoch: self.epoch, }; + // The collection starts before the waiters learn that the object is + // indexed. + self.update_reclaim_boundary(); + if batch.reclaim_boundary > self.durable_reclaim_boundary { + self.durable_reclaim_boundary = batch.reclaim_boundary; + self.collect(); + } for waiter in batch.waiters { let _ = waiter.response.send(Ok(AppendBatchResponse { last_entry_ids: waiter.last_entry_ids, @@ -1434,6 +1601,7 @@ impl Actor { } }; for batch in self.sealed.drain(..) { + self.unindexed.insert(batch.object_seq); batch.fail(failure); } self.reset_open_batch(); @@ -1546,9 +1714,105 @@ impl Actor { let result = self.raise_sequence_floor(region_id, entry_id).map(|_| { record_obsolete(&self.obsolete_entry_ids, region_id, watermark); }); + if result.is_ok() { + self.update_reclaim_boundary(); + self.collect(); + } let _ = response.send(result); } + /// Raises the reclaim boundary to what the catalog and the obsolete + /// watermarks allow. + fn update_reclaim_boundary(&mut self) { + let boundary = { + let obsolete_entry_ids = self + .obsolete_entry_ids + .lock() + .unwrap_or_else(PoisonError::into_inner); + self.catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .reclaim_boundary(self.reclaim_boundary, &obsolete_entry_ids) + }; + self.reclaim_boundary = boundary; + } + + /// Starts a collection, which also attempts the deletes that failed since + /// the last one again. + fn collect(&mut self) { + self.failed_deletes.clear(); + self.start_deletes(); + } + + /// Starts deleting the objects below the durable reclaim boundary while + /// fewer than [`MAX_IN_FLIGHT_DELETES`] deletes are in flight, skipping + /// objects whose create is in flight or whose delete failed since the + /// collection started. + fn start_deletes(&mut self) { + if self.is_stopped() || terminal(&self.terminal_error).is_some() { + return; + } + let free = MAX_IN_FLIGHT_DELETES.saturating_sub(self.deletes.len()); + let boundary = self.durable_reclaim_boundary; + let candidates = { + let catalog = self.catalog.read().unwrap_or_else(PoisonError::into_inner); + catalog + .objects_in_order() + .map(|(object_seq, _)| object_seq) + .take_while(|object_seq| *object_seq < boundary) + .chain(self.unindexed.range(..boundary).copied()) + .filter(|object_seq| { + !self.deleting.contains(*object_seq) + && !self.failed_deletes.contains(object_seq) + && !self.creating.contains(object_seq) + }) + .take(free) + .collect::>() + }; + for object_seq in candidates { + self.deleting.start(object_seq); + let io = self.io.clone(); + self.deletes.push(Box::pin(async move { + (object_seq, io.delete(object_seq).await) + })); + } + } + + /// Unindexes the object `object_seq` and drops its holes once its delete + /// succeeded, then settles its attempt, so a read that the attempt + /// releases sees the outcome in the catalog. + fn on_delete_completed(&mut self, object_seq: u64, result: Result<()>) { + match result { + Ok(()) => { + self.unindexed.remove(&object_seq); + let mut catalog = self.catalog.write().unwrap_or_else(PoisonError::into_inner); + let footer = catalog.remove_object(object_seq); + let mut wal_holes = self + .wal_holes + .lock() + .unwrap_or_else(PoisonError::into_inner); + for entry in footer.iter().flatten() { + if let Some(holes) = wal_holes.get_mut(&entry.region_id) { + holes.remove(&object_seq); + } + } + METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.inc(); + } + Err(error) => { + warn!( + error; + "Failed to delete WAL object {} (sequence {}), it is deleted again at the next collection", + self.io.object_path(object_seq), + object_seq + ); + METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.inc(); + self.failed_deletes.insert(object_seq); + } + } + self.start_deletes(); + self.deleting.finish(object_seq); + } + /// Moves the sequence of the next object above the object that holds /// `entry_id` unless it is there already. The move fails while the open /// batch has handed out ids under the next sequence. A floor that does @@ -1633,9 +1897,14 @@ impl Actor { } /// Answers the callers of `stop` once every sealed batch is settled and - /// every create has completed. Returns true when the actor is done. + /// every create and delete has completed. Returns true when the actor is + /// done. fn finish_stop(&mut self) -> bool { - if self.stop.is_empty() || !self.sealed.is_empty() || !self.creates.is_empty() { + if self.stop.is_empty() + || !self.sealed.is_empty() + || !self.creates.is_empty() + || !self.deletes.is_empty() + { return false; } for waiter in self.durable_waiters.drain(..) { @@ -1686,6 +1955,14 @@ impl Actor { } } +impl Drop for Actor { + /// Settles the deletes the actor still had in flight, so no read waits + /// for an attempt that will never complete. + fn drop(&mut self) { + self.deleting.abandon_all(); + } +} + fn terminal(terminal_error: &TerminalError) -> Option> { terminal_error .lock() @@ -1715,6 +1992,25 @@ fn shared(error: &Arc) -> Error { ObjectStoreWalSnafu.into_error(error.clone()) } +/// Returns whether collection deleted the object `object_seq` after a read +/// listed it. The set of deleting objects is observed before the catalog, the +/// reverse of the order the actor updates them, so an object that no longer +/// looks deleting but still looks indexed was never collected. An undecided +/// delete is waited for. +async fn is_collected( + catalog: &RwLock, + deleting: &DeletingObjects, + object_seq: u64, +) -> bool { + if let Some(mut settled) = deleting.attempt(object_seq) { + let _ = settled.wait_for(|settled| *settled).await; + } + !catalog + .read() + .unwrap_or_else(PoisonError::into_inner) + .contains_object(object_seq) +} + /// Fetches and decodes the segment `entry` describes in the object /// `object_seq`. A segment that does not decode is fetched once more, since /// the download rather than the object may be damaged; when it still does not @@ -1735,14 +2031,27 @@ async fn fetch_segment( } /// Records the segment `entry` describes in the object `object_seq` as a hole -/// of its region, once per object, counts the skip and warns about it. +/// of its region, once per object, counts the skip and warns about it. An +/// object collected meanwhile is skipped without a hole: the catalog is held +/// while the hole is recorded, as a collection holds it while it drops the +/// holes of the object it removes. fn record_skipped_segment( + catalog: &RwLock, wal_holes: &WalHoles, io: &dyn WalObjectIo, object_seq: u64, entry: &FooterEntry, error: &Error, ) { + let catalog = catalog.read().unwrap_or_else(PoisonError::into_inner); + if !catalog.contains_object(object_seq) { + debug!( + "Skipped WAL object {} (sequence {}), collected after the read listed it: {error}", + io.object_path(object_seq), + object_seq + ); + return; + } let hole = WalHole { path: io.object_path(object_seq), object_seq, @@ -1791,6 +2100,19 @@ struct Recovered { tip: Option, /// Largest epoch any present object carries, zero on an empty prefix. max_epoch: u64, + /// The reclaim boundary of the tip, zero on an empty prefix. + reclaim_boundary: u64, + /// Present objects that are not indexed, in sequence order. + unindexed: Vec, +} + +/// The objects recovery fetched, ordered by sequence, and the listed objects +/// that were gone when their header or footer was fetched, with the error the +/// fetch met. A collection of an earlier instance may complete after the +/// listing. +struct FetchedObjects { + objects: Vec, + vanished: Vec<(u64, Error)>, } /// An object as recovery fetched it: its key, header and footer. @@ -1812,18 +2134,22 @@ async fn recover(io: &dyn WalObjectIo) -> Result { finish_recovery(fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?) } -/// Indexes the chain of `objects`, which are ordered by sequence. Objects off -/// the chain are orphans: they are not indexed, but no later object takes -/// their sequences. -fn finish_recovery(objects: Vec) -> Result { +/// Indexes the chain of the fetched objects. Objects off the chain, orphans +/// and objects below the reclaim boundary of the tip, are not indexed, but no +/// later object takes their sequences, nor the sequences of vanished objects. +/// A vanished object at or above the boundary of the tip fails recovery with +/// the error its fetch met: collection never deletes such an object. +fn finish_recovery(fetched: FetchedObjects) -> Result { + let FetchedObjects { objects, vanished } = fetched; let headers = objects .iter() .map(|fetched| (fetched.object.object_seq, &fetched.header)) .collect::>(); let chain = select_chain(&headers); // Every open writes an object that starts a chain or extends a complete - // one, and nothing removes objects, so present objects without any - // complete chain mean objects of the chain are gone. + // one, and collection removes only objects below the reclaim boundary of + // a later object, so present objects without any complete chain mean + // objects of the chain are gone. ensure!( headers.is_empty() || !chain.is_empty(), CorruptedWalObjectSnafu { @@ -1842,20 +2168,35 @@ fn finish_recovery(objects: Vec) -> Result { .map(|header| header.epoch) .max() .unwrap_or(0); - let after_listed = match headers.last_key_value() { + let last_listed = headers + .keys() + .copied() + .chain(vanished.iter().map(|(object_seq, _)| *object_seq)) + .max(); + let after_listed = match last_listed { None => 0, - Some((&last_object_seq, _)) => last_object_seq + Some(last_object_seq) => last_object_seq .checked_add(1) .context(WalObjectSequenceExhaustedSnafu { last_object_seq })?, }; + let reclaim_boundary = tip.map_or(0, |tip| headers[&tip.object_seq].reclaim_boundary); + if let Some((_, error)) = vanished + .into_iter() + .find(|(object_seq, _)| *object_seq >= reclaim_boundary) + { + return Err(error); + } let chain = chain.into_iter().collect::>(); let mut catalog = ObjectCatalog::default(); + let mut unindexed = Vec::new(); for FetchedObject { object, footer, .. } in objects { if chain.contains(&object.object_seq) { catalog .insert_object(object.object_seq, footer) .with_context(|_| InvalidWalObjectSnafu { path: object.path })?; + } else { + unindexed.push(object.object_seq); } } let next_object_seq = catalog.next_object_seq()?.max(after_listed); @@ -1872,15 +2213,20 @@ fn finish_recovery(objects: Vec) -> Result { durable_entry_ids, tip, max_epoch, + reclaim_boundary, + unindexed, }) } /// Returns, in sequence order, the objects on the chain that ends at the tip: /// the complete object with the largest epoch, then the largest sequence. /// -/// An object is complete when every link on its chain holds: the chain starts -/// at an object without a predecessor, and every other link names a present -/// object that carries the epoch the link records. Only one instance writes +/// An object is complete when every link on its chain holds, judged with the +/// reclaim boundary of that object: the chain starts at an object without a +/// predecessor or reaches a predecessor below the boundary, whether or not it +/// is present, and every other link names a present object that carries the +/// epoch the link records. The chain of the tip ends there, so objects below +/// its boundary are not replayed. Only one instance writes /// under an epoch, so the epoch tells an object of the linking instance from /// another object under the same sequence. A create that /// was reported as failed may still leave its object, but every object @@ -1888,43 +2234,43 @@ fn finish_recovery(objects: Vec) -> Result { /// an epoch above every object present when it opened, so neither such an /// object nor a late object of an earlier instance ends the chosen chain. fn select_chain(headers: &BTreeMap) -> Vec { - // A predecessor precedes its successor, so one pass in sequence order - // settles every object. - let mut complete = HashMap::with_capacity(headers.len()); - for (&object_seq, header) in headers { - let holds = match header.prev { - None => true, - Some(link) if link.object_seq >= object_seq => false, - Some(link) => match headers.get(&link.object_seq) { - Some(prev) => prev.epoch == link.epoch && complete[&link.object_seq], - None => false, - }, - }; - complete.insert(object_seq, holds); - } - let Some(tip) = headers - .iter() - .filter(|(object_seq, _)| complete[*object_seq]) - .max_by_key(|(object_seq, header)| (header.epoch, **object_seq)) - .map(|(object_seq, _)| *object_seq) - else { - return Vec::new(); - }; + let mut candidates = headers.iter().collect::>(); + candidates.sort_unstable_by_key(|(object_seq, header)| Reverse((header.epoch, **object_seq))); + candidates + .into_iter() + .find_map(|(object_seq, _)| chain_ending_at(headers, *object_seq)) + .unwrap_or_default() +} + +/// Returns, in sequence order, the chain that ends at `tip` if `tip` is +/// complete, see [`select_chain`]. +fn chain_ending_at(headers: &BTreeMap, tip: u64) -> Option> { + let reclaim_boundary = headers[&tip].reclaim_boundary; let mut chain = vec![tip]; - while let Some(link) = headers[chain.last().expect("the chain holds the tip")].prev - && headers.contains_key(&link.object_seq) - { + let mut header = headers[&tip]; + while let Some(link) = header.prev { + if link.object_seq >= header.object_seq { + return None; + } + if link.object_seq < reclaim_boundary { + break; + } + header = headers + .get(&link.object_seq) + .filter(|prev| prev.epoch == link.epoch)?; chain.push(link.object_seq); } chain.reverse(); - chain + Some(chain) } /// Writes the empty object that starts the epoch of this instance at -/// `object_seq`, linked to the recovered `tip`, and returns it as the tip -/// later objects extend. The epoch is one above the sequence the create -/// claims, so no two instances share one however their opens interleave, and -/// it is above the epoch of every object recovery listed. +/// `object_seq`, linked to the recovered `tip` and carrying the recovered +/// `reclaim_boundary`, and returns it as the tip later objects extend, +/// together with the sequences of earlier-epoch objects it moved past. The +/// epoch is one above the sequence the create claims, so no two instances +/// share one however their opens interleave, and it is above the epoch of +/// every object recovery listed. /// /// An object an earlier instance left at that sequence after recovery listed /// the prefix has a lower epoch and never ends a chain, so the start object @@ -1938,7 +2284,9 @@ async fn start_epoch( io: &dyn WalObjectIo, mut object_seq: u64, tip: Option, -) -> Result { + reclaim_boundary: u64, +) -> Result<(ChainLink, Vec)> { + let mut moved_past = Vec::new(); loop { ensure!( object_seq < OBJECT_SEQ_LIMIT, @@ -1951,12 +2299,13 @@ async fn start_epoch( object_seq, epoch, prev: tip, + reclaim_boundary, }; match io .put_if_absent(object_seq, encode_object(header, &[])?.bytes) .await { - Ok(PutResult::Created) => return Ok(ChainLink { object_seq, epoch }), + Ok(PutResult::Created) => return Ok((ChainLink { object_seq, epoch }, moved_past)), Ok(PutResult::AlreadyPresent) => { return UnconfirmedWalEpochStartSnafu { path: io.object_path(object_seq), @@ -1968,6 +2317,7 @@ async fn start_epoch( if epoch_of(io, object_seq).await? >= epoch { return Err(error); } + moved_past.push(object_seq); object_seq += 1; } Err(error) => return Err(error), @@ -2007,27 +2357,43 @@ async fn stale_conflict( /// Fetches and verifies the headers and footers of `objects`, up to /// `concurrency` objects at a time, and returns them ordered by object sequence -/// whatever the order the fetches complete in. The first failure abandons the +/// whatever the order the fetches complete in. An object the object store no +/// longer has is returned as vanished; any other failure abandons the /// remaining fetches. async fn fetch_footers( io: &dyn WalObjectIo, objects: Vec, concurrency: usize, -) -> Result> { - let mut fetched = futures::stream::iter(objects) +) -> Result { + let outcomes = futures::stream::iter(objects) .map(|object| async move { - let (header, footer) = fetch_footer(io, &object).await?; - Ok(FetchedObject { - object, - header, - footer, - }) + match fetch_footer(io, &object).await { + Ok((header, footer)) => Ok(Ok(FetchedObject { + object, + header, + footer, + })), + Err(error) if is_not_found(&error) => Ok(Err((object.object_seq, error))), + Err(error) => Err(error), + } }) .buffer_unordered(concurrency) .try_collect::>() .await?; - fetched.sort_unstable_by_key(|fetched| fetched.object.object_seq); - Ok(fetched) + let mut objects = Vec::with_capacity(outcomes.len()); + let mut vanished = Vec::new(); + for outcome in outcomes { + match outcome { + Ok(object) => objects.push(object), + Err(gone) => vanished.push(gone), + } + } + objects.sort_unstable_by_key(|fetched| fetched.object.object_seq); + Ok(FetchedObjects { objects, vanished }) +} + +fn is_not_found(error: &Error) -> bool { + matches!(error, Error::WalObjectStore { error, .. } if error.kind() == object_store::ErrorKind::NotFound) } /// Reads the header, trailer and footer of `object` and verifies them: the @@ -2148,6 +2514,8 @@ pub(crate) trait WalObjectIo: Send + Sync { async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result; + async fn delete(&self, object_seq: u64) -> Result<()>; + async fn list(&self) -> Result>; fn object_path(&self, object_seq: u64) -> String; @@ -2167,6 +2535,10 @@ impl WalObjectIo for ObjectStoreIo { ObjectStoreIo::get_range(self, object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + ObjectStoreIo::delete(self, object_seq).await + } + async fn list(&self) -> Result> { ObjectStoreIo::list(self).await } @@ -2182,6 +2554,7 @@ mod tests { use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, RetryHint}; + use futures::FutureExt; use object_store::services::Memory; use store_api::logstore::entry::{MultiplePartEntry, MultiplePartHeader}; use tokio::time::timeout; @@ -2197,6 +2570,9 @@ mod tests { /// counter, so that tests running in one process do not interleave /// their increments. static SKIPPED_SEGMENTS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + /// Held by every test that fails a delete, so that a test reading the + /// process-wide failed deletes counter sees only its own failures. + static FAILED_DELETES: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); fn memory_store() -> ObjectStore { ObjectStore::new(Memory::default()).unwrap() @@ -2318,6 +2694,14 @@ mod tests { .collect() } + /// Lists the objects once no delete is in flight. + async fn collected_object_seqs(store: &ObjectStoreLogStore) -> Vec { + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + object_seqs(store.io.as_ref()).await + } + fn unwrap_shared(error: &Error) -> &Error { match error { Error::ObjectStoreWal { source, .. } => source, @@ -2332,22 +2716,22 @@ mod tests { ); } - /// Waits for the next create that `ParkedIo` parked. - async fn next_create( + /// Waits for the next operation that `ParkedIo` parked. + async fn next_parked_operation( parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, ) -> (u64, oneshot::Sender) { timeout(WAIT, parked.recv()).await.unwrap().unwrap() } - /// Waits for `count` parked creates and returns their releases by + /// Waits for `count` parked operations and returns their releases by /// sequence. - async fn parked_creates( + async fn parked_operations( parked: &mut mpsc::UnboundedReceiver<(u64, oneshot::Sender)>, count: usize, ) -> HashMap> { let mut releases = HashMap::new(); for _ in 0..count { - let (object_seq, release) = next_create(parked).await; + let (object_seq, release) = next_parked_operation(parked).await; releases.insert(object_seq, release); } releases @@ -2496,7 +2880,10 @@ mod tests { footer: decoded.footer, }); } - finish_recovery(objects) + finish_recovery(FetchedObjects { + objects, + vanished: Vec::new(), + }) } fn assert_same_recovery(expected: &Recovered, actual: &Recovered) { @@ -2508,6 +2895,8 @@ mod tests { assert_eq!(expected.durable_entry_ids, actual.durable_entry_ids); assert_eq!(expected.tip, actual.tip); assert_eq!(expected.max_epoch, actual.max_epoch); + assert_eq!(expected.reclaim_boundary, actual.reclaim_boundary); + assert_eq!(expected.unindexed, actual.unindexed); } fn catalog_contents(catalog: &ObjectCatalog) -> Vec<(u64, Vec)> { @@ -2559,6 +2948,7 @@ mod tests { object_seq: prev.object_seq, epoch: prev.epoch, }), + reclaim_boundary: 0, } } @@ -2581,6 +2971,7 @@ mod tests { object_seq, epoch, prev: None, + reclaim_boundary: 0, }; put_object_with_header(&io, header, &[]).await; } @@ -3017,6 +3408,7 @@ mod tests { catalog: Arc::default(), obsolete_entry_ids: ObsoleteEntryIds::default(), wal_holes: WalHoles::default(), + deleting: Arc::default(), terminal_error: Arc::default(), stopped: Arc::new(AtomicBool::new(false)), command_tx, @@ -3105,6 +3497,7 @@ mod tests { fetch_footers(io.as_ref(), objects, RECOVERY_CONCURRENCY) .await .unwrap() + .objects .into_iter() .map(|fetched| fetched.object.object_seq) .collect::>() @@ -3663,7 +4056,7 @@ mod tests { HashMap::from([(region_one, id(2, 1)), (region_two, id(2, 1))]), third.await.unwrap().unwrap().last_entry_ids ); - assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); assert_eq!( expected_entries( region_one, @@ -3704,7 +4097,7 @@ mod tests { HashMap::from([(region_id, id(1, POSITION_LIMIT - 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); // An append that alone runs past the range fits no object: it seals // the open batch like any append the batch cannot take, then it is @@ -3730,7 +4123,7 @@ mod tests { HashMap::from([(region_id, id(3, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1, 2, 3], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); assert_eq!( expected_entries(region_id, &[(id(2, 1), "next"), (id(3, 1), "after")]), read_entries(&store, region_id, id(2, 1)).await @@ -3753,16 +4146,16 @@ mod tests { .await .unwrap() .unwrap(); - let expected_seqs = (0..=index as u64 + 1).collect::>(); - assert_eq!(expected_seqs, object_seqs(store.io.as_ref()).await); + // Object 1 carries the boundary past the start object. + let expected_seqs = (1..=index as u64 + 1).collect::>(); + assert_eq!(expected_seqs, collected_object_seqs(&store).await); } // Ticks with an empty open batch do not create objects. tokio::time::sleep(interval * 5).await; - let seqs = object_seqs(store.io.as_ref()).await; - assert_eq!(vec![0, 1, 2, 3], seqs); - // Object 0 is the start object of the store. - for object_seq in seqs.into_iter().skip(1) { + let seqs = collected_object_seqs(&store).await; + assert_eq!(vec![1, 2, 3], seqs); + for object_seq in seqs { let bytes = store.io.get(object_seq).await.unwrap(); let decoded = decode_object(&bytes).unwrap(); assert_eq!(1, decoded.records.len(), "object {object_seq} is empty"); @@ -3829,7 +4222,7 @@ mod tests { let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 2).await; // At most the limit of creates run at a time; the rest wait for a slot. - let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; assert_eq!( (1..=MAX_IN_FLIGHT_CREATES as u64).collect::>(), releases.keys().copied().collect::>() @@ -3842,7 +4235,7 @@ mod tests { // each, but nothing is acknowledged ahead of object 1. for (released, expected_next) in [(3, 5), (2, 6)] { releases.remove(&released).unwrap().send(true).unwrap(); - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(expected_next, object_seq); releases.insert(object_seq, release); } @@ -3894,7 +4287,7 @@ mod tests { MAX_IN_FLIGHT_CREATES, io.max_in_flight.load(Ordering::SeqCst) ); - assert_eq!(vec![0, 1, 2, 3, 4, 5, 6], object_seqs(io.as_ref()).await); + assert_eq!(vec![1, 2, 3, 4, 5, 6], collected_object_seqs(&store).await); // Every batch extends the batch sealed before it, whichever was // durable first. for object_seq in 1..=6 { @@ -3922,11 +4315,11 @@ mod tests { #[tokio::test] async fn test_store_transient_failure_fails_every_batch_that_is_not_indexed() { - let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); // Every slot is taken; the last batch waits for one. let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; - let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; round_trip_actor(&store).await; assert!(parked.try_recv().is_err()); @@ -3948,14 +4341,14 @@ mod tests { // one retry starts until they complete; what they store is off the // chain. let retries = spawn_appends(&store, region_id, 2).await; - let (object_seq, first) = next_create(&mut parked).await; + let (object_seq, first) = next_parked_operation(&mut parked).await; assert_eq!(6, object_seq); round_trip_actor(&store).await; assert!(parked.try_recv().is_err()); for release in releases.into_values() { release.send(true).unwrap(); } - let (object_seq, second) = next_create(&mut parked).await; + let (object_seq, second) = next_parked_operation(&mut parked).await; assert_eq!(7, object_seq); first.send(true).unwrap(); second.send(true).unwrap(); @@ -3966,7 +4359,7 @@ mod tests { response.last_entry_ids ); } - assert_eq!(vec![0, 2, 3, 4, 6, 7], object_seqs(io.as_ref()).await); + assert_eq!(vec![2, 3, 4, 6, 7], collected_object_seqs(&store).await); assert_eq!( expected_entries(region_id, &[(id(6, 1), "a1"), (id(7, 1), "a2")]), read_entries(&store, region_id, 1).await @@ -3979,7 +4372,7 @@ mod tests { let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); let appends = spawn_appends(&store, region_id, 2).await; - let mut releases = parked_creates(&mut parked, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; // Object 2 is created while object 1 failed: both batches fail, and // the store keeps serving. @@ -4003,7 +4396,7 @@ mod tests { // The retry extends the start object, not the failed batches. let retry = spawn_appends(&store, region_id, 1).await; - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(3, object_seq); release.send(true).unwrap(); let response = timeout(WAIT, retry.into_iter().next().unwrap()) @@ -4045,7 +4438,7 @@ mod tests { // Another writer of the same epoch took sequence 1. put_foreign(&object_store, 1, 1).await; let appends = spawn_appends(&store, region_id, 2).await; - let mut releases = parked_creates(&mut parked, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; // Object 2 is created; object 1 conflicts with the foreign object. releases.remove(&2).unwrap().send(true).unwrap(); @@ -4077,12 +4470,16 @@ mod tests { #[tokio::test] async fn test_store_conflict_with_an_earlier_epoch_fails_the_batch_without_poisoning() { let object_store = memory_store(); - let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); // A late object of an earlier epoch lands under sequence 1. put_foreign(&object_store, 1, 0).await; let append = spawn_appends(&store, region_id, 1).await; - next_create(&mut parked).await.1.send(true).unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); let error = timeout(WAIT, append.into_iter().next().unwrap()) .await .unwrap() @@ -4103,7 +4500,7 @@ mod tests { // The store keeps serving: the retry takes the next sequence. let retry = spawn_appends(&store, region_id, 1).await; - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(2, object_seq); release.send(true).unwrap(); let response = timeout(WAIT, retry.into_iter().next().unwrap()) @@ -4115,7 +4512,7 @@ mod tests { HashMap::from([(region_id, id(2, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); store.stop().await.unwrap(); // The stale object is off the chain after a restart. @@ -4197,7 +4594,8 @@ mod tests { HashMap::from([(region_one, id(2, 2))]), response.last_entry_ids ); - assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + // Object 2 carries the boundary past the watermark of region one. + assert_eq!(vec![2], collected_object_seqs(&store).await); store .obsolete(&provider(region_two), region_two, id(3, 7)) .await @@ -4214,7 +4612,7 @@ mod tests { HashMap::from([(region_two, id(4, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1, 2, 4], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![2, 4], collected_object_seqs(&store).await); // The watermark hides nothing of this prefix; the region has no // entry at or below it. assert_eq!( @@ -4260,11 +4658,12 @@ mod tests { response.last_entry_ids ); } - assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + // Object 3 carries the boundary at object 2, so the failed object 1 + // is collected. + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); store.stop().await.unwrap(); - // Object 2 extends the start object, so recovery leaves object 1 off - // the chain: only the retries replay. + // Only the retries replay. let store = open_over(io, &manual()).await; assert_eq!( expected_entries(region_one, &[(id(2, 1), "a1")]), @@ -4279,11 +4678,11 @@ mod tests { #[tokio::test] async fn test_store_floor_applies_while_a_create_is_in_flight() { - let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_one = region(1); let region_two = region(2); let pending = spawn_append_batch(&store, vec![entry(&store, region_one, "a1")]); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; // Object 1 is in flight and the open batch is empty: the floor moves // the next sequence, which a failure of object 1 does not move back. @@ -4295,13 +4694,17 @@ mod tests { timeout(WAIT, pending).await.unwrap().unwrap().unwrap_err(); let write = spawn_append_batch(&store, vec![entry(&store, region_two, "b1")]); - next_create(&mut parked).await.1.send(true).unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); let response = timeout(WAIT, write).await.unwrap().unwrap().unwrap(); assert_eq!( HashMap::from([(region_two, id(6, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 6], object_seqs(io.as_ref()).await); + assert_eq!(vec![6], collected_object_seqs(&store).await); } #[tokio::test] @@ -4349,7 +4752,7 @@ mod tests { assert_eq!(vec![0], object_seqs(store.io.as_ref()).await); store.release_creates(); timeout(WAIT, held).await.unwrap().unwrap().unwrap(); - assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); store.fail_creates(); let error = append(&store, region_id, "a2").await.unwrap_err(); @@ -4357,7 +4760,7 @@ mod tests { matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), "unexpected error: {error:?}" ); - assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1], object_seqs(store.io.as_ref()).await); store.begin_stop(); assert_stopped(&append(&store, region_id, "a3").await.unwrap_err()); @@ -4386,7 +4789,7 @@ mod tests { HashMap::from([(region_id, id(2, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1, 2], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); store.stop().await.unwrap(); } @@ -4455,7 +4858,7 @@ mod tests { #[tokio::test] async fn test_store_rollback_and_poison_drop_the_open_batch() { let object_store = memory_store(); - let (store, io, mut parked) = open_parking_creates(object_store.clone(), &manual()).await; + let (store, _, mut parked) = open_parking_creates(object_store.clone(), &manual()).await; let region_id = region(1); let spawn_seal = || { let store = store.clone(); @@ -4467,7 +4870,7 @@ mod tests { let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); store.wait_for_admitted_appends(1).await.unwrap(); let seal = spawn_seal(); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; let open = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); store.wait_for_admitted_appends(2).await.unwrap(); release.send(false).unwrap(); @@ -4485,7 +4888,11 @@ mod tests { let retry = spawn_append_batch(&store, vec![entry(&store, region_id, "b1")]); store.wait_for_admitted_appends(3).await.unwrap(); let seal = spawn_seal(); - next_create(&mut parked).await.1.send(true).unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); timeout(WAIT, seal).await.unwrap().unwrap().unwrap(); let response = timeout(WAIT, retry).await.unwrap().unwrap().unwrap(); assert_eq!( @@ -4501,7 +4908,7 @@ mod tests { let sealed = spawn_append_batch(&store, vec![entry(&store, region_id, "c1")]); store.wait_for_admitted_appends(4).await.unwrap(); let seal = spawn_seal(); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; let open = spawn_append_batch(&store, vec![entry(&store, region_id, "c2")]); store.wait_for_admitted_appends(5).await.unwrap(); put_foreign(&object_store, 3, 1).await; @@ -4514,7 +4921,7 @@ mod tests { ); } timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); - assert_eq!(vec![0, 2, 3], object_seqs(io.as_ref()).await); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); store.stop().await.unwrap(); } @@ -4562,7 +4969,7 @@ mod tests { let entries = vec![entry(&store, region_id, &format!("w{index}"))]; appends.push(spawn_append_batch(&store, entries)); } - let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; round_trip_actor(&store).await; assert_eq!(MAX_SEALED_BATCHES - 1, admitted()); @@ -4581,7 +4988,11 @@ mod tests { for release in releases.into_values() { release.send(true).unwrap(); } - next_create(&mut parked).await.1.send(true).unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); timeout(WAIT, stop).await.unwrap().unwrap().unwrap(); let mut acknowledged = 0; for append in appends { @@ -4611,7 +5022,7 @@ mod tests { async fn test_store_stop_during_failed_flush_reports_stopped() { let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; let pending = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; let stop = begin_spawned_stop(&store).await; assert!(!stop.is_finished()); @@ -4629,7 +5040,7 @@ mod tests { let region_id = region(1); // Four creates are in flight, the fifth batch waits for a slot. let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 1).await; - let mut releases = parked_creates(&mut parked, MAX_IN_FLIGHT_CREATES).await; + let mut releases = parked_operations(&mut parked, MAX_IN_FLIGHT_CREATES).await; assert!(parked.try_recv().is_err()); let stop = begin_spawned_stop(&store).await; @@ -4682,7 +5093,7 @@ mod tests { let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); let pending = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; let stop = begin_spawned_stop(&store).await; // Another writer of the same epoch takes the sequence before the // create runs. @@ -4719,6 +5130,7 @@ mod tests { object_seq, epoch, prev: prev.map(|(object_seq, epoch)| ChainLink { object_seq, epoch }), + reclaim_boundary: 0, } } @@ -4802,6 +5214,29 @@ mod tests { ); // A predecessor at or above the object itself never holds. assert!(chain_of(&[chain_header(5, 1, Some((5, 1)))]).is_empty()); + + // A link below the reclaim boundary of the object it is judged for + // holds whether or not the predecessor is present; above it a + // missing predecessor still breaks the link. + let bounded = |object_seq, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, 1, prev) + }; + assert_eq!( + vec![3, 4], + chain_of(&[bounded(3, Some((2, 1)), 3), bounded(4, Some((3, 1)), 3)]) + ); + assert!(chain_of(&[bounded(4, Some((3, 1)), 3)]).is_empty()); + // Object 3 is incomplete with its own boundary, but object 4 only + // needs it to be present. + assert_eq!( + vec![3, 4], + chain_of(&[bounded(3, Some((2, 1)), 0), bounded(4, Some((3, 1)), 3)]) + ); + assert_eq!( + vec![4], + chain_of(&[bounded(3, Some((2, 1)), 0), bounded(4, Some((3, 1)), 4)]) + ); } #[tokio::test] @@ -4948,10 +5383,15 @@ mod tests { }); let result = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await; if moves { - result.unwrap().stop().await.unwrap(); + let store = result.unwrap(); let start = decode_header(&fixtures.get(2).await.unwrap()).unwrap(); assert_eq!(3, start.epoch); assert_eq!(Some(0), start.prev.map(|prev| prev.object_seq)); + // Object 3 carries the boundary past the late object, which + // recovery never listed. + append(&store, region(1), "a").await.unwrap(); + assert_eq!(vec![3], collected_object_seqs(&store).await); + store.stop().await.unwrap(); } else { let error = result.unwrap_err(); assert!( @@ -5024,9 +5464,15 @@ mod tests { // Open A resumes at the sequence it recovered: the epoch follows from // the sequence it claims, so it differs from the epoch of open B. - let a = start_epoch(&io, recovered.next_object_seq, recovered.tip) - .await - .unwrap(); + let a = start_epoch( + &io, + recovered.next_object_seq, + recovered.tip, + recovered.reclaim_boundary, + ) + .await + .unwrap() + .0; assert_eq!( ChainLink { object_seq: 1, @@ -5113,7 +5559,7 @@ mod tests { ); store.seal_open_batch().await.unwrap(); - assert_eq!(vec![0, 1], object_seqs(store.io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); assert_eq!(id(1, 2), latest(&store, region_id)); assert_eq!( id(1, 2), @@ -5127,8 +5573,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_durable_id_advances_after_create_and_indexing() { - let (store, io, mut parked) = - open_parking_creates(memory_store(), &enqueued(eager())).await; + let (store, _, mut parked) = open_parking_creates(memory_store(), &enqueued(eager())).await; let region_one = region(1); let region_two = region(2); @@ -5140,7 +5585,7 @@ mod tests { HashMap::from([(region_one, id(1, 1))]), response.last_entry_ids ); - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(1, object_seq); assert_eq!(0, store.durable_entry_id(&provider(region_one)).unwrap()); let wait = |entry_id| { @@ -5169,7 +5614,7 @@ mod tests { store.durable_entry_id(&provider(region_one)).unwrap() ); assert_eq!(0, store.durable_entry_id(&provider(region_two)).unwrap()); - assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); } #[tokio::test] @@ -5181,7 +5626,7 @@ mod tests { // Object 1 holds region A's id(1, 1) and object 2 region B's id(2, 1). for (region_id, data) in [(region_a, "a1"), (region_b, "b1")] { let response = append(&store, region_id, data).await.unwrap(); - let (_, release) = next_create(&mut parked).await; + let (_, release) = next_parked_operation(&mut parked).await; release.send(true).unwrap(); let entry_id = response.last_entry_ids[®ion_id]; timeout(WAIT, store.wait_durable(&provider(region_id), entry_id)) @@ -5194,7 +5639,7 @@ mod tests { HashMap::from([(region_a, id(3, 1))]), response.last_entry_ids ); - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(3, object_seq); // Every entry of region A up to id(2, 1) is durable, so the wait does @@ -5284,7 +5729,7 @@ mod tests { let stalled = spawn_append_batch(&store, vec![entry(&store, region_id, "a2")]); // The stall seals the open batch so that an upload is in flight. - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(1, object_seq); for _ in 0..16 { tokio::task::yield_now().await; @@ -5300,7 +5745,7 @@ mod tests { max_unpersisted_bytes: ReadableSize(1), ..enqueued(manual()) }; - let (store, io, mut parked, stalled, release) = stall_second_append(config).await; + let (store, _, mut parked, stalled, release) = stall_second_append(config).await; let region_id = region(1); // Two more appends queue behind the stalled one. let third = spawn_append_batch(&store, vec![entry(&store, region_id, "a3")]); @@ -5320,10 +5765,10 @@ mod tests { HashMap::from([(region_id, id(2, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); assert_eq!(id(1, 1), latest(&store, region_id)); for (append, object_seq) in [(third, 2), (fourth, 3)] { - let (parked_seq, release) = next_create(&mut parked).await; + let (parked_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(object_seq, parked_seq); release.send(true).unwrap(); let response = timeout(WAIT, append).await.unwrap().unwrap().unwrap(); @@ -5332,7 +5777,7 @@ mod tests { response.last_entry_ids ); } - assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); assert_eq!( expected_entries( region_id, @@ -5348,7 +5793,7 @@ mod tests { max_unpersisted_age: Duration::from_millis(50), ..enqueued(manual()) }; - let (store, io, _parked, stalled, release) = stall_second_append(config).await; + let (store, _, _parked, stalled, release) = stall_second_append(config).await; let region_id = region(1); release.send(true).unwrap(); @@ -5357,7 +5802,7 @@ mod tests { HashMap::from([(region_id, id(2, 1))]), response.last_entry_ids ); - assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); assert_eq!( expected_entries(region_id, &[(id(1, 1), "a1")]), read_entries(&store, region_id, 1).await @@ -5401,7 +5846,7 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(vec![0, 1], object_seqs(io.as_ref()).await); + assert_eq!(vec![1], collected_object_seqs(&store).await); assert_eq!( expected_entries(region_id, &[(id(1, 1), "a1")]), read_entries(&store, region_id, 1).await @@ -5415,7 +5860,7 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); } #[tokio::test] @@ -5466,7 +5911,7 @@ mod tests { let store = store.clone(); tokio::spawn(async move { store.seal_open_batch().await }) }; - let (object_seq, release) = next_create(&mut parked).await; + let (object_seq, release) = next_parked_operation(&mut parked).await; assert_eq!(1, object_seq); // Stop began, but the actor has not received the stop command when @@ -5644,6 +6089,521 @@ mod tests { assert_stopped(&timeout(WAIT, wait).await.unwrap().unwrap().unwrap_err()); } + /// The reclaim boundary in the header of the object `object_seq`. + async fn reclaim_boundary_of(io: &dyn WalObjectIo, object_seq: u64) -> u64 { + decode_header(&io.get(object_seq).await.unwrap()) + .unwrap() + .reclaim_boundary + } + + #[tokio::test] + async fn test_store_stamps_the_reclaim_boundary_and_collects_below_it() { + let store = open(memory_store(), &eager()).await; + let io = store.io.clone(); + let (one, two, three) = (region(1), region(2), region(3)); + let deleted = METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get(); + + // Object 1 carries the boundary past the indexed start object, which + // is collected; region one then holds the boundary at object 1. + append(&store, one, "a1").await.unwrap(); + append(&store, two, "b1").await.unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!(1, reclaim_boundary_of(io.as_ref(), 1).await); + assert_eq!(1, reclaim_boundary_of(io.as_ref(), 2).await); + + // A watermark moves the boundary; the next indexed object makes it + // durable. + store.obsolete(&provider(one), one, id(1, 1)).await.unwrap(); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + append(&store, one, "a2").await.unwrap(); + assert_eq!(2, reclaim_boundary_of(io.as_ref(), 3).await); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + + // Region one dropped and region two flushed through object 2: object + // 4 carries the sequence after the last indexed object. + store.obsolete(&provider(two), two, id(2, 1)).await.unwrap(); + store.obsolete_all(&provider(one), one).await.unwrap(); + append(&store, two, "b2").await.unwrap(); + assert_eq!(4, reclaim_boundary_of(io.as_ref(), 4).await); + assert_eq!(vec![4], collected_object_seqs(&store).await); + + // A region that appears later does not move the boundary back. + append(&store, three, "c1").await.unwrap(); + assert_eq!(4, reclaim_boundary_of(io.as_ref(), 5).await); + assert_eq!(vec![4, 5], collected_object_seqs(&store).await); + // Other stores of the process count their deletes as well. + assert!(METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get() >= deleted + 4); + assert_eq!( + expected_entries(two, &[(id(4, 1), "b2")]), + read_entries(&store, two, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_failed_batches_below_a_durable_boundary() { + let (io, _) = RecordingIo::over(memory_store()); + let store = open_over(io.clone(), &eager()).await; + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(2, 1)) + .await + .unwrap(); + + // Object 3 carries boundary 3 but its create is reported as failed, + // and the create of object 4 fails before it writes: neither boundary + // permits a deletion. + io.fail_after_next_put.store(true, Ordering::Relaxed); + append(&store, region_id, "a3").await.unwrap_err(); + assert_eq!(3, reclaim_boundary_of(io.as_ref(), 3).await); + io.fail_next_put_persistently(); + append(&store, region_id, "a4").await.unwrap_err(); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + + // Object 5 makes boundary 3 durable. The failed batches go once a + // durable boundary passes them, whether their object exists or not. + append(&store, region_id, "a5").await.unwrap(); + assert_eq!(vec![3, 5], collected_object_seqs(&store).await); + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + append(&store, region_id, "a6").await.unwrap(); + assert_eq!(vec![6], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(6, 1), "a6")]), + read_entries(&store, region_id, 0).await + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_attempts_a_failed_delete_again_at_the_next_collection() { + let _serialized = FAILED_DELETES.lock().await; + let (io, mut parked) = ParkedIo::parking_creates(memory_store()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + let failed = METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(false).unwrap(); + + // The object stays indexed, and an indexing that leaves the durable + // boundary where it is starts no collection. + timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(); + assert_eq!( + failed + 1, + METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get() + ); + assert!(store.catalog.read().unwrap().contains_object(0)); + append(&store, region_id, "a3").await.unwrap(); + assert!(parked.try_recv().is_err()); + assert_eq!(vec![0, 1, 2, 3], object_seqs(io.as_ref()).await); + + // The next `obsolete` starts a collection, which deletes it again. + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(0, object_seq); + release.send(true).unwrap(); + assert_eq!(vec![1, 2, 3], collected_object_seqs(&store).await); + assert!(!store.catalog.read().unwrap().contains_object(0)); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collection_ends_with_stop_and_poison() { + for poison in [false, true] { + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + for index in 1..=5 { + append(&store, region_id, &format!("a{index}")) + .await + .unwrap(); + } + store + .obsolete(&provider(region_id), region_id, id(5, 1)) + .await + .unwrap(); + append(&store, region_id, "a6").await.unwrap(); + // Objects 0 to 5 are below the durable boundary; four are deleted + // at a time. + let releases = parked_operations(&mut parked, MAX_IN_FLIGHT_DELETES).await; + assert!(parked.try_recv().is_err()); + + let mut stop = None; + if poison { + put_foreign(&object_store, 7, 100).await; + append(&store, region_id, "a7").await.unwrap_err(); + } else { + // The first poll queues the stop command; the actor answers + // the command of `obsolete_all` queued after it only once it + // handled the stop. + let mut stopping = Box::pin(store.stop()); + assert!((&mut stopping).now_or_never().is_none()); + store + .obsolete_all(&provider(region_id), region_id) + .await + .unwrap(); + assert!((&mut stopping).now_or_never().is_none()); + stop = Some(stopping); + } + for release in releases.into_values() { + release.send(true).unwrap(); + } + match stop { + Some(stop) => timeout(WAIT, stop).await.unwrap().unwrap(), + None => timeout(WAIT, store.wait_for_garbage_collection()) + .await + .unwrap(), + } + + // The deletes in flight completed and no other started. + assert!(parked.try_recv().is_err()); + assert_eq!( + MAX_IN_FLIGHT_DELETES, + io.max_in_flight.load(Ordering::SeqCst) + ); + let mut expected = vec![4, 5, 6]; + if poison { + expected.push(7); + } + assert_eq!(expected, object_seqs(io.as_ref()).await); + } + } + + #[tokio::test] + async fn test_store_drops_the_holes_of_a_collected_object() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let object_store = memory_store(); + write_corrupted_segment(&object_store).await; + let store = open(object_store, &eager()).await; + read_entries(&store, region(2), 0).await; + assert_eq!(1, store.wal_holes(&provider(region(2))).unwrap().len()); + + store + .obsolete(&provider(region(1)), region(1), 1) + .await + .unwrap(); + store + .obsolete(&provider(region(2)), region(2), 10) + .await + .unwrap(); + append(&store, region(1), "a").await.unwrap(); + assert_eq!(vec![3], collected_object_seqs(&store).await); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_records_no_hole_for_a_collected_object() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let object_store = memory_store(); + let (_, footer) = write_corrupted_segment(&object_store).await; + let store = open(object_store, &eager()).await; + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + + // A read met the damaged segment of object 0, which a collection + // removed before the read recorded the hole. + store.catalog.write().unwrap().remove_object(0); + let error = CorruptedWalObjectSnafu { reason: "damaged" }.build(); + record_skipped_segment( + &store.catalog, + &store.wal_holes, + store.io.as_ref(), + 0, + &footer[1], + &error, + ); + assert!(store.wal_holes(&provider(region(2))).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_a_failed_batch_whose_create_completes_late() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let region_id = region(1); + // Object 1 fails while the create of object 2 is in flight. + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&1).unwrap().send(false).unwrap(); + for append in appends { + timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + } + + // Objects 3 and 4 make boundary 3 durable; object 2 is skipped while + // its create is in flight. + for object_seq in [3, 4] { + let append = spawn_append_batch(&store, vec![entry(&store, region_id, "a")]); + let (parked_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(object_seq, parked_seq); + release.send(true).unwrap(); + timeout(WAIT, append).await.unwrap().unwrap().unwrap(); + } + assert_eq!(vec![3, 4], collected_object_seqs(&store).await); + + // The create completes and its object is deleted without another + // collection. + io.deletes_parked.store(true, Ordering::SeqCst); + releases.remove(&2).unwrap().send(true).unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(2, object_seq); + release.send(true).unwrap(); + assert_eq!(vec![3, 4], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + + async fn recover_listed(io: &ObjectStoreIo, listed: Vec) -> Result { + finish_recovery(fetch_footers(io, listed, RECOVERY_CONCURRENCY).await?) + } + + #[tokio::test] + async fn test_store_recovery_tolerates_objects_collected_after_the_listing() { + let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); + let bounded = |object_seq, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, 1, prev) + }; + put_object_with_header(&io, chain_header(0, 1, None), &[]).await; + put_object_with_header(&io, chain_header(1, 1, Some((0, 1))), &[]).await; + put_object_with_header(&io, bounded(2, Some((1, 1)), 2), &[]).await; + put_object_with_header(&io, bounded(3, Some((2, 1)), 2), &[]).await; + let listed = io.list().await.unwrap(); + + // A delete of an earlier instance removes object 1, below the + // boundary of the tip, after the listing. + io.delete(1).await.unwrap(); + let recovered = recover_listed(&io, listed.clone()).await.unwrap(); + assert_eq!(Some(3), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(2, recovered.reclaim_boundary); + assert_eq!(vec![0], recovered.unindexed); + assert_eq!(4, recovered.next_object_seq); + + // Object 2 is needed above the boundary: its disappearance fails + // recovery. + io.delete(2).await.unwrap(); + let error = recover_listed(&io, listed).await.unwrap_err(); + assert!(is_not_found(&error), "unexpected error: {error:?}"); + } + + #[tokio::test] + async fn test_store_read_skips_an_object_collected_after_it_was_listed() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let store = open(memory_store(), &eager()).await; + let region_id = region(1); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + // The read lists objects 1 and 2 before a collection deletes object 1. + let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + assert_eq!(vec![2, 3], collected_object_seqs(&store).await); + + let entries = stream + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect::>(); + assert_eq!(expected_entries(region_id, &[(id(2, 1), "a2")]), entries); + assert!(store.wal_holes(&provider(region_id)).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_read_waits_for_an_undecided_delete() { + let _serialized = SKIPPED_SEGMENTS.lock().await; + let _failing = FAILED_DELETES.lock().await; + for deleted in [true, false] { + let object_store = memory_store(); + let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + io.deletes_parked.store(true, Ordering::SeqCst); + let region_id = region(1); + let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); + append(&store, region_id, "a1").await.unwrap(); + next_parked_operation(&mut parked) + .await + .1 + .send(true) + .unwrap(); + append(&store, region_id, "a2").await.unwrap(); + + // The read lists objects 1 and 2. The delete of object 1 removes + // it, but its outcome is not reported yet. + let stream = store.read(&provider(region_id), 0, None).await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + object_store + .delete(&object_path(&object_store, 1)) + .await + .unwrap(); + let read = tokio::spawn(stream.try_collect::>()); + let waiting = || { + store + .deleting + .lock() + .get(&1) + .map_or(0, |settled| settled.receiver_count()) + }; + timeout(WAIT, async { + while waiting() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(!read.is_finished()); + + // A delete that succeeded skips the object; one that failed + // leaves it indexed, and the read fails on the fetch. + release.send(deleted).unwrap(); + let result = timeout(WAIT, read).await.unwrap().unwrap(); + if deleted { + let entries = result.unwrap().into_iter().flatten().collect::>(); + assert_eq!(expected_entries(region_id, &[(id(2, 1), "a2")]), entries); + } else { + let error = result.unwrap_err(); + assert!(is_not_found(&error), "unexpected error: {error:?}"); + } + assert!(store.wal_holes(&provider(region_id)).unwrap().is_empty()); + assert_eq!( + skipped, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get() + ); + store.stop().await.unwrap(); + } + } + + #[tokio::test] + async fn test_store_recovers_from_the_reclaim_boundary_of_the_tip() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let region_id = region(1); + let store = open(object_store.clone(), &eager()).await; + append(&store, region_id, "a1").await.unwrap(); + store + .obsolete(&provider(region_id), region_id, id(1, 1)) + .await + .unwrap(); + append(&store, region_id, "a2").await.unwrap(); + assert_eq!(vec![2], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + // Object 2 links to object 1, below its boundary. The start object + // carries the recovered boundary. + let store = open(object_store.clone(), &eager()).await; + assert_eq!(2, reclaim_boundary_of(&io, 3).await); + assert_eq!( + expected_entries(region_id, &[(id(2, 1), "a2")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(2, 1)) + .await + .unwrap(); + append(&store, region_id, "a3").await.unwrap(); + assert_eq!(vec![4], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + + let store = open(object_store, &eager()).await; + assert_eq!(4, reclaim_boundary_of(&io, 5).await); + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "a3")]), + read_entries(&store, region_id, 0).await + ); + assert_eq!(vec![4, 5], object_seqs(&io).await); + store.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_store_collects_a_late_object_of_an_earlier_epoch_below_the_boundary() { + let object_store = memory_store(); + let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let region_id = region(1); + let record = |entry_id, data: &'static str| Record { + region_id, + entry_id, + payload: Bytes::from_static(data.as_bytes()), + }; + let bounded = |object_seq, epoch, prev, reclaim_boundary| Header { + reclaim_boundary, + ..chain_header(object_seq, epoch, prev) + }; + put_object_with_header(&io, chain_header(0, 1, None), &[]).await; + put_object_with_header( + &io, + chain_header(1, 1, Some((0, 1))), + &[record(id(1, 1), "a")], + ) + .await; + // Epoch 4 started at object 3, and object 4 carries boundary 4. + // Object 2 of epoch 1 landed late; its link holds, but its epoch is + // lower. + put_object_with_header( + &io, + bounded(2, 1, Some((1, 1)), 2), + &[record(id(2, 1), "late")], + ) + .await; + put_object_with_header(&io, chain_header(3, 4, Some((1, 1))), &[]).await; + put_object_with_header( + &io, + bounded(4, 4, Some((3, 4)), 4), + &[record(id(4, 1), "b")], + ) + .await; + let recovered = recover(&io).await.unwrap(); + assert_eq!(Some(4), recovered.tip.map(|tip| tip.object_seq)); + assert_eq!(4, recovered.reclaim_boundary); + assert_eq!(vec![0, 1, 2, 3], recovered.unindexed); + + // The start object makes boundary 4 durable, and the first + // collection deletes everything below it. + let store = open(object_store, &eager()).await; + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "b")]), + read_entries(&store, region_id, 0).await + ); + store + .obsolete(&provider(region_id), region_id, id(4, 1)) + .await + .unwrap(); + assert_eq!(vec![4, 5], collected_object_seqs(&store).await); + store.stop().await.unwrap(); + } + /// Object access whose next create stores the object but reports a /// transient failure, as a create whose response is lost. struct LostResponseIo { @@ -5676,6 +6636,10 @@ mod tests { self.inner.get_range(object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + async fn list(&self) -> Result> { self.inner.list().await } @@ -5710,6 +6674,10 @@ mod tests { self.inner.get_range(object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + async fn list(&self) -> Result> { self.inner.list().await } @@ -5719,16 +6687,18 @@ mod tests { } } - /// Object access whose whole-object reads, or conditional creates, park - /// until the test releases them, counting how many are in flight. A - /// release of false fails the operation with a transient error before it - /// reaches the object store. + /// Object access whose whole-object reads, or conditional creates, and + /// on request deletes, park until the test releases them, counting how + /// many are in flight. A release of false fails the operation with a + /// transient error before it reaches the object store. struct ParkedIo { inner: ObjectStoreIo, parked: mpsc::UnboundedSender<(u64, oneshot::Sender)>, park_creates: bool, /// Whether creates park yet; set once the store opened. creates_parked: AtomicBool, + /// Whether deletes park. + deletes_parked: AtomicBool, in_flight: AtomicUsize, max_in_flight: AtomicUsize, } @@ -5766,6 +6736,7 @@ mod tests { parked, park_creates, creates_parked: AtomicBool::new(false), + deletes_parked: AtomicBool::new(false), in_flight: AtomicUsize::new(0), max_in_flight: AtomicUsize::new(0), }), @@ -5809,6 +6780,13 @@ mod tests { self.inner.get_range(object_seq, offset, len).await } + async fn delete(&self, object_seq: u64) -> Result<()> { + if self.deletes_parked.load(Ordering::SeqCst) && !self.park(object_seq).await { + return injected_failure("delete", self.object_path(object_seq)); + } + self.inner.delete(object_seq).await + } + async fn list(&self) -> Result> { self.inner.list().await } @@ -5896,6 +6874,10 @@ mod tests { Ok(bytes) } + async fn delete(&self, object_seq: u64) -> Result<()> { + self.inner.delete(object_seq).await + } + async fn list(&self) -> Result> { self.inner.list().await } diff --git a/src/mito2/src/engine/object_store_wal_recovery_test.rs b/src/mito2/src/engine/object_store_wal_recovery_test.rs index 917ad614a07..126e38e4b26 100644 --- a/src/mito2/src/engine/object_store_wal_recovery_test.rs +++ b/src/mito2/src/engine/object_store_wal_recovery_test.rs @@ -373,6 +373,98 @@ async fn test_reopen_after_partial_flush_replays_only_unflushed_regions(#[case] assert_eq!(5, engine.get_region_statistic(REGION_B).unwrap().num_rows); } +/// Waits until no WAL object delete is in flight. +async fn wait_for_collection(store: &ObjectStoreLogStore) { + tokio::time::timeout(WAIT, store.wait_for_garbage_collection()) + .await + .expect("collection must complete"); +} + +#[rstest] +#[case(AckMode::Durable)] +#[case(AckMode::Enqueued)] +#[tokio::test] +async fn test_flush_collects_objects_below_the_watermark(#[case] ack_mode: AckMode) { + let mut env = TestEnv::with_prefix("object-store-wal-collection").await; + let object_store = memory_store(); + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + let (table_dir_a, schema_a) = create_region(&engine, REGION_A, &[]).await; + let (table_dir_b, schema_b) = create_region(&engine, REGION_B, &[]).await; + + // Object 1 holds entries of both regions, object 2 of region A and + // object 3 of region B. Object 1 carries the boundary past the start + // object, which is collected. + let mut writer = SealedWriter::new(&store); + writer + .put_and_seal( + &engine, + vec![ + (REGION_A, rows(&schema_a, 0, 2)), + (REGION_B, rows(&schema_b, 0, 3)), + ], + ) + .await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 2, 4))]) + .await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 3, 5))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![1, 2, 3], wal_object_seqs(&object_store).await); + + // Region B has no watermark and still needs object 1, so flushing + // region A collects nothing. + flush_region(&engine, REGION_A, None).await; + writer + .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 4, 6))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![1, 2, 3, 4], wal_object_seqs(&object_store).await); + + // Once region B is flushed, the first object a region needs is object 4, + // the unflushed entry of region A; object 5 carries that boundary. + flush_region(&engine, REGION_B, None).await; + writer + .put_and_seal(&engine, vec![(REGION_B, rows(&schema_b, 5, 7))]) + .await; + wait_for_collection(&store).await; + assert_eq!(vec![4, 5], wal_object_seqs(&object_store).await); + let rows_a = scan_rows(&engine, REGION_A).await; + let rows_b = scan_rows(&engine, REGION_B).await; + engine.stop().await.unwrap(); + drop(engine); + drop(writer); + drop(store); + + // Recovery stops at the boundary and each region replays its entry above + // its flushed id. + let store = open_store(&object_store, ack_mode).await; + let engine = new_engine(&mut env, store.clone()).await; + open_region(&engine, REGION_A, &table_dir_a, &[]) + .await + .unwrap(); + open_region(&engine, REGION_B, &table_dir_b, &[]) + .await + .unwrap(); + for (region_id, flushed_entry_id, last_entry_id) in [ + (REGION_A, entry_id(2, 1), entry_id(4, 1)), + (REGION_B, entry_id(3, 1), entry_id(5, 1)), + ] { + let state = region_recovery_state(&engine, region_id).await; + assert_eq!(flushed_entry_id, state.flushed_entry_id); + assert_eq!(last_entry_id, state.last_entry_id); + assert_eq!(2, state.memtable_rows); + } + assert_eq!(rows_a, scan_rows(&engine, REGION_A).await); + assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); + assert_eq!(6, engine.get_region_statistic(REGION_A).unwrap().num_rows); + assert_eq!(7, engine.get_region_statistic(REGION_B).unwrap().num_rows); + wait_for_collection(&store).await; + assert_eq!(vec![4, 5, 6], wal_object_seqs(&object_store).await); +} + #[rstest] #[case(AckMode::Durable)] #[case(AckMode::Enqueued)] @@ -600,12 +692,14 @@ async fn test_reopen_on_empty_prefix_without_durable_entries() { assert_eq!(rows_b, scan_rows(&engine, REGION_B).await); // The region is usable and its first entry lands in the first object - // after the start object. + // after the start object, which carries the boundary past the start + // object. let mut writer = SealedWriter::new(&store); writer .put_and_seal(&engine, vec![(REGION_A, rows(&schema_a, 0, 2))]) .await; - assert_eq!(vec![0, 1], wal_object_seqs(&object_store).await); + wait_for_collection(&store).await; + assert_eq!(vec![1], wal_object_seqs(&object_store).await); assert_eq!(entry_id(1, 1), latest(&store, REGION_A)); assert_eq!( entry_id(1, 1), @@ -766,7 +860,10 @@ async fn test_reopen_after_a_failed_create_shows_only_the_acknowledged_write() { writer .put_and_seal(&engine, vec![(REGION_A, key_rows(2))]) .await; - assert_eq!(vec![0, 1, 2], wal_object_seqs(&object_store).await); + // Object 2 carries the boundary past the start object, which is + // collected. + wait_for_collection(&store).await; + assert_eq!(vec![1, 2], wal_object_seqs(&object_store).await); assert_eq!(entry_id(2, 1), latest(&store, REGION_A)); let region_a = region(&engine, REGION_A); assert_eq!(1, region_a.version_control.committed_sequence()); @@ -979,7 +1076,10 @@ async fn test_enqueued_crash_before_the_object_exists_replays_durable_entries() entry_id(2, 1), region_recovery_state(&engine, REGION_A).await.last_entry_id ); - assert_eq!(vec![0, 1, 2], wal_object_seqs(&object_store).await); + // Opening the region recorded its watermark, so object 2 carries the + // boundary past both start objects. + wait_for_collection(&store).await; + assert_eq!(vec![2], wal_object_seqs(&object_store).await); let rows_before = scan_rows(&engine, REGION_A).await; assert_eq!(2, engine.get_region_statistic(REGION_A).unwrap().num_rows); drop(writer); diff --git a/tests-integration/tests/object_store_wal.rs b/tests-integration/tests/object_store_wal.rs index 1ce37992b5f..0aef1e145b0 100644 --- a/tests-integration/tests/object_store_wal.rs +++ b/tests-integration/tests/object_store_wal.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::path::Path; +use std::time::Duration; use common_procedure::options::ProcedureConfig; use common_query::Output; @@ -67,6 +68,18 @@ async fn count_wal_objects(standalone: &GreptimeDbStandalone) -> usize { .count() } +/// Waits until collection leaves `expected` WAL objects under [`WAL_ROOT`]. +async fn wait_for_wal_objects(standalone: &GreptimeDbStandalone, expected: usize) { + let collected = async { + while count_wal_objects(standalone).await != expected { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + tokio::time::timeout(Duration::from_secs(30), collected) + .await + .unwrap_or_else(|_| panic!("expected {expected} WAL objects after collection")); +} + /// Stops the procedure manager, the region server and the WAL of the /// instance. async fn shutdown(standalone: &mut GreptimeDbStandalone) { @@ -103,7 +116,8 @@ async fn restart( } /// Writes batches that land in separate WAL objects, restarts before and -/// after a flush, and checks that the rows survive both restarts. +/// after a flush, and checks that the rows survive both restarts and that the +/// flush lets collection delete every earlier object. async fn run_restarts_on_object_store_wal(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let wal_config = ObjectStoreWalConfig { @@ -129,29 +143,14 @@ async fn run_restarts_on_object_store_wal(store_type: StorageType) { // Every insert returns once its entries are durable, so the batches land // in separate WAL objects after the object that started the epoch. for batch in 0..BATCHES { - let values = (0..ROWS_PER_BATCH) - .map(|row| { - let i = batch * ROWS_PER_BATCH + row; - format!( - "('host_{i}', {i}.0, {})", - 1_686_567_600_000 + i as i64 * 1000 - ) - }) - .collect::>() - .join(", "); - execute_sql( - standalone.fe_instance(), - &format!("INSERT INTO cpu VALUES {values}"), - ) - .await; + insert_batch(&standalone, batch).await; } let expected = query_rows(&standalone).await; assert_eq!(BATCHES * ROWS_PER_BATCH + 4, expected.lines().count()); - let objects = count_wal_objects(&standalone).await; - assert!( - objects > BATCHES, - "expected more than {BATCHES} WAL objects" - ); + // The region has no watermark, so its first object holds the boundary + // and only the object that started the epoch is collected. + wait_for_wal_objects(&standalone, BATCHES).await; + let objects = BATCHES; // No Raft Engine log store is created as a fallback. assert!( !Path::new(&standalone.opts.storage.data_home) @@ -165,13 +164,37 @@ async fn run_restarts_on_object_store_wal(store_type: StorageType) { assert_eq!(expected, query_rows(&standalone).await); assert!(count_wal_objects(&standalone).await > objects); - // The flushed rows come back from the SST files. + // After the flush the next object carries a reclaim boundary past every + // earlier object, so collection leaves that object alone. The flushed + // rows come back from the SST files and the last batch from the WAL. execute_sql(standalone.fe_instance(), "ADMIN FLUSH_TABLE('cpu')").await; + insert_batch(&standalone, BATCHES).await; + wait_for_wal_objects(&standalone, 1).await; + let expected = query_rows(&standalone).await; let mut standalone = restart(&builder, standalone).await; assert_eq!(expected, query_rows(&standalone).await); shutdown(&mut standalone).await; } +/// Inserts the rows of `batch`, which land in one WAL object. +async fn insert_batch(standalone: &GreptimeDbStandalone, batch: usize) { + let values = (0..ROWS_PER_BATCH) + .map(|row| { + let i = batch * ROWS_PER_BATCH + row; + format!( + "('host_{i}', {i}.0, {})", + 1_686_567_600_000 + i as i64 * 1000 + ) + }) + .collect::>() + .join(", "); + execute_sql( + standalone.fe_instance(), + &format!("INSERT INTO cpu VALUES {values}"), + ) + .await; +} + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_object_store_wal_survives_restarts_on_file() { run_restarts_on_object_store_wal(StorageType::File).await; From 198d92407c49bbcf12e0c116f1f81f56fcf13319 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 9 Oct 2026 02:28:32 +0000 Subject: [PATCH 109/133] fix(operator): bound FlowMirrorTask pending rows to prevent frontend OOM (#9421) * fix(operator): bound FlowMirrorTask pending rows to prevent frontend OOM When flownode is slow or unreachable, FlowMirrorTask::detach spawns unbounded tasks that each pin a full insert batch for ~10s (gRPC timeout). With continuous writes, in-flight memory grows unboundedly (throughput x 10s x peers), causing frontend-write OOMKill and a million-per-hour ERROR storm. Add a MAX_MIRROR_PENDING_ROWS cap (1M rows) on Inserter: detach fetch_adds the batch's row count before spawning; if the cap is exceeded the reservation is rolled back, the batch is dropped, a DIST_MIRROR_DROPPED_ROW_COUNT metric is bumped, and a warn is logged instead of spawning. On normal completion the reservation is released in both success and error paths. This keeps core writes alive while degrading flow mirror writes to best-effort, which is the intended trade-off: protect the primary insert path at the cost of possible flow aggregation gaps. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix: address review - balanced reservation, gauge race, log rate limit P1: Compute pending_rows() from actual per-peer payloads in self.requests instead of the pre-fanout num_rows field. This fixes under-reservation for same-table multi-region batches and over-release in multi-flownode fanout. Each peer task now releases exactly its own share so reservations and releases balance. P2: Replace load+set gauge publication with add/sub symmetry. reserve_mirror_pending_rows does fetch_add + gauge.add, release does fetch_update + gauge.sub, eliminating the stale-value race. P3: Rate-limit the dropped-batch warning to at most once per 10s via MirrorDropLog (atomic compare_exchange), accumulating batch/row counts between reports. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: fix race in dual-peer test - wait for gauge release before unlocking The dual-peer test previously only waited for the pending budget to return to its start value before releasing MIRROR_METRIC_TEST_LOCK. But release_mirror_pending_rows moves the budget before the gauge, so a worker's gauge.sub could still be in flight when the next test acquires the lock and records its baseline, polluting the result. Now waits for both pending == start AND gauge == gauge_before. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(operator): preserve instant writes and limit mirror failure logs Reject saturated requests containing instant-TTL rows with the existing retryable write-admission error before mirror or datanode dispatch. Keep normal source-table inserts independent of best-effort mirror drops. Reuse a smaller process-wide event limiter for both dropped batches and failed peer tasks, reporting at most once per 10 seconds with failure count, peer context, and error chain. Cover actual mirror construction, mixed and instant admission, fast/slow peer shares, failed-task release, and log concurrency. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(operator): reject oversized instant mirror batches permanently Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/operator/src/insert.rs | 898 +++++++++++++++++++++++++++++++++++- src/operator/src/metrics.rs | 5 + 2 files changed, 885 insertions(+), 18 deletions(-) diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index e1fbe39a9f4..80866b45b66 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -13,7 +13,9 @@ // limitations under the License. use std::collections::BTreeMap; -use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::v1::alter_table_expr::Kind; @@ -104,6 +106,10 @@ pub struct Inserter { /// per-request `auto_create_table` hint. When `true`, the hint still applies. auto_create_table: bool, pending_rows_batcher: Option>, + /// Rows handed to detached flow mirror tasks that have not finished yet. + /// Bounds how much row data in-flight mirror writes may keep alive, see + /// [`MAX_MIRROR_PENDING_ROWS`]. + mirror_pending_rows: Arc, } pub type InserterRef = Arc; @@ -300,6 +306,7 @@ impl Inserter { table_flownode_set_cache, auto_create_table, pending_rows_batcher: None, + mirror_pending_rows: Arc::new(AtomicU64::new(0)), } } @@ -804,7 +811,17 @@ impl Inserter { .chain(instant_requests.requests.iter()), ) .await?; - flow_mirror_task.detach(self.node_manager.clone())?; + let has_instant_rows = instant_requests.requests.iter().any(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + flow_mirror_task.detach( + self.node_manager.clone(), + self.mirror_pending_rows.clone(), + has_instant_rows, + )?; // Write requests to datanode and wait for response let write_tasks = self @@ -2004,9 +2021,65 @@ struct CreateAlterTableResult { table_infos: HashMap>, } +/// Upper bound of rows buffered by detached flow mirror tasks on this frontend. +/// +/// Rows retained by in-flight mirror RPCs count against this bound. If ordinary +/// persisted source-table mirroring exceeds it, that mirror batch is best-effort +/// dropped while its datanode write continues. Instant-TTL batches fail retryably +/// on temporary saturation; oversized batches fail non-retryably so the client +/// can reduce the batch size. +const MAX_MIRROR_PENDING_ROWS: usize = 1_000_000; + +/// How often at most each mirror event is reported. +const MIRROR_LOG_INTERVAL: Duration = Duration::from_secs(10); +const MIRROR_LOG_NEVER_REPORTED: u64 = u64::MAX; + +struct MirrorLog { + last_log_millis: AtomicU64, + events: AtomicU64, +} + +impl MirrorLog { + const fn new() -> Self { + Self { + last_log_millis: AtomicU64::new(MIRROR_LOG_NEVER_REPORTED), + events: AtomicU64::new(0), + } + } + + /// Returns the count since the previous report, including this event, to + /// the single caller that claims the interval. + fn claim_report(&self, now_millis: u64) -> Option { + self.events.fetch_add(1, Ordering::Relaxed); + let last = self.last_log_millis.load(Ordering::Relaxed); + if last != MIRROR_LOG_NEVER_REPORTED + && now_millis.saturating_sub(last) < MIRROR_LOG_INTERVAL.as_millis() as u64 + { + return None; + } + if self + .last_log_millis + .compare_exchange(last, now_millis, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return None; + } + Some(self.events.swap(0, Ordering::Relaxed)) + } +} + +static MIRROR_DROP_LOG: MirrorLog = MirrorLog::new(); +static MIRROR_FAILURE_LOG: MirrorLog = MirrorLog::new(); + +/// Milliseconds since the first call, so the drop log is rate limited on a +/// monotonic clock instead of the wall clock, which can jump. +fn mirror_drop_log_millis() -> u64 { + static START: LazyLock = LazyLock::new(Instant::now); + START.elapsed().as_millis() as u64 +} + struct FlowMirrorTask { requests: HashMap, - num_rows: usize, } impl FlowMirrorTask { @@ -2016,7 +2089,6 @@ impl FlowMirrorTask { ) -> Result { let mut src_table_reqs: HashMap, RegionInsertRequests)>> = HashMap::new(); - let mut num_rows = 0; for req in requests { let table_id = RegionId::from_u64(req.region_id).table_id(); @@ -2040,11 +2112,6 @@ impl FlowMirrorTask { if !peers.is_empty() { let mut reqs = RegionInsertRequests::default(); reqs.requests.push(req.clone()); - num_rows += reqs - .requests - .iter() - .map(|r| r.rows.as_ref().unwrap().rows.len()) - .sum::(); src_table_reqs.insert(table_id, Some((peers, reqs))); } else { // insert a empty entry to avoid repeat query @@ -2080,16 +2147,69 @@ impl FlowMirrorTask { } } - Ok(Self { - requests: inserts, - num_rows, - }) + Ok(Self { requests: inserts }) } - fn detach(self, node_manager: NodeManagerRef) -> Result<()> { - crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(self.num_rows as i64); + /// Rows of row data this task keeps alive once detached: the payload it + /// spawns per peer, not the source requests it was built from. + /// + /// A source table mapped to several flownodes clones its requests per peer + /// and every clone owns its own copy of the rows, so reserving from + /// `self.requests` is what the per-peer shares released by [`Self::detach`] + /// sum to. Counting the source requests instead under-reserves whenever a + /// batch holds several requests for the same source table. + fn pending_rows(&self) -> u64 { + self.requests.values().map(region_inserts_rows).sum() + } + + fn detach( + self, + node_manager: NodeManagerRef, + mirror_pending_rows: Arc, + has_instant_rows: bool, + ) -> Result<()> { + // Reserve the cloned per-peer payload before spawning. A full budget + // drops normal mirrors or rejects instant writes before datanode dispatch. + let num_rows = self.pending_rows(); + if num_rows == 0 { + return Ok(()); + } + if has_instant_rows && num_rows > MAX_MIRROR_PENDING_ROWS as u64 { + return InvalidInsertRequestSnafu { + reason: format!( + "flow mirror batch has {num_rows} peer-fan-out rows, exceeding the limit of {}; reduce the batch size", + MAX_MIRROR_PENDING_ROWS + ), + } + .fail(); + } + + let pending = reserve_mirror_pending_rows(&mirror_pending_rows, num_rows); + if pending > MAX_MIRROR_PENDING_ROWS as u64 { + release_mirror_pending_rows(&mirror_pending_rows, num_rows); + if has_instant_rows { + return Err(meter_core::collect::WriteRejected::new( + "flow mirror pending rows limit exceeded for instant-TTL insert", + )) + .context(WriteRejectedSnafu); + } + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.inc_by(num_rows); + if let Some(events) = MIRROR_DROP_LOG.claim_report(mirror_drop_log_millis()) { + warn!( + "Flow mirror write dropped: {} pending rows exceeds limit {}, {} drop events since last report", + pending, MAX_MIRROR_PENDING_ROWS, events + ); + } + return Ok(()); + } + for (peer, inserts) in self.requests { + // Each spawned task releases exactly its own share of the + // reservation. The shares are computed over the same per-peer + // payload as `pending_rows`, so they sum to the reservation. + let peer_rows = region_inserts_rows(&inserts); let node_manager = node_manager.clone(); + let mirror_pending_rows = mirror_pending_rows.clone(); common_runtime::spawn_global(async move { let result = node_manager .flownode(&peer) @@ -2102,12 +2222,18 @@ impl FlowMirrorTask { Ok(resp) => { let affected_rows = resp.affected_rows; crate::metrics::DIST_MIRROR_ROW_COUNT.inc_by(affected_rows); - crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(affected_rows as _); } Err(err) => { - error!(err; "Failed to insert data into flownode {}", peer); + if let Some(events) = + MIRROR_FAILURE_LOG.claim_report(mirror_drop_log_millis()) + { + error!(err; flownode_id = peer.id, flownode_addr = %peer.addr, "Failed to insert data into flownode ({} total mirror failures across peers since last report)", events); + } } } + // Release the reservation on both the success and the failure + // path, otherwise a broken flownode would fill the budget forever. + release_mirror_pending_rows(&mirror_pending_rows, peer_rows); }); } @@ -2115,17 +2241,62 @@ impl FlowMirrorTask { } } +/// Rows of row data carried by `inserts`. A request without a `rows` payload +/// carries nothing to keep buffered and contributes zero. +fn region_inserts_rows(inserts: &RegionInsertRequests) -> u64 { + inserts + .requests + .iter() + .filter_map(|req| req.rows.as_ref()) + .map(|rows| rows.rows.len() as u64) + .sum() +} + +/// Reserves `rows` of the pending mirror budget and returns the resulting +/// pending row count. The caller must [`release_mirror_pending_rows`] exactly +/// these rows once they are no longer buffered, whether it spawns the batch or +/// drops it. +fn reserve_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) -> u64 { + let pending = mirror_pending_rows.fetch_add(rows, Ordering::Relaxed) + rows; + // Keep the gauge in step by applying the same increment, never by setting an + // absolute value: another task may move the budget in between. + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(rows as i64); + pending +} + +/// Releases `rows` of the pending mirror budget. +/// +/// Every release mirrors the reservation it hands back, so a well-behaved +/// release never subtracts more than is pending; the saturating update only +/// keeps a duplicated release from wrapping the shared budget around. +fn release_mirror_pending_rows(mirror_pending_rows: &AtomicU64, rows: u64) { + let prev = mirror_pending_rows + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |pending| { + Some(pending.saturating_sub(rows)) + }) + .unwrap_or_else(|prev| prev); + // Subtract from the gauge with the opposite sign of + // [`reserve_mirror_pending_rows`]; reading the budget here and `set`-ting it + // would publish a stale value whenever another task releases concurrently. + let released = rows.min(prev); + if released > 0 { + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(released as i64); + } +} + #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicU64, Ordering}; use api::helper::ColumnDataTypeWrapper; + use api::v1::flow::FlowResponse; use api::v1::helper::{field_column_schema, time_index_column_schema}; use api::v1::{Row, RowInsertRequest, Rows, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_meta::cache::new_table_flownode_set_cache; use common_meta::ddl::test_util::datanode_handler::NaiveDatanodeHandler; - use common_meta::test_util::MockDatanodeManager; + use common_meta::test_util::{MockDatanodeManager, MockFlownodeHandler, MockFlownodeManager}; use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; use common_query::prelude::{greptime_native_histogram, set_default_prefix}; use datatypes::data_type::ConcreteDataType; @@ -3340,4 +3511,695 @@ mod tests { .unwrap(); assert!(matches!(output.data, OutputData::AffectedRows(0))); } + + /// Flownode handler that keeps mirror inserts in flight until the test + /// releases the gate, so the pending mirror budget can be observed + /// deterministically. + #[derive(Clone)] + struct GatedFlownodeHandler { + gate: Arc, + affected_rows: u64, + } + + #[async_trait::async_trait] + impl MockFlownodeHandler for GatedFlownodeHandler { + async fn handle_inserts( + &self, + _peer: &Peer, + _requests: api::v1::region::InsertRequests, + ) -> common_meta::error::Result { + let _permit = self.gate.acquire().await.unwrap(); + Ok(FlowResponse { + affected_rows: self.affected_rows, + ..Default::default() + }) + } + } + + fn flownode_peer() -> Peer { + Peer { + id: 1, + addr: "127.0.0.1:4001".to_string(), + } + } + + #[derive(Clone)] + struct RecordingNodeManager { + datanode_dispatch: tokio::sync::mpsc::UnboundedSender, + flownode_dispatch: tokio::sync::mpsc::UnboundedSender, + gates: Arc>>>, + fail_flownode: bool, + } + + #[async_trait::async_trait] + impl common_meta::node_manager::DatanodeManager for RecordingNodeManager { + async fn datanode(&self, peer: &Peer) -> common_meta::node_manager::DatanodeRef { + Arc::new(RecordingNode { + peer: peer.clone(), + manager: self.clone(), + }) + } + } + + #[async_trait::async_trait] + impl common_meta::node_manager::FlownodeManager for RecordingNodeManager { + async fn flownode(&self, peer: &Peer) -> common_meta::node_manager::FlownodeRef { + Arc::new(RecordingNode { + peer: peer.clone(), + manager: self.clone(), + }) + } + } + + struct RecordingNode { + peer: Peer, + manager: RecordingNodeManager, + } + + #[async_trait::async_trait] + impl common_meta::node_manager::Datanode for RecordingNode { + async fn handle( + &self, + _request: api::v1::region::RegionRequest, + ) -> common_meta::error::Result { + let rows = match _request.body.as_ref() { + Some(api::v1::region::region_request::Body::Inserts(requests)) => requests + .requests + .iter() + .filter_map(|request| request.rows.as_ref()) + .map(|rows| rows.rows.len()) + .sum(), + _ => 0, + }; + let _ = self.manager.datanode_dispatch.send(rows); + Ok(api::region::RegionResponse::new(rows)) + } + + async fn handle_query( + &self, + _request: common_query::request::QueryRequest, + ) -> common_meta::error::Result { + unreachable!() + } + } + + #[async_trait::async_trait] + impl common_meta::node_manager::Flownode for RecordingNode { + async fn handle( + &self, + _request: api::v1::flow::FlowRequest, + ) -> common_meta::error::Result { + unreachable!() + } + + async fn handle_inserts( + &self, + _request: api::v1::region::InsertRequests, + ) -> common_meta::error::Result { + let _ = self.manager.flownode_dispatch.send(self.peer.clone()); + if self.manager.fail_flownode { + return Err(common_meta::error::UnexpectedSnafu { + err_msg: "test flownode failure".to_string(), + } + .build()); + } + let gate = self + .manager + .gates + .lock() + .unwrap() + .get(&self.peer.id) + .cloned(); + if let Some(gate) = gate { + let _permit = gate.acquire().await.unwrap(); + } + Ok(FlowResponse::default()) + } + + async fn handle_mark_window_dirty( + &self, + _request: api::v1::flow::DirtyWindowRequests, + ) -> common_meta::error::Result { + unreachable!() + } + } + + fn mirror_requests(peer: &Peer, num_rows: usize) -> HashMap { + HashMap::from_iter([( + peer.clone(), + RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(1, 1).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { values: vec![] }; num_rows], + }), + ..Default::default() + }], + }, + )]) + } + + /// Serializes the mirror metric tests below: the metrics are process-wide + /// singletons, so these tests would otherwise observe each other's updates + /// when the test binary runs tests in parallel (they are isolated per + /// process only under nextest). + static MIRROR_METRIC_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + + async fn lock_mirror_metrics() -> tokio::sync::MutexGuard<'static, ()> { + MIRROR_METRIC_TEST_LOCK.lock().await + } + + // The tests below share the process-wide mirror metrics, so they take + // `MIRROR_METRIC_TEST_LOCK` and assert on the difference their own budget + // publishes instead of on absolute values. + #[tokio::test] + async fn flow_mirror_saturation_preserves_instant_and_normal_contracts() { + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + + let _guard = lock_mirror_metrics().await; + let kv_backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(kv_backend.clone()).await; + let (datanode_tx, mut datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates: Arc::new(std::sync::Mutex::new(HashMap::new())), + fail_flownode: false, + }); + let flow_cache = Cache::new(10); + let flow_cache_backend = prepare_mocked_backend().await; + let inserter = Inserter::new( + catalog::memory::MemoryCatalogManager::new(), + partition_manager, + node_manager, + Arc::new(new_table_flownode_set_cache( + String::new(), + flow_cache.clone(), + flow_cache_backend, + )), + true, + ); + let peer = flownode_peer(); + inserter + .table_flownode_set_cache + .invalidate(&[common_meta::instruction::CacheIdent::CreateFlow( + common_meta::instruction::CreateFlow { + flow_id: 1, + source_table_ids: vec![1, 2], + partition_to_peer_mapping: vec![(0, peer)], + }, + )]) + .await + .unwrap(); + let mut normal_info = new_test_table_info(1, "normal_table", [1].into_iter()); + normal_info.catalog_name = DEFAULT_CATALOG_NAME.to_string(); + normal_info.schema_name = DEFAULT_SCHEMA_NAME.to_string(); + let mut instant_info = new_test_table_info(2, "instant_table", [1].into_iter()); + instant_info.meta.options.ttl = Some(common_time::ttl::TimeToLive::Instant); + let table_infos = + HashMap::from_iter([(1, Arc::new(normal_info)), (2, Arc::new(instant_info))]); + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + let insert = |table_id, num_rows| RegionInsertRequests { + requests: vec![RegionInsertRequest { + region_id: RegionId::new(table_id, 1).as_u64(), + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); num_rows], + }), + ..Default::default() + }], + }; + let make_request = |normal, instant| InstantAndNormalInsertRequests { + normal_requests: if normal == 0 { + RegionInsertRequests::default() + } else { + insert(1, normal) + }, + instant_requests: if instant == 0 { + RegionInsertRequests::default() + } else { + insert(2, instant) + }, + }; + + let pending = inserter.mirror_pending_rows.clone(); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64); + pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed); + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + for request in [make_request(0, 1), make_request(1, 1)] { + let error = inserter + .do_request(request, &table_infos, &ctx) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert!(datanode_rx.try_recv().is_err()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!( + pending.load(Ordering::Relaxed), + MAX_MIRROR_PENDING_ROWS as u64 + ); + } + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + ); + + for pending_before in [0, MAX_MIRROR_PENDING_ROWS as u64] { + if pending_before == 0 { + pending.store(0, Ordering::Relaxed); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64); + } + let gauge_before_rejection = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + for (normal, instant) in [ + (0, MAX_MIRROR_PENDING_ROWS + 1), + (MAX_MIRROR_PENDING_ROWS / 2, MAX_MIRROR_PENDING_ROWS / 2 + 1), + ] { + let error = inserter + .do_request(make_request(normal, instant), &table_infos, &ctx) + .await + .unwrap_err(); + assert_eq!(error.status_code(), StatusCode::InvalidArguments); + assert_eq!(error.retry_hint(), RetryHint::NonRetryable); + let message = error.to_string(); + assert!(message.contains("1000001"), "{message}"); + assert!(message.contains("1000000"), "{message}"); + assert!(message.contains("reduce the batch size"), "{message}"); + assert!(datanode_rx.try_recv().is_err()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!(pending.load(Ordering::Relaxed), pending_before); + assert_eq!( + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + gauge_before_rejection + ); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + ); + } + if pending_before == 0 { + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.add(MAX_MIRROR_PENDING_ROWS as i64); + pending.store(MAX_MIRROR_PENDING_ROWS as u64, Ordering::Relaxed); + } + } + + let result = inserter + .do_request(make_request(1, 0), &table_infos, &ctx) + .await + .unwrap(); + assert!(matches!(result.data, OutputData::AffectedRows(1))); + assert_eq!(1, datanode_rx.try_recv().unwrap()); + assert!(flownode_rx.try_recv().is_err()); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + 1 + ); + pending.store(0, Ordering::Relaxed); + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.sub(MAX_MIRROR_PENDING_ROWS as i64); + assert_eq!( + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + gauge_before + ); + } + + #[tokio::test] + async fn flow_mirror_new_counts_regions_and_peer_clones() { + use common_meta::instruction::{CacheIdent, CreateFlow}; + + let peer1 = flownode_peer(); + let peer2 = Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }; + let cache_backend = prepare_mocked_backend().await; + let cache = Arc::new(new_table_flownode_set_cache( + String::new(), + Cache::new(10), + cache_backend, + )); + cache + .invalidate(&[CacheIdent::CreateFlow(CreateFlow { + flow_id: 1, + source_table_ids: vec![1], + partition_to_peer_mapping: vec![(0, peer1), (1, peer2)], + })]) + .await + .unwrap(); + let request = |region_id| RegionInsertRequest { + region_id, + rows: Some(Rows { + schema: vec![], + rows: vec![api::v1::Row { values: vec![] }; 3], + }), + ..Default::default() + }; + let requests = [ + request(RegionId::new(1, 0).as_u64()), + request(RegionId::new(1, 1).as_u64()), + ]; + let task = FlowMirrorTask::new(&cache, requests.iter()).await.unwrap(); + assert_eq!(12, task.pending_rows()); + assert_eq!(2, task.requests.len()); + } + + #[tokio::test] + async fn flow_mirror_dropped_when_pending_exceeds_limit() { + let _guard = lock_mirror_metrics().await; + let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64)); + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + let num_rows = 100; + + let task = FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + }; + // The budget is already full, so nothing may be spawned for this batch. + task.detach( + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + pending.clone(), + false, + ) + .unwrap(); + + assert_eq!( + num_rows as u64, + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get() - dropped_before + ); + // The dropped batch must not leak its reservation. + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + // The reservation taken before the limit check is released again, so the + // gauge nets out and keeps matching the pending budget. + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + "a dropped batch must not change the pending gauge" + ); + + let task = FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + }; + let error = task + .detach( + Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)), + pending.clone(), + true, + ) + .err() + .unwrap(); + use common_error::ext::{ErrorExt, RetryHint}; + use common_error::status_code::StatusCode; + assert_eq!(error.status_code(), StatusCode::RateLimited); + assert_eq!(error.retry_hint(), RetryHint::Retryable); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + num_rows as u64, + "rejected instant rows are not successful best-effort drops" + ); + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() + ); + } + + /// A source table mapped to several flownodes clones its requests per peer, + /// so the reservation must cover every clone and each spawned task must hand + /// back exactly its own share; otherwise the shared budget drifts. + #[tokio::test] + async fn flow_mirror_reserves_and_releases_per_peer_shares() { + let _guard = lock_mirror_metrics().await; + let per_peer_rows = 10; + let total_rows = (2 * per_peer_rows) as u64; + // Exactly enough room for the whole batch, cloned rows included. + let start = MAX_MIRROR_PENDING_ROWS as u64 - total_rows; + let pending = Arc::new(AtomicU64::new(start)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler { + gate: gate.clone(), + affected_rows: per_peer_rows as u64, + })); + + let mut requests = mirror_requests(&flownode_peer(), per_peer_rows); + requests.extend(mirror_requests( + &Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }, + per_peer_rows, + )); + + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { requests } + .detach(node_manager, pending.clone(), false) + .unwrap(); + + // Both per-peer payloads are counted, so the budget is exactly full + // instead of over-reserved by the duplicated rows. + assert_eq!( + MAX_MIRROR_PENDING_ROWS as u64, + pending.load(Ordering::Relaxed) + ); + + // Let both tasks finish and check that the budget lands back on its + // starting value: over-releasing a shared reservation would saturate it. + // The gauge is checked too because `release_mirror_pending_rows` moves + // the budget before the gauge; exiting on the budget alone could leak a + // pending `gauge.sub` into the next test holding the metrics lock. + gate.add_permits(2); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != start + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "the mirror tasks must release exactly their own shares" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_fast_peer_releases_only_its_share_before_drop() { + let _guard = lock_mirror_metrics().await; + let peer1 = flownode_peer(); + let peer2 = Peer { + id: 2, + addr: "127.0.0.1:4002".to_string(), + }; + let gate1 = Arc::new(tokio::sync::Semaphore::new(0)); + let gate2 = Arc::new(tokio::sync::Semaphore::new(0)); + let gates = Arc::new(std::sync::Mutex::new(HashMap::from_iter([ + (peer1.id, gate1.clone()), + (peer2.id, gate2.clone()), + ]))); + let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates, + fail_flownode: false, + }); + let pending = Arc::new(AtomicU64::new(MAX_MIRROR_PENDING_ROWS as u64 - 20)); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + let mut requests = mirror_requests(&peer1, 10); + requests.extend(mirror_requests(&peer2, 10)); + let task = FlowMirrorTask { requests }; + task.detach(node_manager.clone(), pending.clone(), false) + .unwrap(); + let first = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_ne!(first.id, second.id); + let fast_gate = if first.id == peer1.id { + gate1.clone() + } else { + gate2.clone() + }; + let slow_gate = if first.id == peer1.id { + gate2.clone() + } else { + gate1.clone() + }; + fast_gate.add_permits(1); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 10 { + assert!( + std::time::Instant::now() < deadline, + "fast peer did not release" + ); + std::thread::sleep(Duration::from_millis(1)); + } + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + FlowMirrorTask { + requests: mirror_requests(&peer1, 11), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + assert_eq!( + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + dropped_before + 11 + ); + assert!(flownode_rx.try_recv().is_err()); + slow_gate.add_permits(1); + while pending.load(Ordering::Relaxed) != MAX_MIRROR_PENDING_ROWS as u64 - 20 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "slow peer did not release" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_failure_releases_budget_and_gauge() { + let _guard = lock_mirror_metrics().await; + let (datanode_tx, _datanode_rx) = tokio::sync::mpsc::unbounded_channel(); + let (flownode_tx, mut flownode_rx) = tokio::sync::mpsc::unbounded_channel(); + let node_manager = Arc::new(RecordingNodeManager { + datanode_dispatch: datanode_tx, + flownode_dispatch: flownode_tx, + gates: Arc::new(std::sync::Mutex::new(HashMap::new())), + fail_flownode: true, + }); + let pending = Arc::new(AtomicU64::new(0)); + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), 5), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(10), flownode_rx.recv()) + .await + .unwrap() + .unwrap() + .id, + flownode_peer().id + ); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != 0 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "failed task leaked reservation" + ); + std::thread::sleep(Duration::from_millis(1)); + } + } + + #[tokio::test] + async fn flow_mirror_pending_counter_consistent() { + let _guard = lock_mirror_metrics().await; + let pending = Arc::new(AtomicU64::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let node_manager = Arc::new(MockFlownodeManager::new(GatedFlownodeHandler { + gate: gate.clone(), + affected_rows: 7, + })); + + // A batch without rows short-circuits and leaves the budget untouched. + let gauge_before = crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(); + FlowMirrorTask { + requests: HashMap::new(), + } + .detach(node_manager.clone(), pending.clone(), false) + .unwrap(); + assert_eq!(0, pending.load(Ordering::Relaxed)); + assert_eq!( + gauge_before, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get(), + "an empty batch must not touch the pending gauge" + ); + + let dropped_before = crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(); + let num_rows = 7; + + FlowMirrorTask { + requests: mirror_requests(&flownode_peer(), num_rows), + } + .detach(node_manager, pending.clone(), false) + .unwrap(); + + // The spawned task is still waiting on the gated flownode, so the + // reservation made by `detach` is the value observable here. + assert_eq!(num_rows as u64, pending.load(Ordering::Relaxed)); + assert_eq!( + num_rows as i64, + crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() - gauge_before, + "the pending gauge must track the pending rows" + ); + assert_eq!( + dropped_before, + crate::metrics::DIST_MIRROR_DROPPED_ROW_COUNT.get(), + "a batch within the limit must not be counted as dropped" + ); + + // Release the in-flight task and wait for it to hand its reservation + // back, so no task is left holding the shared metrics after this test. + gate.add_permits(1); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + while pending.load(Ordering::Relaxed) != 0 + || crate::metrics::DIST_MIRROR_PENDING_ROW_COUNT.get() != gauge_before + { + assert!( + std::time::Instant::now() < deadline, + "the mirror task did not release its pending reservation" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } + + #[test] + fn mirror_logs_are_rate_limited() { + let log = Arc::new(MirrorLog::new()); + assert_eq!(Some(1), log.claim_report(1_000)); + assert_eq!(None, log.claim_report(1_500)); + assert_eq!(None, log.claim_report(10_999)); + assert_eq!(Some(3), log.claim_report(11_000)); + assert_eq!(None, log.claim_report(11_001)); + assert_eq!(Some(2), log.claim_report(21_000)); + + let concurrent = Arc::new(MirrorLog::new()); + assert_eq!(Some(1), concurrent.claim_report(50_000)); + let barrier = Arc::new(std::sync::Barrier::new(8)); + let threads = (0..8) + .map(|_| { + let log = concurrent.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + log.claim_report(60_000) + }) + }) + .collect::>(); + let winners = threads + .into_iter() + .filter_map(|thread| thread.join().unwrap()) + .collect::>(); + assert_eq!(winners.len(), 1); + let residual = concurrent.events.swap(0, Ordering::Relaxed); + assert_eq!(winners[0] + residual, 8); + } } diff --git a/src/operator/src/metrics.rs b/src/operator/src/metrics.rs index 3e0b2fed2a6..438904ed36c 100644 --- a/src/operator/src/metrics.rs +++ b/src/operator/src/metrics.rs @@ -47,6 +47,11 @@ lazy_static! { "table operator mirror pending rows" ) .unwrap(); + pub static ref DIST_MIRROR_DROPPED_ROW_COUNT: IntCounter = register_int_counter!( + "greptime_operator_dist_mirror_dropped_rows", + "table operator mirror rows dropped because the pending mirror row limit was exceeded" + ) + .unwrap(); pub static ref DIST_DELETE_ROW_COUNT: IntCounterVec = register_int_counter_vec!( "greptime_table_operator_delete_rows", "table operator delete rows", From 27c1aa76d31649747b7c7f6de01bd66f250decdb Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 03:13:28 +0000 Subject: [PATCH 110/133] chore(deps): bump xxhash-rust from 0.8.15 to 0.8.19 (#9465) Bumps [xxhash-rust](https://github.com/DoumanAsh/xxhash-rust) from 0.8.15 to 0.8.19. - [Commits](https://github.com/DoumanAsh/xxhash-rust/commits) --- updated-dependencies: - dependency-name: xxhash-rust dependency-version: 0.8.19 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 79e325b19b7..c08ac6f1ca8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17340,9 +17340,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "550a2b930b62486a393c52d5c3b84bff264b28aa437ed64694d31e93b1757af7" [[package]] name = "yaml-rust" From fe8a647bc8f87897520f10ed2c2bcbb8629642ff Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Fri, 9 Oct 2026 03:17:22 +0000 Subject: [PATCH 111/133] fix(json2): support settings in ALTER TABLE ADD COLUMN (#9472) * fix(json2): support settings in ALTER TABLE ADD COLUMN Signed-off-by: fys * refactor(sql): reuse Column for ADD COLUMN definitions Signed-off-by: fys * chore: better method name * refactor(sql): share column parsing between CREATE TABLE and ADD COLUMN Signed-off-by: fys * chore: simplify tests Signed-off-by: fys * fix(sql): accept trailing commas in ALTER TABLE ADD COLUMN Signed-off-by: fys * fix(sql): preserve ALTER TABLE ADD COLUMN compatibility Signed-off-by: fys --------- Signed-off-by: fys --- src/operator/src/expr_helper.rs | 40 +++- src/sql/src/parsers/alter_parser.rs | 211 +++++++++++++----- src/sql/src/parsers/create_parser.rs | 43 ++-- src/sql/src/statements.rs | 111 +++++---- src/sql/src/statements/alter.rs | 12 +- src/sql/src/statements/create.rs | 11 +- .../src/statements/transform/type_alias.rs | 2 +- .../common/types/json/json2_alter.result | 87 ++++++++ .../common/types/json/json2_alter.sql | 27 +++ .../common/types/json/json2_limit.result | 21 -- .../common/types/json/json2_limit.sql | 13 -- 11 files changed, 404 insertions(+), 174 deletions(-) diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 02c1af46706..d69277a1c90 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -68,8 +68,8 @@ use sql::statements::create::{ CreateView, TableConstraint, }; use sql::statements::{ - OptionMap, column_to_schema, concrete_data_type_to_sql_data_type, - sql_column_def_to_grpc_column_def, sql_data_type_to_concrete_data_type, value_to_sql_value, + OptionMap, column_to_schema, concrete_data_type_to_sql_data_type, sql_col_to_grpc_col_def, + sql_data_type_to_concrete_data_type, value_to_sql_value, }; use sql::util::extract_tables_from_query; use store_api::mito_engine_options::{COMPACTION_OVERRIDE, COMPACTION_TYPE}; @@ -828,12 +828,10 @@ pub(crate) fn to_alter_table_expr( add_columns: add_columns .into_iter() .map(|add_column| { - let column_def = sql_column_def_to_grpc_column_def( - &add_column.column_def, - Some(&query_ctx.timezone()), - ) - .map_err(BoxedError::new) - .context(ExternalSnafu)?; + let column_def = + sql_col_to_grpc_col_def(&add_column.column, Some(&query_ctx.timezone())) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; if is_interval_type(&column_def.data_type()) { return NotSupportedSnafu { feat: "Add column with interval type", @@ -1225,6 +1223,32 @@ mod tests { assert!(is_json2_extension_type(&schema.arrow_schema().fields()[0])); } + #[test] + fn test_to_alter_add_json2_column_preserves_settings() { + let sql = "ALTER TABLE monitor ADD COLUMN payload JSON2(max_auto_expanded_paths = 0, service STRING, nested.value BIGINT)"; + let stmt = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .pop() + .unwrap(); + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let expr = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap(); + let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = expr.kind.unwrap() else { + unreachable!() + }; + let column_schema = + try_as_column_schema(add_columns[0].column_def.as_ref().unwrap()).unwrap(); + let metadata: datatypes::extension::json::JsonMetadata = + serde_json::from_str(&column_schema.metadata()["ARROW:extension:metadata"]).unwrap(); + let settings = metadata.json_settings(); + assert_eq!(Some(0), settings.max_auto_expanded_paths()); + assert_eq!(2, settings.type_hints().len()); + assert_eq!(vec!["service"], settings.type_hints()[0].path); + assert_eq!(vec!["nested", "value"], settings.type_hints()[1].path); + } + #[test] fn test_create_flow_tql_expr() { let sql = r#" diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index 3e5c27fda73..62c417d2537 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -26,7 +26,7 @@ use sqlparser::parser::IsOptional::Mandatory; use sqlparser::parser::{Parser, ParserError}; use sqlparser::tokenizer::{Token, TokenWithSpan}; -use crate::ast::ObjectNamePartExt; +use crate::ast::{ColumnOption, ObjectNamePartExt}; use crate::error::{self, InvalidColumnOptionSnafu, Result, SetFulltextOptionSnafu}; use crate::parser::ParserContext; use crate::parsers::create_parser::{INVERTED, parse_json2_type_and_options}; @@ -379,14 +379,67 @@ impl ParserContext<'_> { Ok(AlterTableOperation::AddConstraint(constraint)) } else { self.parser.prev_token(); - let add_columns = self - .parser - .parse_comma_separated(parse_add_columns) - .context(error::SyntaxSnafu)?; + let mut add_columns = vec![self.parse_add_column()?]; + while self.parser.consume_token(&Token::Comma) { + if matches!( + self.parser.peek_token().token, + Token::SemiColon | Token::EOF + ) || self.matches_keyword(Keyword::WITH) + { + break; + } + add_columns.push(self.parse_add_column()?); + } Ok(AlterTableOperation::AddColumns { add_columns }) } } + fn parse_add_column(&mut self) -> Result { + self.parser + .expect_keyword(Keyword::ADD) + .context(error::SyntaxSnafu)?; + let _ = self.parser.parse_keyword(Keyword::COLUMN); + let add_if_not_exists = + self.parser + .parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); + let column = self.parse_column_def(true)?; + + // These attributes need ALTER-specific execution support. + ensure!( + column.extensions.fulltext_index_options.is_none() + && column.extensions.skipping_index_options.is_none() + && column.extensions.inverted_index_options.is_none() + && !column + .options() + .iter() + .any(|option| matches!(option.option, ColumnOption::DialectSpecific(_))), + InvalidColumnOptionSnafu { + name: column.name().to_string(), + msg: "ADD COLUMN does not support inline indexes or TIME INDEX", + } + ); + + let location = if self.parser.parse_keyword(Keyword::FIRST) { + Some(AddColumnLocation::First) + } else if matches!(self.parser.peek_token().token, Token::Word(word) if word.value.eq_ignore_ascii_case("AFTER")) + { + self.parser.next_token(); + let name = Self::canonicalize_identifier( + self.parser.parse_identifier().context(error::SyntaxSnafu)?, + ); + Some(AddColumnLocation::After { + column_name: name.value, + }) + } else { + None + }; + Ok(AddColumn { + column, + location, + add_if_not_exists, + }) + } + fn parse_alter_table_drop_default( &mut self, column_name: Ident, @@ -704,34 +757,6 @@ fn parse_string_options(parser: &mut Parser) -> std::result::Result<(String, Str Ok((name, value)) } -fn parse_add_columns(parser: &mut Parser) -> std::result::Result { - parser.expect_keyword(Keyword::ADD)?; - let _ = parser.parse_keyword(Keyword::COLUMN); - let add_if_not_exists = parser.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); - let mut column_def = parser.parse_column_def()?; - column_def.name = ParserContext::canonicalize_identifier(column_def.name); - let location = if parser.parse_keyword(Keyword::FIRST) { - Some(AddColumnLocation::First) - } else if let Token::Word(word) = parser.peek_token().token { - if word.value.eq_ignore_ascii_case("AFTER") { - let _ = parser.next_token(); - let name = ParserContext::canonicalize_identifier(parser.parse_identifier()?); - Some(AddColumnLocation::After { - column_name: name.value, - }) - } else { - None - } - } else { - None - }; - Ok(AddColumn { - column_def, - location, - add_if_not_exists, - }) -} - /// Parses a comma separated list of string literals. fn parse_string_option_names(parser: &mut Parser) -> std::result::Result { parser.parse_literal_string() @@ -750,6 +775,75 @@ mod tests { use crate::dialect::GreptimeDbDialect; use crate::parser::ParseOptions; use crate::statements::alter::AlterDatabaseOperation; + use crate::statements::create::{Column, ColumnExtensions}; + + #[test] + fn test_add_column_matches_create_table() { + for (definition, location) in [ + ( + "payload JSON2(max_auto_expanded_paths = 100, service STRING) NOT NULL", + "AFTER ts", + ), + ("payload JSON2(service STRING) NULL", "FIRST"), + ("embedding VECTOR(3) NOT NULL", "AFTER ts"), + ("reading INT DEFAULT 42 COMMENT 'value'", ""), + ("\"select\" STRING NULL", ""), + ] { + let parse = |sql: &str| { + ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .pop() + .unwrap() + }; + let Statement::CreateTable(create) = parse(&format!( + "CREATE TABLE t ({definition}, n INT, ts TIMESTAMP TIME INDEX)" + )) else { + unreachable!() + }; + let Statement::AlterTable(alter) = parse(&format!( + "ALTER TABLE t ADD COLUMN {definition} {location}, ADD COLUMN n INT" + )) else { + unreachable!() + }; + let AlterTableOperation::AddColumns { add_columns } = alter.alter_operation() else { + unreachable!() + }; + assert_eq!(add_columns.len(), 2); + for (expected, actual) in create.columns.iter().zip(add_columns) { + assert_eq!(expected, &actual.column); + } + let Statement::AlterTable(roundtrip) = parse(&alter.to_string()) else { + unreachable!() + }; + assert_eq!(alter, roundtrip); + } + } + + #[test] + fn test_add_column_rejects_unsupported_attributes() { + for definition in [ + "payload STRING FULLTEXT INDEX", + "payload STRING SKIPPING INDEX", + "payload STRING INVERTED INDEX", + "ts TIMESTAMP TIME INDEX", + ] { + let err = ParserContext::create_with_dialect( + &format!("ALTER TABLE t ADD COLUMN {definition}"), + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("ADD COLUMN does not support inline indexes or TIME INDEX"), + "{err}" + ); + } + } #[test] fn test_parse_alter_database() { @@ -830,10 +924,14 @@ mod tests { match alter_operation { AlterTableOperation::AddColumns { add_columns } => { assert_eq!(add_columns.len(), 1); - assert_eq!("tagk_i", add_columns[0].column_def.name.value); - assert_eq!(DataType::String(None), add_columns[0].column_def.data_type); + assert_eq!("tagk_i", add_columns[0].column.column_def.name.value); + assert_eq!( + DataType::String(None), + add_columns[0].column.column_def.data_type + ); assert!( add_columns[0] + .column .column_def .options .iter() @@ -866,10 +964,14 @@ mod tests { assert_matches!(alter_operation, AlterTableOperation::AddColumns { .. }); match alter_operation { AlterTableOperation::AddColumns { add_columns } => { - assert_eq!("tagk_i", add_columns[0].column_def.name.value); - assert_eq!(DataType::String(None), add_columns[0].column_def.data_type); + assert_eq!("tagk_i", add_columns[0].column.column_def.name.value); + assert_eq!( + DataType::String(None), + add_columns[0].column.column_def.data_type + ); assert!( add_columns[0] + .column .column_def .options .iter() @@ -931,7 +1033,7 @@ mod tests { .zip(expecteds) .collect::, ColumnDef))>>() { - assert_eq!(add_column.column_def, expected.1); + assert_eq!(add_column.column.column_def, expected.1); assert_eq!(&expected.0, &add_column.location); } } @@ -962,28 +1064,37 @@ mod tests { AlterTableOperation::AddColumns { add_columns } => { let expected = [ AddColumn { - column_def: ColumnDef { - name: Ident::new("a"), - data_type: DataType::Integer(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("a"), + data_type: DataType::Integer(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: true, }, AddColumn { - column_def: ColumnDef { - name: Ident::new("b"), - data_type: DataType::String(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("b"), + data_type: DataType::String(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: false, }, AddColumn { - column_def: ColumnDef { - name: Ident::new("c"), - data_type: DataType::Int(None), - options: vec![], + column: Column { + column_def: ColumnDef { + name: Ident::new("c"), + data_type: DataType::Int(None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }, location: None, add_if_not_exists: true, diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 32606679a1c..c87d0274382 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -41,8 +41,8 @@ use table::requests::{validate_database_option, validate_database_option_value}; use crate::ast::{ColumnDef, Ident, ObjectNamePartExt}; use crate::error::{ self, InvalidColumnOptionSnafu, InvalidDatabaseOptionSnafu, InvalidDatabaseOptionValueSnafu, - InvalidFlowQuerySnafu, InvalidIntervalSnafu, InvalidSqlSnafu, InvalidTimeIndexSnafu, - MissingTimeIndexSnafu, Result, SyntaxSnafu, UnexpectedSnafu, UnsupportedSnafu, + InvalidFlowQuerySnafu, InvalidIntervalSnafu, InvalidTimeIndexSnafu, MissingTimeIndexSnafu, + Result, SyntaxSnafu, UnexpectedSnafu, UnsupportedSnafu, }; use crate::parser::{FLOW, ParserContext}; use crate::parsers::tql_parser; @@ -685,7 +685,7 @@ impl<'a> ParserContext<'a> { columns: &mut Vec, constraints: &mut Vec, ) -> Result<()> { - let mut column = self.parse_column_def()?; + let mut column = self.parse_column_def(false)?; let mut time_index_opt_idx = None; for (index, opt) in column.options().iter().enumerate() { @@ -775,22 +775,16 @@ impl<'a> ParserContext<'a> { Ok(name) } - pub fn parse_column_def(&mut self) -> Result { - let name = self.parse_column_name().context(SyntaxSnafu)?; + /// Parses a column definition, optionally allowing unquoted keyword names + /// for compatibility with ALTER TABLE ADD COLUMN. + pub fn parse_column_def(&mut self, allow_keyword_name: bool) -> Result { + let name = if allow_keyword_name { + self.parser.parse_identifier().context(SyntaxSnafu)? + } else { + self.parse_column_name().context(SyntaxSnafu)? + }; let parser = &mut self.parser; - ensure!( - !(name.quote_style.is_none() && - // "ALL_KEYWORDS" are sorted. - ALL_KEYWORDS.binary_search(&name.value.to_uppercase().as_str()).is_ok()), - InvalidSqlSnafu { - msg: format!( - "Cannot use keyword '{}' as column name. Hint: add quotes to the name.", - name.value - ), - } - ); - let mut extensions = ColumnExtensions::default(); let data_type = @@ -2924,11 +2918,16 @@ ENGINE=mito"; #[test] fn test_invalid_column_name() { - let sql = "create table foo(user string, i timestamp time index)"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()); - let err = result.unwrap_err().output_msg(); - assert!(err.contains("Cannot use keyword 'user' as column name")); + for name in ["user", "at", "value", "select"] { + let sql = format!("create table foo({name} string, i timestamp time index)"); + let result = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ); + let err = result.unwrap_err().output_msg(); + assert!(err.contains(&format!("Cannot use keyword '{name}' as column name"))); + } // If column name is quoted, it's valid even same with keyword. let sql = r#" diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index 8a001510045..c2c97c34567 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -173,11 +173,12 @@ fn set_json2_extension(column_schema: &mut ColumnSchema, column: &Column) -> Res Ok(()) } -/// Convert `ColumnDef` in sqlparser to `ColumnDef` in gRPC proto. -pub fn sql_column_def_to_grpc_column_def( - col: &ColumnDef, +/// Converts a SQL [`Column`] to a `ColumnDef` in gRPC proto. +pub fn sql_col_to_grpc_col_def( + column: &Column, timezone: Option<&Timezone>, ) -> Result { + let col = &column.column_def; let name = col.name.value.clone(); let data_type = sql_data_type_to_concrete_data_type(&col.data_type)?; @@ -208,9 +209,11 @@ pub fn sql_column_def_to_grpc_column_def( SemanticType::Field }; - // TODO(fys): Extend the ALTER TABLE ADD COLUMN parser to support JSON2 - // type hints and pass the parsed JsonSettings through this conversion. - let options = json2_extension(&col.data_type, JsonSettings::new_v2()).map(|extension| { + let settings = column + .extensions + .build_json_settings()? + .unwrap_or_else(JsonSettings::new_v2); + let options = json2_extension(&col.data_type, settings).map(|extension| { let mut options = ColumnOptions::default(); options.options.insert( EXTENSION_TYPE_NAME_KEY.to_string(), @@ -489,15 +492,18 @@ mod tests { } #[test] - pub fn test_sql_column_def_to_grpc_column_def() { + pub fn test_sql_col_to_grpc_col_def() { // test basic - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![], + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!("col", grpc_column_def.name); assert!(grpc_column_def.is_nullable); // nullable when options are empty @@ -506,56 +512,65 @@ mod tests { assert_eq!(grpc_column_def.semantic_type, SemanticType::Field as i32); // test not null - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::NotNull, - }], + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![ColumnOptionDef { + name: None, + option: ColumnOption::NotNull, + }], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert!(!grpc_column_def.is_nullable); // test primary key - let column_def = ColumnDef { - name: "col".into(), - data_type: SqlDataType::Double(ExactNumberInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::PrimaryKey(PrimaryKeyConstraint { + let column = Column { + column_def: ColumnDef { + name: "col".into(), + data_type: SqlDataType::Double(ExactNumberInfo::None), + options: vec![ColumnOptionDef { name: None, - index_name: None, - index_type: None, - columns: vec![], - index_options: vec![], - characteristics: None, - }), - }], + option: ColumnOption::PrimaryKey(PrimaryKeyConstraint { + name: None, + index_name: None, + index_type: None, + columns: vec![], + index_options: vec![], + characteristics: None, + }), + }], + }, + extensions: ColumnExtensions::default(), }; - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!(grpc_column_def.semantic_type, SemanticType::Tag as i32); } #[test] - pub fn test_sql_column_def_to_grpc_column_def_with_timezone() { - let column_def = ColumnDef { - name: "col".into(), - // MILLISECOND - data_type: SqlDataType::Timestamp(Some(3), TimezoneInfo::None), - options: vec![ColumnOptionDef { - name: None, - option: ColumnOption::Default(Expr::Value( - SqlValue::SingleQuotedString("2024-01-30T00:01:01".to_string()).into(), - )), - }], + pub fn test_sql_col_to_grpc_col_def_with_timezone() { + let column = Column { + column_def: ColumnDef { + name: "col".into(), + // MILLISECOND + data_type: SqlDataType::Timestamp(Some(3), TimezoneInfo::None), + options: vec![ColumnOptionDef { + name: None, + option: ColumnOption::Default(Expr::Value( + SqlValue::SingleQuotedString("2024-01-30T00:01:01".to_string()).into(), + )), + }], + }, + extensions: ColumnExtensions::default(), }; // with timezone "Asia/Shanghai" - let grpc_column_def = sql_column_def_to_grpc_column_def( - &column_def, + let grpc_column_def = sql_col_to_grpc_col_def( + &column, Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()), ) .unwrap(); @@ -575,7 +590,7 @@ mod tests { ); // without timezone - let grpc_column_def = sql_column_def_to_grpc_column_def(&column_def, None).unwrap(); + let grpc_column_def = sql_col_to_grpc_col_def(&column, None).unwrap(); assert_eq!("col", grpc_column_def.name); assert!(grpc_column_def.is_nullable); // nullable when options are empty assert_eq!( diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index 4fc45c2e69b..781c65ab528 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -22,11 +22,11 @@ use common_query::AddColumnLocation; use datatypes::schema::{FulltextOptions, SkippingIndexOptions}; use itertools::Itertools; use serde::Serialize; -use sqlparser::ast::{ColumnDef, DataType, Expr, Ident, ObjectName, TableConstraint}; +use sqlparser::ast::{DataType, Expr, Ident, ObjectName, TableConstraint}; use sqlparser_derive::{Visit, VisitMut}; use crate::statements::OptionMap; -use crate::statements::create::{Json2Options, Partitions}; +use crate::statements::create::{Column, Json2Options, Partitions}; #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct AlterTable { @@ -229,18 +229,18 @@ pub enum UnsetIndexOperation { #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct AddColumn { - pub column_def: ColumnDef, + pub column: Column, pub location: Option, pub add_if_not_exists: bool, } impl Display for AddColumn { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.column)?; if let Some(location) = &self.location { - write!(f, "{} {location}", self.column_def) - } else { - write!(f, "{}", self.column_def) + write!(f, " {location}")?; } + Ok(()) } } diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index 9bd6c341a65..919c1cd31f9 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -137,7 +137,7 @@ pub struct CreateTable { pub partitions: Option, } -/// Column definition in `CREATE TABLE` statement. +/// Column definition in `CREATE TABLE` and `ALTER TABLE ADD COLUMN`. #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] pub struct Column { /// `ColumnDef` from `sqlparser::ast` @@ -267,10 +267,11 @@ impl Column { && let Some(vector_options) = &self.extensions.vector_options && let Some(dim) = vector_options.get(VECTOR_OPT_DIM) { - return write!(f, "{} VECTOR({})", self.column_def.name, dim); + write!(f, "{} VECTOR({})", self.column_def.name, dim)?; + } else { + write!(f, "{} ", self.column_def.name)?; + write_data_type(f, &self.column_def.data_type)?; } - write!(f, "{} ", self.column_def.name)?; - write_data_type(f, &self.column_def.data_type)?; if let Some(options) = &self.extensions.json2_options { write!(f, "{options}")?; } @@ -771,7 +772,7 @@ mod tests { name: None, option: sqlparser::ast::ColumnOption::Comment("private-value".repeat(8192)), }); - assert!(!vector.to_string().contains("private-value")); + assert!(vector.to_string().contains("private-value")); assert!(!vector.consume_sql_size_budget(&mut 4096)); vector.columns[1].column_def.options.clear(); assert!(vector.consume_sql_size_budget(&mut 4096)); diff --git a/src/sql/src/statements/transform/type_alias.rs b/src/sql/src/statements/transform/type_alias.rs index b2d1803fb37..90157901f3c 100644 --- a/src/sql/src/statements/transform/type_alias.rs +++ b/src/sql/src/statements/transform/type_alias.rs @@ -72,7 +72,7 @@ impl TransformRule for TypeAliasTransformRule { alter_table.alter_operation_mut() { for add_column in add_columns { - replace_type_alias(&mut add_column.column_def.data_type); + replace_type_alias(add_column.column.mut_data_type()); } } } diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result index 153aa9e7902..055ea90346a 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.result +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -150,3 +150,90 @@ DROP TABLE application_logs; Affected Rows: 0 +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +-- ADD COLUMN preserves JSON2 settings and type hints, including for existing rows. +insert into json2_alter_add_settings values (1); + +Affected Rows: 1 + +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +Affected Rows: 0 + +alter table json2_alter_add_settings add column k json2(service string); + +Affected Rows: 0 + +alter table json2_alter_add_settings add column m json2(max_auto_expanded_paths = 0, nested.value bigint); + +Affected Rows: 0 + +show create table json2_alter_add_settings; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| json2_alter_add_settings | CREATE TABLE IF NOT EXISTS "json2_alter_add_settings" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 100 | +| | ) NULL, | +| | "k" JSON2( | +| | max_auto_expanded_paths = 100, | +| | "service" STRING | +| | ) NULL, | +| | "m" JSON2( | +| | max_auto_expanded_paths = 0, | +| | "nested"."value" BIGINT | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +insert into json2_alter_add_settings values + (2, '{"a":1}', '{"service":"api"}', '{"nested":{"value":42},"extra":true}'); + +Affected Rows: 1 + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| ts | j | json_get(json2_alter_add_settings.k,Utf8("$.service"),Utf8View(NULL)) | json_get(json2_alter_add_settings.m,Utf8("$.nested.value"),Int64(NULL)) | m | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| 1970-01-01T00:00:00.001 | | | | | +| 1970-01-01T00:00:00.002 | {"a":1} | api | 42 | {"extra":true,"nested":{"value":42}} | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ + +admin flush_table('json2_alter_add_settings'); + ++-----------------------------------------------+ +| ADMIN flush_table('json2_alter_add_settings') | ++-----------------------------------------------+ +| 0 | ++-----------------------------------------------+ + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| ts | j | json_get(json2_alter_add_settings.k,Utf8("$.service"),Utf8View(NULL)) | json_get(json2_alter_add_settings.m,Utf8("$.nested.value"),Int64(NULL)) | m | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ +| 1970-01-01T00:00:00.001 | | | | | +| 1970-01-01T00:00:00.002 | {"a":1} | api | 42 | {"extra":true,"nested":{"value":42}} | ++-------------------------+---------+-----------------------------------------------------------------------+-------------------------------------------------------------------------+--------------------------------------+ + +drop table json2_alter_add_settings; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql index 457943e93f6..83b514f6895 100644 --- a/tests/cases/standalone/common/types/json/json2_alter.sql +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -74,3 +74,30 @@ ORDER BY ts; ADMIN FLUSH_TABLE('application_logs'); DROP TABLE application_logs; + +create table json2_alter_add_settings ( + ts timestamp time index +) with ( + 'append_mode' = 'true' +); + +-- ADD COLUMN preserves JSON2 settings and type hints, including for existing rows. +insert into json2_alter_add_settings values (1); +alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); + +alter table json2_alter_add_settings add column k json2(service string); + +alter table json2_alter_add_settings add column m json2(max_auto_expanded_paths = 0, nested.value bigint); + +show create table json2_alter_add_settings; + +insert into json2_alter_add_settings values + (2, '{"a":1}', '{"service":"api"}', '{"nested":{"value":42},"extra":true}'); + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + +admin flush_table('json2_alter_add_settings'); + +select ts, j, k.service, m.nested.value, m from json2_alter_add_settings order by ts; + +drop table json2_alter_add_settings; diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 32bbfbaf0a7..555e76f1e5e 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -95,27 +95,6 @@ drop table json2_alter_non_append; Affected Rows: 0 -create table json2_alter_add_settings ( - ts timestamp time index -) with ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - --- ALTER TABLE ADD COLUMN does not yet support JSON2 settings or type hints. -alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); - -Error: 2000(InvalidSyntax), Invalid SQL syntax: sql parser error: Expected: type modifiers, found: = at Line: 1, Column: 81 - -alter table json2_alter_add_settings add column k json2(service string); - -Error: 2000(InvalidSyntax), SQL data type not supported yet: Custom(ObjectName([Identifier(Ident { value: "json2", quote_style: None, span: Span(Location(1,51)..Location(1,56)) })]), ["service", "string"]) - -drop table json2_alter_add_settings; - -Affected Rows: 0 - create table json2_set_append_mode_false ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 1bbadfd5adc..653fdcec12c 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -58,19 +58,6 @@ alter table json2_alter_non_append add column j json2; drop table json2_alter_non_append; -create table json2_alter_add_settings ( - ts timestamp time index -) with ( - 'append_mode' = 'true' -); - --- ALTER TABLE ADD COLUMN does not yet support JSON2 settings or type hints. -alter table json2_alter_add_settings add column j json2(max_auto_expanded_paths = 100); - -alter table json2_alter_add_settings add column k json2(service string); - -drop table json2_alter_add_settings; - create table json2_set_append_mode_false ( ts timestamp time index, j json2 From 0d2cb1aa7e6668eb4a29a3d5772e3db072794a5b Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 9 Oct 2026 03:50:54 +0000 Subject: [PATCH 112/133] fix(promql): align subqueries to the absolute evaluation grid (#9400) * fix(promql): align subqueries to the absolute evaluation grid Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): streamline subquery comments Remove redundant arithmetic narration and restated test comments. Keep grid, fallback, validation and overflow contracts, and preserve all regression tests and assertions. Regenerate SQLness output only for comment echoes. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/error.rs | 11 + src/query/src/promql/planner.rs | 86 ++- src/query/src/promql/planner/test.rs | 599 ++++++++++++++++++ .../standalone/common/promql/subquery.result | 167 +++++ .../standalone/common/promql/subquery.sql | 63 ++ 5 files changed, 917 insertions(+), 9 deletions(-) diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index 3b24d33612d..121b5e81a83 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -142,6 +142,16 @@ pub enum Error { location: Location, }, + #[snafu(display( + "The subquery child window is outside the representable millisecond range: {}", + timestamp + ))] + SubqueryTimestampOutOfRange { + timestamp: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Cannot find column {col}"))] ColumnNotFound { col: String, @@ -244,6 +254,7 @@ impl ErrorExt for Error { | ExpectRangeSelector { .. } | ZeroRangeSelector { .. } | InvalidTimeRange { .. } + | SubqueryTimestampOutOfRange { .. } | ColumnNotFound { .. } | FunctionInvalidArgument { .. } | UnsupportedVectorMatch { .. } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 07faf3b53dd..2d52cfb83f7 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -48,7 +48,7 @@ use datafusion::functions_window::row_number::RowNumber; use datafusion::logical_expr::expr::{Alias, ScalarFunction, WindowFunction}; use datafusion::logical_expr::expr_rewriter::normalize_cols; use datafusion::logical_expr::{ - BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, + BinaryExpr, Cast, EmptyRelation, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; use datafusion::prelude as df_prelude; @@ -107,9 +107,10 @@ use crate::promql::error::{ CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, - MultipleVectorSnafu, NoMetricMatcherSnafu, Result, TableNameNotFoundSnafu, - TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, - UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, Result, SubqueryTimestampOutOfRangeSnafu, + TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, + UnknownTableSnafu, UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, + ZeroRangeSelectorSnafu, }; use crate::query_engine::QueryEngineState; @@ -410,9 +411,7 @@ impl PromPlanner { .. } = subquery_expr; - // Prometheus evaluates the inner expression over `(start - offset - range, end - offset]` - // (`subqueryTimeRange`). Shift the inner window back here; `RangeManipulate` maps the - // samples forward again by the same offset. + // Shift the child window back; `RangeManipulate` restores the outer timeline. let offset_ms = match offset { Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), @@ -423,14 +422,42 @@ impl PromPlanner { if let Some(step) = step { self.ctx.interval = step.as_millis() as _; } + ensure!(self.ctx.interval > 0, ZeroRangeSelectorSnafu); let current_start = self.ctx.start; let current_end = self.ctx.end; - self.ctx.start -= offset_ms + range.as_millis() as i64 - self.ctx.interval; - self.ctx.end -= offset_ms; + // Reject an invalid parent window rather than treating it as empty. + ensure!( + current_start <= current_end, + InvalidTimeRangeSnafu { + start: current_start, + end: current_end, + } + ); + // An empty child needs only its schema; plan it over the valid caller window. + let child_window = Self::subquery_child_window( + current_start, + current_end, + current_interval, + offset_ms, + range.as_millis() as Millisecond, + self.ctx.interval, + )?; + if let Some((child_start, child_end)) = child_window { + self.ctx.start = child_start; + self.ctx.end = child_end; + } let input = self.prom_expr_to_plan(expr, query_engine_state).await?; self.ctx.interval = current_interval; self.ctx.start = current_start; self.ctx.end = current_end; + // Return no rows with the child schema so enclosing fallbacks still work. + let input = match child_window { + Some(_) => input, + None => LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: input.schema().clone(), + }), + }; ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); let range_ms = range.as_millis() as _; @@ -1656,6 +1683,47 @@ impl PromPlanner { } } + /// Returns the first absolute step multiple strictly after `start - offset - range` + /// through the last parent evaluation instant minus `offset`. Requires a positive step. + /// Returns `None` for an empty window; an unrepresentable nonempty bound is an error. + fn subquery_child_window( + start: Millisecond, + end: Millisecond, + interval: Millisecond, + offset: Millisecond, + range: Millisecond, + step: Millisecond, + ) -> Result> { + let window_start = i128::from(start) - i128::from(offset) - i128::from(range); + let grid_start = (window_start.div_euclid(i128::from(step)) + 1) * i128::from(step); + let grid_end = if interval > 0 { + let steps = (i128::from(end) - i128::from(start)) / i128::from(interval); + i128::from(start) + steps * i128::from(interval) + } else { + i128::from(end) + } - i128::from(offset); + // Check emptiness before converting bounds to i64. + if grid_start > grid_end { + return Ok(None); + } + let bound = |name: &str, value: i128| { + i64::try_from(value).map_err(|_| { + SubqueryTimestampOutOfRangeSnafu { + timestamp: format!( + "subquery child window {name} = {value} ms (start = {start} ms, \ + end = {end} ms, interval = {interval} ms, offset = {offset} ms, \ + range = {range} ms, step = {step} ms)" + ), + } + .build() + }) + }; + Ok(Some(( + bound("grid start", grid_start)?, + bound("grid end", grid_end)?, + ))) + } + /// The columns that identify one series, which is the series key expected by the PromQL plan /// nodes that hold exactly one series per input batch. fn series_key_columns(&self) -> Vec { diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 215acb85cf9..642255532bd 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -5134,6 +5134,605 @@ async fn count_over_time_subquery_with_offset() { indie_query_plan_compare(query, expected).await; } +/// Time-varying samples on and off the subquery step grid expose incorrect window folding. +async fn build_subquery_offset_probe_provider() -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let probe_rows = (0..=6_i64) + .map(|step| (step * 10_000, (step + 1) as f64)) + .collect::>(); + let offgrid_rows = (0..6_i64) + .map(|step| (step * 10_000 + 5_000, (step + 1) as f64)) + .collect::>(); + + for (table_id, name, rows) in [ + (7_001_u32, "subquery_offset_probe", &probe_rows), + (7_002_u32, "subquery_offset_offgrid", &offgrid_rows), + ] { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "host".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "val".to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec!["a"; rows.len()])), + Arc::new(TimestampMillisecondArray::from_iter_values( + rows.iter().map(|(ts, _)| *ts), + )), + Arc::new(Float64Array::from_iter_values( + rows.iter().map(|(_, value)| *value), + )), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(3) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(table_meta) + .build() + .unwrap(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: name.to_string(), + table_id, + table, + }) + .is_ok() + ); + } + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + +/// Plans on a seconds-based evaluation grid. +async fn plan_subquery_offset_probe( + query: &str, + start: u64, + end: u64, + step: u64, +) -> Result { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: UNIX_EPOCH + Duration::from_secs(start), + end: UNIX_EPOCH + Duration::from_secs(end), + interval: Duration::from_secs(step), + lookback_delta: Duration::from_secs(300), + }; + PromPlanner::stmt_to_plan( + build_subquery_offset_probe_provider().await, + &eval_stmt, + &build_query_engine_state(), + ) + .await +} + +/// Evaluates seconds-based bounds and returns sorted `(timestamp_ms, value)` rows. +async fn run_subquery_offset_probe( + query: &str, + start: u64, + end: u64, + step: u64, +) -> Vec<(i64, f64)> { + let state = build_query_engine_state(); + let plan = plan_subquery_offset_probe(query, start, end, step) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + let mut rows = Vec::new(); + for batch in &batches { + let timestamp_index = batch + .schema() + .fields() + .iter() + .position(|field| matches!(field.data_type(), ArrowDataType::Timestamp(..))) + .expect("no timestamp column"); + let value_index = batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .expect("no Float64 value column"); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .expect("timestamp column is not a millisecond timestamp"); + let values = batch + .column(value_index) + .as_any() + .downcast_ref::() + .expect("value column is not Float64"); + rows.extend( + timestamps + .iter() + .zip(values.iter()) + .map(|(timestamp, value)| (timestamp.unwrap(), value.unwrap())), + ); + } + rows.sort_by_key(|(timestamp, _)| *timestamp); + rows +} + +#[tokio::test] +async fn subquery_own_offset_shifts_the_inner_window_but_keeps_evaluation_timestamps() { + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset 30s)", + 45, + 45, + 1 + ) + .await, + vec![(45_000, 3.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 15, 15, 1).await, + vec![(15_000, 3.0)], + ); + + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset 30s)", + 30, + 60, + 30 + ) + .await, + vec![(30_000, 1.0), (60_000, 7.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 0, 30, 30).await, + vec![(0, 1.0), (30_000, 7.0)], + ); + + // Negative offsets look ahead while retaining outer timestamps. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_probe[20s:10s] offset -30s)", + 0, + 30, + 30 + ) + .await, + vec![(0, 7.0), (30_000, 13.0)], + ); + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_probe[20s:10s])", 30, 60, 30) + .await, + vec![(30_000, 7.0), (60_000, 13.0)], + ); +} + +/// Expectations match Prometheus v3.14.0. +#[tokio::test] +async fn subquery_inner_points_are_anchored_on_absolute_step_multiples() { + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[20s:10s])", 45, 45, 1) + .await, + vec![(45_000, 7.0)], + ); + + // An off-grid offset distinguishes absolute anchoring from re-phasing. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 30, + 30, + 1 + ) + .await, + vec![(30_000, 3.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 45, + 45, + 1 + ) + .await, + vec![(45_000, 7.0)], + ); + + // The window starts before the first sample. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s)", + 30, + 30, + 1 + ) + .await, + vec![(30_000, 1.0)], + ); + + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s)", + 15, + 15, + 1 + ) + .await, + vec![(15_000, 3.0)], + ); + + // The negative grid start exercises Euclidean alignment. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[40s:10s])", 15, 15, 1) + .await, + vec![(15_000, 1.0)], + ); + + // Multi-step evaluation starts off the subquery grid. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s])", + 35, + 65, + 30 + ) + .await, + vec![(35_000, 5.0), (65_000, 11.0)], + ); + + // Child points reach no sample at the first evaluation. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s)", + 30, + 60, + 30 + ) + .await, + vec![(60_000, 5.0)], + ); +} + +#[test] +fn subquery_child_window_is_the_child_grid_window() { + let child_window = PromPlanner::subquery_child_window; + // Exclusive lower bound; upper bound is the last parent evaluation step. + assert_eq!( + child_window(30_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((20_000, 60_000)) + ); + assert_eq!( + child_window(45_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((30_000, 55_000)) + ); + assert_eq!( + child_window(25_000, 60_000, 10_000, 0, 20_000, 10_000).unwrap(), + Some((10_000, 55_000)) + ); + assert_eq!( + child_window(35_000, 65_000, 30_000, 0, 20_000, 10_000).unwrap(), + Some((20_000, 65_000)) + ); + assert_eq!( + child_window(0, 55_000, 30_000, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 30_000)) + ); + // A nonempty child window can contribute no parent rows. + assert_eq!( + child_window(35_000, 65_000, 10_000, 0, 5_000, 10_000).unwrap(), + Some((40_000, 65_000)) + ); + // The first child point is past the parent’s last evaluation. + assert_eq!( + child_window(35_000, 44_000, 10_000, 0, 5_000, 10_000).unwrap(), + None + ); + assert_eq!( + child_window(30_000, 30_000, 1_000, 5_000, 20_000, 10_000).unwrap(), + Some((10_000, 25_000)) + ); + assert_eq!( + child_window(15_000, 15_000, 1_000, -5_000, 20_000, 10_000).unwrap(), + Some((10_000, 20_000)) + ); + // Boundaries before the epoch floor towards minus infinity. + assert_eq!( + child_window(-15_000, -15_000, 1_000, 0, 20_000, 10_000).unwrap(), + Some((-30_000, -15_000)) + ); + assert_eq!( + child_window(-45_000, -45_000, 1_000, -10_000, 20_000, 10_000).unwrap(), + Some((-50_000, -35_000)) + ); + // A grid start before the first sample is not clamped to it. + assert_eq!( + child_window(0, 0, 1_000, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 0)) + ); + // A zero parent interval leaves the upper bound unaligned. + assert_eq!( + child_window(0, 55_000, 0, 0, 20_000, 10_000).unwrap(), + Some((-10_000, 55_000)) + ); + + // Only the derived start overflows. + let err = child_window(i64::MIN + 1_000, i64::MIN + 1_000, 1_000, 0, 2_000, 10).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") && timestamp.contains("-9223372036854776800") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Only the derived end overflows. + let err = child_window(i64::MAX - 1_000, i64::MAX - 1_000, 1, -2_000, 3_000, 10).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid end") && timestamp.contains("9223372036854776807") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // The wide window is nonempty even though its start overflows. + let err = child_window(i64::MIN, i64::MAX, 10_000, i64::MAX, i64::MAX, 1_000).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") + ), + "{err}" + ); + // Both derived bounds overflow; report the exact start without clipping. + let err = child_window(i64::MAX - 1_000, i64::MAX - 1_000, 1, i64::MIN, 10, 1).unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start") + && timestamp.contains("18446744073709550606") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Check emptiness before converting overflowing bounds. + assert_eq!( + child_window(i64::MAX, i64::MAX, 1_000, i64::MIN, 0, 1_000).unwrap(), + None + ); +} + +#[tokio::test] +async fn subquery_child_window_without_a_step_reports_no_row() { + // The first child point is after the parent’s only evaluation. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 35, 44, 10) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 1.0)], + ); + // Multi-step: every step of the grid reports no child point and falls back. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1)", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 1.0), (45_000, 1.0), (55_000, 1.0), (65_000, 1.0)], + ); + // A nonempty child grid can yield no parent rows; a longer range does. + assert_eq!( + run_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 35, 65, 10) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s])", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 5.0), (45_000, 7.0), (55_000, 9.0), (65_000, 11.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s)", + 35, + 65, + 10 + ) + .await, + vec![(35_000, 5.0), (45_000, 7.0), (55_000, 9.0), (65_000, 11.0)], + ); + // A negative offset makes the short window reach a child point. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 4.0)], + ); + + // An empty outer window suppresses inner results, but not an outer fallback. + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s])", + 35, + 44, + 10 + ) + .await, + vec![], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1)", + 35, + 44, + 10 + ) + .await, + vec![(35_000, 1.0)], + ); + assert_eq!( + run_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1))[5s:10s])", + 35, + 44, + 10 + ) + .await, + vec![], + ); +} + +#[tokio::test] +async fn subquery_rejects_invalid_range_and_zero_step() { + let err = + plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:10s])", 44, 35, 10) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::InvalidTimeRange { + start: 44_000, + end: 35_000, + .. + } + ), + "{err}" + ); + let err = plan_subquery_offset_probe( + "sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s])", + 44, + 35, + 10, + ) + .await + .unwrap_err(); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + + // `[5s:0]` parses to zero; planning rejects it. + let err = plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:0])", 60, 60, 1) + .await + .unwrap_err(); + assert!( + matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), + "{err}" + ); + // An omitted child step inherits the zero parent step. + let err = plan_subquery_offset_probe("sum_over_time(subquery_offset_offgrid[5s:])", 60, 60, 0) + .await + .unwrap_err(); + assert!( + matches!(&err, crate::promql::error::Error::ZeroRangeSelector { .. }), + "{err}" + ); +} + +/// Rejects nonempty child windows whose derived millisecond bounds are unrepresentable. +#[tokio::test] +async fn subquery_rejects_a_child_window_out_of_the_representable_range() { + // Offset is i64::MAX milliseconds; only the derived start overflows. + let err = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[21366776s:10s] offset 292471208y21366775s807ms)", + 0, + 60, + 10, + ) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid start = -9223372058221550000") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); + // Offset is -i64::MAX milliseconds; only the derived end overflows. + let err = plan_subquery_offset_probe( + "sum_over_time(subquery_offset_offgrid[20s:10s] offset -292471208y21366775s807ms)", + 0, + 60, + 10, + ) + .await + .unwrap_err(); + assert!( + matches!( + &err, + crate::promql::error::Error::SubqueryTimestampOutOfRange { timestamp, .. } + if timestamp.contains("grid end = 9223372036854835807") + ), + "{err}" + ); + assert_eq!(err.status_code(), StatusCode::InvalidArguments, "{err}"); +} + #[tokio::test] async fn test_hash_join() { let mut eval_stmt = EvalStmt { diff --git a/tests/cases/standalone/common/promql/subquery.result b/tests/cases/standalone/common/promql/subquery.result index 61f09258a41..1be8073a541 100644 --- a/tests/cases/standalone/common/promql/subquery.result +++ b/tests/cases/standalone/common/promql/subquery.result @@ -159,7 +159,174 @@ tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30 | 1970-01-01T00:00:30 | 4.0 | a | +---------------------+----------------------------------------------+------+ +-- Offset results match the shifted evaluation while retaining outer timestamps. +-- Verified against Prometheus v3.14.0. +tql eval (30, 60, '30s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 1.0 | a | +| 1970-01-01T00:01:00 | 7.0 | a | ++---------------------+----------------------------------+------+ + +tql eval (0, 30, '30s') sum_over_time(subquery_offset_total[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:30 | 7.0 | a | ++---------------------+----------------------------------+------+ + drop table subquery_offset_total; Affected Rows: 0 +-- Off-grid samples expose absolute step anchoring. +create table subquery_offset_offgrid ( + ts timestamp time index, + host string primary key, + val double, +); + +Affected Rows: 0 + +insert into subquery_offset_offgrid values + (5000, 'a', 1), + (15000, 'a', 2), + (25000, 'a', 3), + (35000, 'a', 4), + (45000, 'a', 5), + (55000, 'a', 6); + +Affected Rows: 6 + +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- Offset moves the window without re-phasing its step points. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 5.0 | a | ++---------------------+----------------------------------+------+ + +-- The window starts before the first sample. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 1.0 | a | ++---------------------+----------------------------------+------+ + +tql eval (15, 15, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:15 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- Multi-step evaluation uses the non-aligned parent timeline. +tql eval (35, 65, '30s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 5.0 | a | +| 1970-01-01T00:01:05 | 11.0 | a | ++---------------------+----------------------------------+------+ + +-- An evaluation instant whose window reaches no sample has no row (30s), unlike 60s -> 2 + 3 +tql eval (30, 60, '30s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:01:00 | 5.0 | a | ++---------------------+----------------------------------+------+ + +-- No child point yields no row; the enclosing fallback still fires. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + ++----+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++----+----------------------------------+------+ ++----+----------------------------------+------+ + +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1); + ++---------------------+------+----------------------------------+ +| ts | host | prom_sum_over_time(ts_range,val) | ++---------------------+------+----------------------------------+ +| 1970-01-01T00:00:35 | | 1.0 | ++---------------------+------+----------------------------------+ + +-- A negative offset makes the short window reach a child point. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 4.0 | a | ++---------------------+----------------------------------+------+ + +-- An empty outer window suppresses inner results, but not an outer fallback. +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]); + ++----+---------------------------------------------------------------+------+ +| ts | prom_sum_over_time(ts_range,prom_sum_over_time(ts_range,val)) | host | ++----+---------------------------------------------------------------+------+ ++----+---------------------------------------------------------------+------+ + +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1); + ++---------------------+------+---------------------------------------------------------------+ +| ts | host | prom_sum_over_time(ts_range,prom_sum_over_time(ts_range,val)) | ++---------------------+------+---------------------------------------------------------------+ +| 1970-01-01T00:00:35 | | 1.0 | ++---------------------+------+---------------------------------------------------------------+ + +-- An aligned instant has a child point in the same short range. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[5s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:30 | 3.0 | a | ++---------------------+----------------------------------+------+ + +-- A longer range covers child steps on the off-grid parent timeline. +tql eval (35, 65, '10s') sum_over_time(subquery_offset_offgrid[20s:10s]); + ++---------------------+----------------------------------+------+ +| ts | prom_sum_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:35 | 5.0 | a | +| 1970-01-01T00:00:45 | 7.0 | a | +| 1970-01-01T00:00:55 | 9.0 | a | +| 1970-01-01T00:01:05 | 11.0 | a | ++---------------------+----------------------------------+------+ + +-- Zero step parses here but is invalid for planning. +tql eval (60, 60, '1s') sum_over_time(subquery_offset_offgrid[5s:0]); + +Error: 1004(InvalidArguments), Zero range in range selector + +tql eval (44, 35, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +Error: 1004(InvalidArguments), The end time must be greater than start time, start: 44000, end: 35000 + +drop table subquery_offset_offgrid; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/subquery.sql b/tests/cases/standalone/common/promql/subquery.sql index 598aa91a06e..c44e5163b8d 100644 --- a/tests/cases/standalone/common/promql/subquery.sql +++ b/tests/cases/standalone/common/promql/subquery.sql @@ -65,4 +65,67 @@ tql eval (60, 60, '1s') predict_linear(subquery_offset_total[20s:10s] offset 30s tql eval (30, 30, '1s') predict_linear(subquery_offset_total[20s:10s] offset -30s, 0); +-- Offset results match the shifted evaluation while retaining outer timestamps. +-- Verified against Prometheus v3.14.0. +tql eval (30, 60, '30s') sum_over_time(subquery_offset_total[20s:10s] offset 30s); + +tql eval (0, 30, '30s') sum_over_time(subquery_offset_total[20s:10s]); + drop table subquery_offset_total; + +-- Off-grid samples expose absolute step anchoring. +create table subquery_offset_offgrid ( + ts timestamp time index, + host string primary key, + val double, +); + +insert into subquery_offset_offgrid values + (5000, 'a', 1), + (15000, 'a', 2), + (25000, 'a', 3), + (35000, 'a', 4), + (45000, 'a', 5), + (55000, 'a', 6); + +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 5s); + +-- Offset moves the window without re-phasing its step points. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- The window starts before the first sample. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 15s); + +tql eval (15, 15, '1s') sum_over_time(subquery_offset_offgrid[20s:10s] offset -5s); + +-- Multi-step evaluation uses the non-aligned parent timeline. +tql eval (35, 65, '30s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- An evaluation instant whose window reaches no sample has no row (30s), unlike 60s -> 2 + 3 +tql eval (30, 60, '30s') sum_over_time(subquery_offset_offgrid[20s:10s] offset 30s); + +-- No child point yields no row; the enclosing fallback still fires. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]) or vector(1); + +-- A negative offset makes the short window reach a child point. +tql eval (35, 44, '10s') sum_over_time(subquery_offset_offgrid[5s:10s] offset -5s); + +-- An empty outer window suppresses inner results, but not an outer fallback. +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]); + +tql eval (35, 44, '10s') sum_over_time((sum_over_time(subquery_offset_offgrid[5s:10s]))[5s:10s]) or vector(1); + +-- An aligned instant has a child point in the same short range. +tql eval (30, 30, '1s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +-- A longer range covers child steps on the off-grid parent timeline. +tql eval (35, 65, '10s') sum_over_time(subquery_offset_offgrid[20s:10s]); + +-- Zero step parses here but is invalid for planning. +tql eval (60, 60, '1s') sum_over_time(subquery_offset_offgrid[5s:0]); + +tql eval (44, 35, '10s') sum_over_time(subquery_offset_offgrid[5s:10s]); + +drop table subquery_offset_offgrid; From c8ea00b13f70712a85214b445856c8ae0e03270b Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 9 Oct 2026 06:07:27 +0000 Subject: [PATCH 113/133] feat(log-store): add object store WAL metrics and the conditional-create fault matrix (#9475) * feat(log-store): add object store WAL metrics and the conditional-create fault matrix Export the remaining fifteen metrics of the object store WAL under greptime_logstore_object_store_wal_: the write path timers and object shape histograms, created objects, the indexed object and byte gauges, recovery and read timers, recovered objects, create conflicts and failures, poisoned stores, and stalled appends with their wait. Add a table-driven fault matrix that drives every outcome of a conditional create and of the start object through the store and asserts the failed waiters, the poison, the sequence, both indexed gauges at the failure, the movement of every metric, and what a read sees afterwards. The test fails if a row of the failure matrix has no case. Every test of the store module now holds one lock, since the metrics are process-wide. Signed-off-by: jeremyhi * fix(log-store): count every conditional create of the object store WAL Count the outcome of the start object's create, and count a create that met a different object as a conflict before the epoch of that object is read, so a failed epoch read no longer hides the conflict. The fault matrix waits until no store actor of an earlier test is alive, checks that a poisoned read returns the recorded terminal error, and covers a conflicting object whose header cannot be read. Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/log-store/src/metrics.rs | 167 ++ src/log-store/src/object_store_wal/format.rs | 2 - src/log-store/src/object_store_wal/store.rs | 1421 +++++++++++++++++- 3 files changed, 1548 insertions(+), 42 deletions(-) diff --git a/src/log-store/src/metrics.rs b/src/log-store/src/metrics.rs index badb5a80b70..03a3dc1f174 100644 --- a/src/log-store/src/metrics.rs +++ b/src/log-store/src/metrics.rs @@ -111,4 +111,171 @@ lazy_static! { "object store logstore failed object deletes total", ) .unwrap(); + + /// How long an object of the object store logstore takes to become durable + /// once its batch is sealed, sampled for every conditional create that + /// succeeds, an identical retry included. The time runs from the seal, so + /// it includes any wait for a create slot and, in the `enqueued` mode, + /// every repeated attempt and the delay before it. + /// + /// A create is a few milliseconds on a local object store and tens of + /// milliseconds on a cloud one, so the buckets start at 1ms and reach 8s. + pub static ref METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_seal_to_durable_seconds", + "object store logstore seconds from sealing a batch to the success of its conditional create, slot waits and retries included", + exponential_buckets(0.001, 2.0, 14).unwrap(), + ) + .unwrap(); + + /// How long an append of the object store logstore waits for its + /// acknowledgement in the `durable` mode, which is the write latency a + /// caller sees. + /// + /// The wait is one seal plus every earlier object plus one create, so with + /// the default 100ms flush interval it sits in the tens of milliseconds; + /// the buckets straddle that and leave room for a stalled uploader. + pub static ref METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_append_ack_seconds", + "object store logstore seconds an append waits for its acknowledgement in the durable mode", + exponential_buckets(0.002, 2.0, 14).unwrap(), + ) + .unwrap(); + + /// How large the objects of the object store logstore are, which says + /// whether batches seal on size or on the flush interval. + /// + /// The buckets are powers of two from 1KiB, the size of a timer-sealed + /// batch of one small entry, to 16MiB, twice the 8MiB default batch size, + /// which is a threshold rather than an upper bound. + pub static ref METRIC_OBJECT_STORE_WAL_OBJECT_BYTES: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_object_bytes", + "object store logstore bytes of each object sealed", + exponential_buckets(1024.0, 2.0, 15).unwrap(), + ) + .unwrap(); + + /// How many entries the objects of the object store logstore hold, which + /// says how many writes one request amortises. + /// + /// The largest finite bucket is 2^20, the position limit of one region in + /// one object. An object holds the entries of many regions, so its total + /// can pass that; such an object falls in the overflow bucket. + pub static ref METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_object_entries", + "object store logstore entries of each object sealed", + exponential_buckets(1.0, 4.0, 11).unwrap(), + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore that + /// succeeded, an identical retry that found its own object included. A + /// create that failed is counted in the create failure or conflict counter + /// instead. + pub static ref METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_created_objects_total", + "object store logstore conditional creates that succeeded, identical retries included", + ) + .unwrap(); + + /// Number of objects the object store logstore has indexed, which is what + /// a recovery of the prefix would have to read. + pub static ref METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS: IntGauge = register_int_gauge!( + "greptime_logstore_object_store_wal_indexed_objects", + "object store logstore objects currently indexed", + ) + .unwrap(); + + /// Bytes of the objects the object store logstore has indexed, which is + /// what the WAL occupies under its prefix. + pub static ref METRIC_OBJECT_STORE_WAL_INDEXED_BYTES: IntGauge = register_int_gauge!( + "greptime_logstore_object_store_wal_indexed_bytes", + "object store logstore bytes currently indexed", + ) + .unwrap(); + + /// How long a recovery of the object store logstore that succeeded took to + /// rebuild its catalog, which is the part of a restart the WAL is + /// responsible for. A recovery that was abandoned records nothing. + /// + /// Recovery reads a few small ranges per object, so the buckets run from + /// 10ms, an empty prefix, to minutes on a prefix with many objects. + pub static ref METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_recovery_seconds", + "object store logstore seconds spent by a recovery that succeeded", + exponential_buckets(0.01, 3.0, 10).unwrap(), + ) + .unwrap(); + + /// Counter of objects a recovery of the object store logstore that + /// succeeded read footers from, which is what its time is spent on. It + /// counts the objects of the same recoveries the timer above measures, so + /// an abandoned one contributes to neither. + pub static ref METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_recovered_objects_total", + "object store logstore objects read by a recovery that succeeded total", + ) + .unwrap(); + + /// How long a read stream of the object store logstore took to serve the + /// replay of one region, recorded once the stream has been read to its + /// end; one that fails or is dropped part way records nothing. + /// + /// A read fetches one segment per object of the region, so the buckets run + /// from 5ms, a region with nothing to replay, to minutes. + pub static ref METRIC_OBJECT_STORE_WAL_READ_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_read_seconds", + "object store logstore seconds spent serving a region read that ran to its end", + exponential_buckets(0.005, 3.0, 10).unwrap(), + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore that found + /// an object with different content. In the `durable` mode an object of an + /// earlier epoch fails the batch like a transient failure; any other + /// conflict poisons the store. + pub static ref METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_create_conflicts_total", + "object store logstore conflicting conditional creates total", + ) + .unwrap(); + + /// Counter of conditional creates of the object store logstore the object + /// store did not confirm, which fail the appends behind them in the + /// `durable` mode and are repeated in the `enqueued` mode. + pub static ref METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_create_failures_total", + "object store logstore transient conditional create failures total", + ) + .unwrap(); + + /// Counter of object store logstores that hit an error they cannot recover + /// from, after which every operation of the store fails until it reopens. + pub static ref METRIC_OBJECT_STORE_WAL_POISONED_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_poisoned_total", + "object store logstores poisoned total", + ) + .unwrap(); + + /// Counter of appends the object store logstore held back because the + /// unpersisted backlog reached a threshold in the `enqueued` mode, which + /// says the write rate is above what the uploads sustain. + pub static ref METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL: IntCounter = register_int_counter!( + "greptime_logstore_object_store_wal_stalled_appends_total", + "object store logstore appends stalled by the backlog thresholds total", + ) + .unwrap(); + + /// How long such an append stays held back, until it is admitted or until + /// it is refused, a stop, a terminal error or the teardown of the store + /// releasing it included, which is the latency the `enqueued` mode pays + /// once the backlog is at a threshold. + /// + /// A stalled append waits for one upload to complete, so the buckets run + /// from 5ms to half a minute of a struggling object store. + pub static ref METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS: Histogram = register_histogram!( + "greptime_logstore_object_store_wal_stalled_append_seconds", + "object store logstore seconds a stalled append is held back before it is admitted or refused", + exponential_buckets(0.005, 3.0, 9).unwrap(), + ) + .unwrap(); } diff --git a/src/log-store/src/object_store_wal/format.rs b/src/log-store/src/object_store_wal/format.rs index bd7e9ae6410..7a255711784 100644 --- a/src/log-store/src/object_store_wal/format.rs +++ b/src/log-store/src/object_store_wal/format.rs @@ -43,7 +43,6 @@ const SEGMENT_HEADER_LEN: usize = 8 + 4; /// offset, segment length and segment CRC32. pub(crate) const FOOTER_ENTRY_LEN: usize = 8 + 8 + 8 + 4 + 8 + 8 + 4; /// Length of the entry count the footer starts with. -#[cfg(test)] const FOOTER_COUNT_LEN: usize = 4; /// Header of a WAL object. @@ -191,7 +190,6 @@ pub(crate) fn encode_object(header: Header, records: &[Record]) -> Result u64 { let segments = footer.iter().map(|entry| entry.segment_len).sum::(); let framing = HEADER_LEN + FOOTER_COUNT_LEN + footer.len() * FOOTER_ENTRY_LEN + TRAILER_LEN; diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index dd3c3958336..d3a3acc6968 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -46,15 +46,23 @@ use crate::error::{ WalObjectSequenceUnsettledSnafu, }; use crate::metrics::{ + METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL, + METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL, METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL, - METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES, METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS, + METRIC_OBJECT_STORE_WAL_OBJECT_BYTES, METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES, + METRIC_OBJECT_STORE_WAL_POISONED_TOTAL, METRIC_OBJECT_STORE_WAL_READ_SECONDS, + METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL, METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS, + METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, + METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL, METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS, + METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL, }; use crate::object_store_wal::batch::{OBJECT_SEQ_LIMIT, OpenBatch, entry_id, sequence_floor}; use crate::object_store_wal::catalog::ObjectCatalog; use crate::object_store_wal::format::{ ChainLink, EncodedObject, FixedTrailer, FooterEntry, HEADER_LEN, Header, MIN_OBJECT_LEN, Record, TRAILER_LEN, decode_footer, decode_header, decode_segment, decode_trailer, - encode_object, footer_range, verify_segment_ranges, + encode_object, footer_range, object_len, verify_segment_ranges, }; use crate::object_store_wal::io::{ListedObject, ObjectStoreIo, PutResult}; @@ -259,6 +267,7 @@ impl ObjectStoreLogStore { reason: "max unpersisted age is zero", } ); + let recovery_started_at = Instant::now(); let Recovered { mut catalog, next_object_seq, @@ -267,6 +276,7 @@ impl ObjectStoreLogStore { max_epoch, reclaim_boundary, unindexed, + recovered_objects, } = recover(io.as_ref()).await?; // Every epoch is one above the sequence of the start object its // instance created, and every object is at or above its start object, @@ -279,6 +289,12 @@ impl ObjectStoreLogStore { ), } ); + // Only a recovery that succeeded is reported, so that the objects + // counted and the time measured are those of the same recoveries. + METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL.inc_by(recovered_objects); + METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS + .observe(recovery_started_at.elapsed().as_secs_f64()); + publish_indexed(&catalog); let (start, moved_past) = start_epoch(io.as_ref(), next_object_seq, tip, reclaim_boundary).await?; let epoch = start.epoch; @@ -292,6 +308,8 @@ impl ObjectStoreLogStore { .with_context(|_| InvalidWalObjectSnafu { path: io.object_path(start.object_seq), })?; + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(object_len(&[]) as i64); // No watermark is recorded yet. let writer_reclaim_boundary = catalog.reclaim_boundary(reclaim_boundary, &HashMap::new()); let catalog = Arc::new(RwLock::new(catalog)); @@ -363,6 +381,8 @@ impl ObjectStoreLogStore { #[cfg(any(test, feature = "testing"))] durability_waits: durability_waits_tx, }; + #[cfg(test)] + LIVE_ACTORS.fetch_add(1, Ordering::SeqCst); common_runtime::spawn_global(actor.run()); Ok(Arc::new(Self { prefix, @@ -680,13 +700,20 @@ impl LogStore for ObjectStoreLogStore { ensure!(entry.is_complete(), IncompleteWalEntrySnafu { region_id }); } + let requested_at = Instant::now(); let (response_tx, response_rx) = oneshot::channel(); self.append_tx .send((entries, response_tx)) .await .ok() .context(ObjectStoreWalStoppedSnafu)?; - response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)? + let response = response_rx.await.ok().context(ObjectStoreWalStoppedSnafu)?; + // Only an acknowledgement is timed; an append that failed never got one. + if self.ack_mode == AckMode::Durable && response.is_ok() { + METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS + .observe(requested_at.elapsed().as_secs_f64()); + } + response } /// Reads the provider's region from `entry_id`, hiding obsolete entries. @@ -734,6 +761,7 @@ impl LogStore for ObjectStoreLogStore { let on_corrupted_segment = self.on_corrupted_segment; let wal_holes = self.wal_holes.clone(); Ok(Box::pin(try_stream! { + let started_at = Instant::now(); for (object_seq, footer_entry) in objects { let records = match fetch_segment(io.as_ref(), object_seq, &footer_entry).await { Ok(records) => records, @@ -769,6 +797,7 @@ impl LogStore for ObjectStoreLogStore { yield entries; } } + METRIC_OBJECT_STORE_WAL_READ_SECONDS.observe(started_at.elapsed().as_secs_f64()); })) } @@ -974,6 +1003,7 @@ struct SealedBatch { bytes: Bytes, footer: Vec, first_admitted_at: Instant, + sealed_at: Instant, /// Number of creates that were attempted for the batch. attempts: u32, waiters: Vec, @@ -1066,8 +1096,9 @@ struct Actor { /// Creates in flight, including those of batches that already failed. creates: FuturesUnordered>, /// The append held back in the `enqueued` mode while the unpersisted - /// backlog is at a threshold. Later appends wait in their channel. - stalled: Option, + /// backlog is at a threshold, and when it was held back. Later appends + /// wait in their channel. + stalled: Option<(QueuedAppend, Instant)>, durable_waiters: Vec, /// Callers of `stop`, answered once nothing is in flight. stop: Vec>>, @@ -1184,7 +1215,8 @@ impl Actor { return; } if self.ack_mode == AckMode::Enqueued && self.backlog_at_threshold() { - self.stalled = Some((entries, response)); + METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.inc(); + self.stalled = Some(((entries, response), Instant::now())); self.ensure_create_in_flight(); return; } @@ -1282,7 +1314,7 @@ impl Actor { return; } if self.is_stopped() { - if let Some((_, response)) = self.stalled.take() { + if let Some((_, response)) = self.take_stalled() { let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); } return; @@ -1295,11 +1327,19 @@ impl Actor { if self.sealed.len() + 2 > MAX_SEALED_BATCHES { return; } - if let Some((entries, response)) = self.stalled.take() { + if let Some((entries, response)) = self.take_stalled() { self.admit(entries, response); } } + /// Takes the held-back append out on whichever path releases it and + /// records how long it was held back. + fn take_stalled(&mut self) -> Option { + let (append, stalled_at) = self.stalled.take()?; + METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.observe(stalled_at.elapsed().as_secs_f64()); + Some(append) + } + /// Seals the open batch as the object `next_object_seq` and starts its /// create when a slot is free. Returns whether a batch was sealed. fn flush_open_batch(&mut self) -> bool { @@ -1321,6 +1361,7 @@ impl Actor { }; let (entries, first_admitted_at) = self.open_batch.seal(); + let entry_count = entries.len(); let header = Header { object_seq, epoch: self.epoch, @@ -1341,6 +1382,8 @@ impl Actor { return false; } }; + METRIC_OBJECT_STORE_WAL_OBJECT_BYTES.observe(encoded.bytes.len() as f64); + METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES.observe(entry_count as f64); self.next_object_seq = object_seq .checked_add(1) .filter(|next_object_seq| *next_object_seq < OBJECT_SEQ_LIMIT); @@ -1361,6 +1404,7 @@ impl Actor { bytes: encoded.bytes, footer: encoded.footer, first_admitted_at, + sealed_at: Instant::now(), attempts: 0, waiters: std::mem::take(&mut self.pending), #[cfg(any(test, feature = "testing"))] @@ -1397,6 +1441,7 @@ impl Actor { batch.attempts += 1; let io = self.io.clone(); let object_seq = batch.object_seq; + let sealed_at = batch.sealed_at; self.creating.insert(object_seq); let bytes = batch.bytes.clone(); let epoch = self.epoch; @@ -1425,11 +1470,24 @@ impl Actor { } #[cfg(any(test, feature = "testing"))] if creates_fail.load(Ordering::Acquire) { - return (object_seq, injected_create_failure(io.as_ref(), object_seq)); + let result = injected_create_failure(io.as_ref(), object_seq); + record_create(&result, Some(sealed_at)); + return (object_seq, result); } + let created = io.put_if_absent(object_seq, bytes).await; + #[cfg(any(test, feature = "testing"))] + let created = match created { + Ok(_) if next_create_fails_after_write.swap(false, Ordering::AcqRel) => { + injected_create_failure(io.as_ref(), object_seq) + } + created => created, + }; + // Counted as the object store answered, before the epoch of a + // conflicting object is read. + record_create(&created, Some(sealed_at)); // Any conflict poisons an `enqueued` store, so only the // `durable` mode reads the epoch of the existing object. - let result = match io.put_if_absent(object_seq, bytes).await { + let result = match created { Err(error @ Error::WalObjectConflict { .. }) if ack_mode == AckMode::Durable => { @@ -1437,10 +1495,6 @@ impl Actor { } result => result, }; - #[cfg(any(test, feature = "testing"))] - if result.is_ok() && next_create_fails_after_write.swap(false, Ordering::AcqRel) { - return (object_seq, injected_create_failure(io.as_ref(), object_seq)); - } (object_seq, result) })); } @@ -1524,6 +1578,8 @@ impl Actor { let Some(batch) = self.sealed.pop_front() else { return false; }; + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.inc(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.add(batch.bytes.len() as i64); self.last_indexed = ChainLink { object_seq: batch.object_seq, epoch: self.epoch, @@ -1653,7 +1709,7 @@ impl Actor { for pending in self.pending.drain(..) { let _ = pending.response.send(Err(error())); } - if let Some((_, response)) = self.stalled.take() { + if let Some((_, response)) = self.take_stalled() { let _ = response.send(Err(error())); } for waiter in self.durable_waiters.drain(..) { @@ -1787,6 +1843,10 @@ impl Actor { self.unindexed.remove(&object_seq); let mut catalog = self.catalog.write().unwrap_or_else(PoisonError::into_inner); let footer = catalog.remove_object(object_seq); + if let Some(footer) = &footer { + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.dec(); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.sub(object_len(footer) as i64); + } let mut wal_holes = self .wal_holes .lock() @@ -1888,7 +1948,7 @@ impl Actor { } } AckMode::Enqueued => { - if let Some((_, response)) = self.stalled.take() { + if let Some((_, response)) = self.take_stalled() { let _ = response.send(Err(ObjectStoreWalStoppedSnafu.build())); } self.flush_open_batch(); @@ -1957,12 +2017,21 @@ impl Actor { impl Drop for Actor { /// Settles the deletes the actor still had in flight, so no read waits - /// for an attempt that will never complete. + /// for an attempt that will never complete, and releases the append still + /// held back, whose caller is gone or about to be. fn drop(&mut self) { self.deleting.abandon_all(); + self.take_stalled(); + #[cfg(test)] + LIVE_ACTORS.fetch_sub(1, Ordering::SeqCst); } } +/// Actors not dropped yet. An actor outlives the test that dropped its store +/// and may still move the process-wide metrics. +#[cfg(test)] +static LIVE_ACTORS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + fn terminal(terminal_error: &TerminalError) -> Option> { terminal_error .lock() @@ -1973,13 +2042,49 @@ fn terminal(terminal_error: &TerminalError) -> Option> { /// Records `error` as the terminal error unless one is already recorded, and /// returns the recorded one. fn set_terminal(terminal_error: &TerminalError, error: Error) -> Arc { - terminal_error + let mut terminal_error = terminal_error .lock() - .unwrap_or_else(PoisonError::into_inner) + .unwrap_or_else(PoisonError::into_inner); + if terminal_error.is_none() { + METRIC_OBJECT_STORE_WAL_POISONED_TOTAL.inc(); + } + terminal_error .get_or_insert_with(|| Arc::new(error)) .clone() } +/// Counts the outcome of a conditional create and, for the object of a batch +/// sealed at `sealed_at`, times a success. The create of a batch is counted +/// whether or not the batch is still there to be indexed: an object a +/// poisoned store gave up on is as durable as any other. +fn record_create(result: &Result, sealed_at: Option) { + match result { + Ok(_) => { + if let Some(sealed_at) = sealed_at { + METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS + .observe(sealed_at.elapsed().as_secs_f64()); + } + METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL.inc(); + } + Err(Error::WalObjectStore { .. }) => METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL.inc(), + Err(Error::WalObjectConflict { .. }) => { + METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL.inc() + } + Err(_) => {} + } +} + +/// Sets both indexed gauges to the objects `catalog` holds. +fn publish_indexed(catalog: &ObjectCatalog) { + let (objects, bytes) = catalog + .objects_in_order() + .fold((0, 0), |(objects, bytes), (_, footer)| { + (objects + 1, bytes + object_len(footer)) + }); + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.set(objects); + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.set(bytes as i64); +} + /// Returns true for a storage error that a later attempt may not meet: one /// the object store reports as temporary, or as persistent, which is how its /// retry layer reports a temporary error that outlasted its retries. @@ -2104,6 +2209,8 @@ struct Recovered { reclaim_boundary: u64, /// Present objects that are not indexed, in sequence order. unindexed: Vec, + /// Number of objects whose header and footer were read. + recovered_objects: u64, } /// The objects recovery fetched, ordered by sequence, and the listed objects @@ -2141,6 +2248,7 @@ async fn recover(io: &dyn WalObjectIo) -> Result { /// the error its fetch met: collection never deletes such an object. fn finish_recovery(fetched: FetchedObjects) -> Result { let FetchedObjects { objects, vanished } = fetched; + let recovered_objects = objects.len() as u64; let headers = objects .iter() .map(|fetched| (fetched.object.object_seq, &fetched.header)) @@ -2215,6 +2323,7 @@ fn finish_recovery(fetched: FetchedObjects) -> Result { max_epoch, reclaim_boundary, unindexed, + recovered_objects, }) } @@ -2301,10 +2410,11 @@ async fn start_epoch( prev: tip, reclaim_boundary, }; - match io + let created = io .put_if_absent(object_seq, encode_object(header, &[])?.bytes) - .await - { + .await; + record_create(&created, None); + match created { Ok(PutResult::Created) => return Ok((ChainLink { object_seq, epoch }, moved_past)), Ok(PutResult::AlreadyPresent) => { return UnconfirmedWalEpochStartSnafu { @@ -2566,13 +2676,10 @@ mod tests { const PREFIX: &str = "wal/datanodes/1/epochs/2"; const WAIT: Duration = Duration::from_secs(30); - /// Held by every test that reads the process-wide skipped segments - /// counter, so that tests running in one process do not interleave - /// their increments. - static SKIPPED_SEGMENTS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); - /// Held by every test that fails a delete, so that a test reading the - /// process-wide failed deletes counter sees only its own failures. - static FAILED_DELETES: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + /// Held by every test of the module: the metrics of the store are + /// process-wide, so tests running in one process would otherwise move the + /// counters and gauges that another test asserts on. + static SERIALIZED: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); fn memory_store() -> ObjectStore { ObjectStore::new(Memory::default()).unwrap() @@ -2770,6 +2877,7 @@ mod tests { #[tokio::test] async fn test_store_rejects_invalid_config() { + let _serialized = SERIALIZED.lock().await; for config in [ config(Duration::from_millis(9), 1), config(Duration::from_secs(1), 0), @@ -2799,6 +2907,7 @@ mod tests { #[tokio::test] async fn test_store_recovers_only_its_node_and_generation() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let config = ObjectStoreWalConfig::default(); let identities = [(1, 2), (3, 2), (1, 4)]; @@ -3033,6 +3142,7 @@ mod tests { #[tokio::test] async fn test_store_footer_recovery_matches_full_decode() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 40, 5).await; let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); @@ -3060,6 +3170,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_corrupted_trailer_version_and_footer() { + let _serialized = SERIALIZED.lock().await; type Corrupt = fn(&mut Vec); let cases: &[(&str, Corrupt)] = &[ ("truncated object", |bytes| { @@ -3098,6 +3209,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_header_sequence_mismatch() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let encoded = encode_object( @@ -3134,6 +3246,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_segment_ranges_that_do_not_tile_the_object() { + let _serialized = SERIALIZED.lock().await; type Corrupt = fn(&mut Vec, &[FooterEntry]); let cases: [(&str, Corrupt); 5] = [ ("overflows the object", |bytes, footer| { @@ -3177,6 +3290,7 @@ mod tests { } #[tokio::test] async fn test_store_recovery_fetches_a_footer_longer_than_the_tail_window() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let regions = (RECOVERY_TAIL_WINDOW / FOOTER_ENTRY_LEN + 100) as u32; let records = (1..=regions) @@ -3254,6 +3368,7 @@ mod tests { #[tokio::test] async fn test_store_resumes_sequence_and_durable_ids() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let recovered = recover(&io).await.unwrap(); @@ -3318,6 +3433,7 @@ mod tests { #[tokio::test] async fn test_store_rejects_exhausted_sequence() { + let _serialized = SERIALIZED.lock().await; for (object_seq, entry_id) in [(u64::MAX, 1), (OBJECT_SEQ_LIMIT - 1, 1), (0, u64::MAX)] { let object_store = memory_store(); put_object(&object_store, object_seq, region(1), &[entry_id]).await; @@ -3333,6 +3449,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_conflicting_entry_ranges() { + let _serialized = SERIALIZED.lock().await; for entries in [[1, 2], [2, 3]] { let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1, 2]).await; @@ -3346,6 +3463,7 @@ mod tests { #[tokio::test] async fn test_store_rejects_foreign_providers() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let region_id = region(1); let other = Provider::object_store_provider(region_id, "other/prefix".to_string()); @@ -3385,6 +3503,7 @@ mod tests { #[tokio::test] async fn test_store_stop_is_idempotent() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let (first, second) = tokio::join!(store.stop(), store.stop()); first.unwrap(); @@ -3425,6 +3544,7 @@ mod tests { #[tokio::test] async fn test_store_stop_with_actor_gone() { + let _serialized = SERIALIZED.lock().await; let (store, command_rx) = store_without_actor(); drop(command_rx); store.stop().await.unwrap(); @@ -3433,6 +3553,7 @@ mod tests { #[tokio::test] async fn test_store_actor_exits_when_the_store_is_dropped() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let io = Arc::downgrade(&store.io); let stopped = Arc::downgrade(&store.stopped); @@ -3449,6 +3570,7 @@ mod tests { #[tokio::test] async fn test_store_retains_first_terminal_error() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let first = set_terminal( &store.terminal_error, @@ -3486,6 +3608,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_bounds_concurrency_and_orders_by_sequence() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 16, 3).await; let (io, mut parked) = ParkedIo::over(object_store); @@ -3540,6 +3663,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_abandons_pending_fetches_on_failure() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); populate(&object_store, 16, 3).await; let (io, mut parked) = ParkedIo::over(object_store.clone()); @@ -3574,6 +3698,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_reads_only_header_and_tail_of_large_object() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_records( &object_store, @@ -3611,6 +3736,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_short_read() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1]).await; let io = ObjectStoreIo::new(object_store, PREFIX).unwrap(); @@ -3653,6 +3779,7 @@ mod tests { #[tokio::test] async fn test_store_reads_only_region_segments_in_entry_order() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let mut expected_ranges = HashMap::<_, Vec<_>>::new(); for (seq, first) in [(1, 10), (3, 30), (7, 70)] { @@ -3727,6 +3854,7 @@ mod tests { #[tokio::test] async fn test_store_obsolete_hides_entries_from_read_only() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); put_object(&object_store, 0, region(1), &[10, 11, 12]).await; put_object(&object_store, 1, region(1), &[20]).await; @@ -3774,6 +3902,7 @@ mod tests { #[tokio::test] async fn test_store_entry_and_watermarks_reject_another_region() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let p = provider(region(2)); let errors = [ @@ -3838,7 +3967,7 @@ mod tests { #[tokio::test] async fn test_store_skips_a_corrupted_segment_and_records_a_hole() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (path, footer) = write_corrupted_segment(&object_store).await; let (io, reads) = RecordingIo::over(object_store); @@ -3903,7 +4032,7 @@ mod tests { #[tokio::test] async fn test_store_corrupted_segment_fails_the_read_that_decodes_it() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (path, footer) = write_corrupted_segment(&object_store).await; let (io, reads) = RecordingIo::over(object_store); @@ -3950,7 +4079,7 @@ mod tests { #[tokio::test] async fn test_store_reads_a_segment_whose_first_download_was_damaged() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1, 2]).await; @@ -3979,7 +4108,7 @@ mod tests { #[tokio::test] async fn test_store_failed_segment_fetch_fails_the_read() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; for on_corrupted_segment in [CorruptedSegmentAction::Skip, CorruptedSegmentAction::Fail] { let object_store = memory_store(); put_object(&object_store, 0, region(1), &[1]).await; @@ -4012,6 +4141,7 @@ mod tests { #[tokio::test] async fn test_store_assigns_ids_from_the_object_sequence() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &manual()).await; let region_one = region(1); let region_two = region(2); @@ -4075,6 +4205,7 @@ mod tests { #[tokio::test] async fn test_store_seals_when_a_region_exhausts_its_positions() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &manual()).await; let region_id = region(1); let entries_of = |count: u64| { @@ -4133,6 +4264,7 @@ mod tests { #[tokio::test] async fn test_store_flushes_at_the_minimum_interval() { + let _serialized = SERIALIZED.lock().await; // An append is acknowledged once its object is durable, so a completed // append proves the tick after it sealed the batch. The appends are // spaced far apart, so every one of them lands in its own tick and @@ -4171,6 +4303,7 @@ mod tests { #[tokio::test] async fn test_store_rejects_entries_of_another_region() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &manual()).await; let region_id = region(1); @@ -4217,6 +4350,7 @@ mod tests { #[tokio::test] async fn test_store_pipelines_creates_and_acknowledges_in_sequence_order() { + let _serialized = SERIALIZED.lock().await; let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); let appends = spawn_appends(&store, region_id, MAX_IN_FLIGHT_CREATES + 2).await; @@ -4315,6 +4449,7 @@ mod tests { #[tokio::test] async fn test_store_transient_failure_fails_every_batch_that_is_not_indexed() { + let _serialized = SERIALIZED.lock().await; let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); // Every slot is taken; the last batch waits for one. @@ -4368,6 +4503,7 @@ mod tests { #[tokio::test] async fn test_store_transient_failure_before_a_created_object_keeps_it_off_the_chain() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (store, io, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); @@ -4431,6 +4567,7 @@ mod tests { #[tokio::test] async fn test_store_permanent_failure_before_a_durable_object_poisons() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); @@ -4469,6 +4606,7 @@ mod tests { #[tokio::test] async fn test_store_conflict_with_an_earlier_epoch_fails_the_batch_without_poisoning() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); @@ -4525,6 +4663,7 @@ mod tests { #[tokio::test] async fn test_store_obsolete_raises_the_sequence_floor() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &manual()).await; let region_one = region(1); let region_two = region(2); @@ -4624,6 +4763,7 @@ mod tests { #[tokio::test] async fn test_store_does_not_reuse_the_sequence_of_a_failed_create() { + let _serialized = SERIALIZED.lock().await; let (io, _) = RecordingIo::over(memory_store()); let store = open_over(io.clone(), &manual()).await; let region_one = region(1); @@ -4678,6 +4818,7 @@ mod tests { #[tokio::test] async fn test_store_floor_applies_while_a_create_is_in_flight() { + let _serialized = SERIALIZED.lock().await; let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_one = region(1); let region_two = region(2); @@ -4709,6 +4850,7 @@ mod tests { #[tokio::test] async fn test_store_obsolete_queued_behind_stop_records_the_watermark() { + let _serialized = SERIALIZED.lock().await; let (store, mut command_rx) = store_without_actor(); let region_id = region(1); let provider_one = provider(region_id); @@ -4742,6 +4884,7 @@ mod tests { #[tokio::test] async fn test_store_hooks_hold_and_fail_creates() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let region_id = region(1); store.hold_creates(); @@ -4769,6 +4912,7 @@ mod tests { #[tokio::test] async fn test_store_hook_fails_the_next_create_after_it_writes() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let region_id = region(1); @@ -4795,6 +4939,7 @@ mod tests { #[tokio::test] async fn test_store_hooks_observe_parked_creates_waits_and_crash() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &enqueued(eager())).await; let region_id = region(1); @@ -4826,6 +4971,7 @@ mod tests { #[tokio::test] async fn test_store_create_held_when_the_store_is_dropped_never_runs() { + let _serialized = SERIALIZED.lock().await; let (io, _) = RecordingIo::over(memory_store()); let store = open_over(io.clone(), &eager()).await; store.hold_creates(); @@ -4857,6 +5003,7 @@ mod tests { #[tokio::test] async fn test_store_rollback_and_poison_drop_the_open_batch() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (store, _, mut parked) = open_parking_creates(object_store.clone(), &manual()).await; let region_id = region(1); @@ -4927,6 +5074,7 @@ mod tests { #[tokio::test] async fn test_store_poisons_once_the_last_sequence_is_taken() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let last = OBJECT_SEQ_LIMIT - 1; // The start object of the store takes the sequence before the last. @@ -4958,6 +5106,7 @@ mod tests { #[tokio::test] async fn test_store_holds_appends_back_while_the_sealed_batches_are_full() { + let _serialized = SERIALIZED.lock().await; let (store, _, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); let admitted = || *store.admitted_appends.borrow(); @@ -5020,6 +5169,7 @@ mod tests { #[tokio::test] async fn test_store_stop_during_failed_flush_reports_stopped() { + let _serialized = SERIALIZED.lock().await; let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; let pending = spawn_append_batch(&store, vec![entry(&store, region(1), "a1")]); let (_, release) = next_parked_operation(&mut parked).await; @@ -5036,6 +5186,7 @@ mod tests { #[tokio::test] async fn test_store_stop_with_several_creates_in_flight() { + let _serialized = SERIALIZED.lock().await; let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); // Four creates are in flight, the fifth batch waits for a slot. @@ -5089,6 +5240,7 @@ mod tests { #[tokio::test] async fn test_store_stop_during_conflicting_flush_reports_stopped_and_poisons() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (store, _, mut parked) = open_parking_creates(object_store.clone(), &eager()).await; let region_id = region(1); @@ -5241,6 +5393,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_indexes_only_the_chain() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let record = |entry_id: EntryId, data: &'static str| Record { @@ -5279,6 +5432,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_never_reuses_the_sequence_of_an_orphan() { + let _serialized = SERIALIZED.lock().await; let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); put_fixture(&io, 0, &[]).await; // Object 2 extends object 1, which never landed. @@ -5297,6 +5451,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_rejects_objects_without_a_complete_chain() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); // Object 1 extends object 0, which is gone. @@ -5314,6 +5469,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_keeps_the_later_epoch_over_a_late_object() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let record = |entry_id: EntryId, data: &'static str| Record { @@ -5356,6 +5512,7 @@ mod tests { #[tokio::test] async fn test_store_open_starts_the_first_epoch_on_an_empty_prefix() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let store = open(object_store.clone(), &eager()).await; store.stop().await.unwrap(); @@ -5368,6 +5525,7 @@ mod tests { #[tokio::test] async fn test_store_start_object_moves_past_an_earlier_epoch_only() { + let _serialized = SERIALIZED.lock().await; for (epoch, moves) in [(1, true), (2, false)] { let object_store = memory_store(); let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); @@ -5405,6 +5563,7 @@ mod tests { #[tokio::test] async fn test_store_open_fails_when_its_start_object_is_already_present() { + let _serialized = SERIALIZED.lock().await; // On an empty prefix and after object 0 of epoch 1, another open that // recovered the same objects wrote the very start object this open // writes. @@ -5450,6 +5609,7 @@ mod tests { #[tokio::test] async fn test_store_opens_racing_across_a_late_object_take_distinct_epochs() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); put_fixture(&io, 0, &[]).await; @@ -5484,6 +5644,7 @@ mod tests { #[tokio::test] async fn test_store_open_rejects_an_epoch_above_the_next_sequence() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); // No instance that starts at or below sequence 0 writes epoch 5. @@ -5501,6 +5662,7 @@ mod tests { #[tokio::test] async fn test_store_open_fails_when_its_start_object_lands_with_an_unknown_outcome() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); put_fixture(&fixtures, 0, &[]).await; @@ -5537,6 +5699,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_append_returns_before_the_object_exists() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &enqueued(manual())).await; let region_id = region(1); @@ -5573,6 +5736,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_durable_id_advances_after_create_and_indexing() { + let _serialized = SERIALIZED.lock().await; let (store, _, mut parked) = open_parking_creates(memory_store(), &enqueued(eager())).await; let region_one = region(1); let region_two = region(2); @@ -5619,6 +5783,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_wait_ignores_later_pending_entries_of_the_region() { + let _serialized = SERIALIZED.lock().await; let (store, _io, mut parked) = open_parking_creates(memory_store(), &enqueued(eager())).await; let region_a = region(1); @@ -5671,6 +5836,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_obsolete_never_passes_the_durable_id() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &enqueued(manual())).await; let region_id = region(1); append(&store, region_id, "a1").await.unwrap(); @@ -5741,6 +5907,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_backlog_bytes_stall_admission_until_an_upload_completes() { + let _serialized = SERIALIZED.lock().await; let config = ObjectStoreWalConfig { max_unpersisted_bytes: ReadableSize(1), ..enqueued(manual()) @@ -5789,6 +5956,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_backlog_age_stalls_admission_until_an_upload_completes() { + let _serialized = SERIALIZED.lock().await; let config = ObjectStoreWalConfig { max_unpersisted_age: Duration::from_millis(50), ..enqueued(manual()) @@ -5811,6 +5979,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_stop_fails_a_stalled_append_and_uploads_the_backlog() { + let _serialized = SERIALIZED.lock().await; let config = ObjectStoreWalConfig { max_unpersisted_bytes: ReadableSize(1), ..enqueued(manual()) @@ -5830,6 +5999,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_repeats_a_create_that_failed_transiently() { + let _serialized = SERIALIZED.lock().await; let (io, _) = RecordingIo::over(memory_store()); let store = open_over(io.clone(), &enqueued(eager())).await; let region_id = region(1); @@ -5865,6 +6035,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_permanent_failure_poisons() { + let _serialized = SERIALIZED.lock().await; // A conflict of any epoch poisons the store without reading the // existing object, since the acknowledged entries cannot move to // another sequence. So does a permanent storage error. @@ -5903,6 +6074,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_stop_reports_a_backlog_lost_before_the_stop_command() { + let _serialized = SERIALIZED.lock().await; let (store, io, mut parked) = open_parking_creates(memory_store(), &enqueued(manual())).await; let region_id = region(1); @@ -5945,6 +6117,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_stop_uploads_the_backlog() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &enqueued(manual())).await; let region_id = region(1); append(&store, region_id, "a1").await.unwrap(); @@ -5987,6 +6160,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_issued_id_needs_no_floor_before_it_is_durable() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &enqueued(manual())).await; let region_id = region(1); append(&store, region_id, "a1").await.unwrap(); @@ -6034,6 +6208,7 @@ mod tests { #[tokio::test] async fn test_store_enqueued_inherited_watermark_raises_the_floor() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let store = open(object_store.clone(), &enqueued(manual())).await; let region_id = region(1); @@ -6071,6 +6246,7 @@ mod tests { #[tokio::test] async fn test_store_durability_wait_pending_at_stop_fails_with_stopped() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &manual()).await; let region_id = region(1); let append = spawn_append_batch(&store, vec![entry(&store, region_id, "a1")]); @@ -6098,6 +6274,7 @@ mod tests { #[tokio::test] async fn test_store_stamps_the_reclaim_boundary_and_collects_below_it() { + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let io = store.io.clone(); let (one, two, three) = (region(1), region(2), region(3)); @@ -6142,6 +6319,7 @@ mod tests { #[tokio::test] async fn test_store_collects_failed_batches_below_a_durable_boundary() { + let _serialized = SERIALIZED.lock().await; let (io, _) = RecordingIo::over(memory_store()); let store = open_over(io.clone(), &eager()).await; let region_id = region(1); @@ -6181,7 +6359,7 @@ mod tests { #[tokio::test] async fn test_store_attempts_a_failed_delete_again_at_the_next_collection() { - let _serialized = FAILED_DELETES.lock().await; + let _serialized = SERIALIZED.lock().await; let (io, mut parked) = ParkedIo::parking_creates(memory_store()); let store = open_over(io.clone(), &eager()).await; io.deletes_parked.store(true, Ordering::SeqCst); @@ -6222,6 +6400,7 @@ mod tests { #[tokio::test] async fn test_store_collection_ends_with_stop_and_poison() { + let _serialized = SERIALIZED.lock().await; for poison in [false, true] { let object_store = memory_store(); let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); @@ -6286,7 +6465,7 @@ mod tests { #[tokio::test] async fn test_store_drops_the_holes_of_a_collected_object() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); write_corrupted_segment(&object_store).await; let store = open(object_store, &eager()).await; @@ -6309,7 +6488,7 @@ mod tests { #[tokio::test] async fn test_store_records_no_hole_for_a_collected_object() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let (_, footer) = write_corrupted_segment(&object_store).await; let store = open(object_store, &eager()).await; @@ -6337,6 +6516,7 @@ mod tests { #[tokio::test] async fn test_store_collects_a_failed_batch_whose_create_completes_late() { + let _serialized = SERIALIZED.lock().await; let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; let region_id = region(1); // Object 1 fails while the create of object 2 is in flight. @@ -6375,6 +6555,7 @@ mod tests { #[tokio::test] async fn test_store_recovery_tolerates_objects_collected_after_the_listing() { + let _serialized = SERIALIZED.lock().await; let io = ObjectStoreIo::new(memory_store(), PREFIX).unwrap(); let bounded = |object_seq, prev, reclaim_boundary| Header { reclaim_boundary, @@ -6404,7 +6585,7 @@ mod tests { #[tokio::test] async fn test_store_read_skips_an_object_collected_after_it_was_listed() { - let _serialized = SKIPPED_SEGMENTS.lock().await; + let _serialized = SERIALIZED.lock().await; let store = open(memory_store(), &eager()).await; let region_id = region(1); let skipped = METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(); @@ -6438,8 +6619,7 @@ mod tests { #[tokio::test] async fn test_store_read_waits_for_an_undecided_delete() { - let _serialized = SKIPPED_SEGMENTS.lock().await; - let _failing = FAILED_DELETES.lock().await; + let _serialized = SERIALIZED.lock().await; for deleted in [true, false] { let object_store = memory_store(); let (io, mut parked) = ParkedIo::parking_creates(object_store.clone()); @@ -6508,6 +6688,7 @@ mod tests { #[tokio::test] async fn test_store_recovers_from_the_reclaim_boundary_of_the_tip() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let region_id = region(1); @@ -6549,6 +6730,7 @@ mod tests { #[tokio::test] async fn test_store_collects_a_late_object_of_an_earlier_epoch_below_the_boundary() { + let _serialized = SERIALIZED.lock().await; let object_store = memory_store(); let io = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); let region_id = region(1); @@ -6801,13 +6983,15 @@ mod tests { /// Object access that records every range read as (sequence, offset, /// length) and, on request, reports the next conditional create as failed /// after it wrote the object, fails it with a given error before it - /// writes, or damages the bytes the next range read returns. + /// writes, stores the same content before it writes, or damages the bytes + /// the next range read returns. struct RecordingIo { inner: ObjectStoreIo, reads: RangeReads, fail_after_next_put: AtomicBool, fail_next_put: Mutex>, damage_next_range_read: AtomicBool, + duplicate_next_put: AtomicBool, } impl RecordingIo { @@ -6819,6 +7003,7 @@ mod tests { fail_after_next_put: AtomicBool::new(false), fail_next_put: Mutex::new(None), damage_next_range_read: AtomicBool::new(false), + duplicate_next_put: AtomicBool::new(false), }; (Arc::new(io), reads) } @@ -6850,6 +7035,11 @@ mod tests { path: self.object_path(object_seq), }); } + if self.duplicate_next_put.swap(false, Ordering::Relaxed) { + self.inner + .put_if_absent(object_seq, content.clone()) + .await?; + } let result = self.inner.put_if_absent(object_seq, content).await?; if self.fail_after_next_put.swap(false, Ordering::Relaxed) { return injected_failure("write", self.object_path(object_seq)); @@ -6886,4 +7076,1155 @@ mod tests { self.inner.object_path(object_seq) } } + + /// Drives every outcome of a conditional create through the store, one + /// case per fault, and checks that every row of the *Failure matrix* has one. + mod fault_matrix { + use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use prometheus::Histogram; + + use super::*; + + /// A row of the *Failure matrix*, named by its *Situation* cell. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum Row { + /// create succeeds, or identical retry + Created, + /// transient I/O error, `durable` mode + TransientDurable, + /// an object of an earlier epoch holds the sequence, `durable` mode + EarlierEpochDurable, + /// transient I/O error, `enqueued` mode + TransientEnqueued, + /// transient I/O error, `enqueued` mode after `stop` began + TransientEnqueuedAfterStop, + /// an object of the same or a later epoch holds the sequence, or + /// any conflicting object in the `enqueued` mode + Conflict, + /// encoding or catalog error + CatalogError, + /// create succeeds at the last representable sequence + LastSequence, + /// the start object meets an object of an earlier epoch + StartMovesPast, + /// the start object meets an object of the same or a later epoch, + /// or its create has an unknown outcome + StartFails, + } + + /// Every row of the matrix, in the order the RFC lists them. + const ROWS: [Row; 10] = [ + Row::Created, + Row::TransientDurable, + Row::EarlierEpochDurable, + Row::TransientEnqueued, + Row::TransientEnqueuedAfterStop, + Row::Conflict, + Row::CatalogError, + Row::LastSequence, + Row::StartMovesPast, + Row::StartFails, + ]; + + /// One failure of the object store access and the row it exercises. + struct Case { + fault: &'static str, + row: Row, + run: fn() -> BoxFuture<'static, ()>, + } + + /// Declares [`Counters`] from one sampling expression per field. + macro_rules! counters { + ($($(#[$doc:meta])* $field:ident: $sample:expr,)*) => { + /// The metrics the cases assert on, sampled around a case. + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] + struct Counters { + $($(#[$doc])* $field: u64,)* + } + + impl Counters { + fn sample() -> Self { + Self { $($field: $sample,)* } + } + + /// The metrics recorded since `before`. + fn since(before: Self) -> Self { + let now = Self::sample(); + Self { $($field: now.$field - before.$field,)* } + } + } + }; + } + + counters! { + created_objects: METRIC_OBJECT_STORE_WAL_CREATED_OBJECTS_TOTAL.get(), + /// Samples of the seal-to-durable histogram. + seal_to_durable: METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS.get_sample_count(), + create_failures: METRIC_OBJECT_STORE_WAL_CREATE_FAILURES_TOTAL.get(), + create_conflicts: METRIC_OBJECT_STORE_WAL_CREATE_CONFLICTS_TOTAL.get(), + poisoned: METRIC_OBJECT_STORE_WAL_POISONED_TOTAL.get(), + /// Samples of the object size and entry histograms, one per seal. + sealed_objects: sealed_objects(), + /// The sums of those two histograms. + object_bytes: sum_of(&METRIC_OBJECT_STORE_WAL_OBJECT_BYTES), + object_entries: sum_of(&METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES), + /// Samples of the `durable` acknowledgement histogram. + acknowledged_appends: METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS.get_sample_count(), + /// Objects recovery read footers from, and recoveries timed. + recovered_objects: METRIC_OBJECT_STORE_WAL_RECOVERED_OBJECTS_TOTAL.get(), + recoveries: METRIC_OBJECT_STORE_WAL_RECOVERY_SECONDS.get_sample_count(), + /// Read streams driven to their end. + reads: METRIC_OBJECT_STORE_WAL_READ_SECONDS.get_sample_count(), + deleted_objects: METRIC_OBJECT_STORE_WAL_DELETED_OBJECTS_TOTAL.get(), + failed_deletes: METRIC_OBJECT_STORE_WAL_FAILED_DELETES_TOTAL.get(), + /// No read of the matrix meets a corrupted segment. + skipped_segments: METRIC_OBJECT_STORE_WAL_SKIPPED_SEGMENTS_TOTAL.get(), + stalled_appends: METRIC_OBJECT_STORE_WAL_STALLED_APPENDS_TOTAL.get(), + stalled_waits: METRIC_OBJECT_STORE_WAL_STALLED_APPEND_SECONDS.get_sample_count(), + } + + /// Samples of the two object shape histograms, which a sealed batch + /// feeds together; a batch that fed only one of them fails the case. + fn sealed_objects() -> u64 { + let bytes = METRIC_OBJECT_STORE_WAL_OBJECT_BYTES.get_sample_count(); + let entries = METRIC_OBJECT_STORE_WAL_OBJECT_ENTRIES.get_sample_count(); + assert_eq!( + bytes, entries, + "an object was sized but its entries were not counted" + ); + bytes + } + + /// What a histogram of whole numbers has recorded so far. The sums + /// are exact at these magnitudes, so a case states them as counts. + fn sum_of(histogram: &Histogram) -> u64 { + let sum = histogram.get_sample_sum(); + assert_eq!( + sum, + sum.round(), + "a histogram of whole numbers recorded {sum}" + ); + sum as u64 + } + + /// The length of the object a batch of `entries` encodes into; the + /// header and the entry ids have a fixed width, so any will do. + fn object_bytes_of(entries: &[(RegionId, &str)]) -> u64 { + let mut positions = HashMap::new(); + let records = entries + .iter() + .map(|(region_id, data)| { + let position = positions.entry(*region_id).or_insert(0); + *position += 1; + Record { + region_id: *region_id, + entry_id: id(1, *position), + payload: Bytes::from(data.as_bytes().to_vec()), + } + }) + .collect::>(); + encode_object(chain_header(1, 1, None), &records) + .unwrap() + .bytes + .len() as u64 + } + + /// The length of the object of one two-byte entry, which most cases seal. + fn one_entry_object() -> u64 { + object_bytes_of(&[(region(1), "a1")]) + } + + /// The sum and the count of a timing histogram before a forced wait. + fn timer_before(histogram: &Histogram) -> (f64, u64) { + (histogram.get_sample_sum(), histogram.get_sample_count()) + } + + /// One sample of a timing histogram, the wait the case forced on it, + /// and what the case itself took around that wait. + #[derive(Debug, Clone, Copy)] + struct Timed { + recorded: f64, + /// The forced wait as measured, which may be far longer than asked. + held: Duration, + elapsed: Duration, + } + + impl Timed { + fn since( + histogram: &Histogram, + before: (f64, u64), + held: Duration, + elapsed: Duration, + ) -> Self { + assert_eq!( + before.1 + 1, + histogram.get_sample_count(), + "expected the timer to take exactly one sample" + ); + Self { + recorded: histogram.get_sample_sum() - before.0, + held, + elapsed, + } + } + + /// Asserts two observations of one timer: each lies between its own + /// wait and its case's elapsed time, and the long wait outlasted the + /// whole short observation, so a timer that records a constant + /// fails one of them however slow the machine is. + fn assert_pair(name: &str, short: Self, long: Self) { + assert!( + long.held > short.elapsed, + "the {name} timer was observed around a long wait of {:?} that did not outlast the short observation's {:?}", + long.held, + short.elapsed + ); + for observed in [short, long] { + assert!( + observed.recorded >= observed.held.as_secs_f64() + && observed.recorded <= observed.elapsed.as_secs_f64(), + "the {name} timer recorded {}s around a wait of {:?}, expected at most {:?}", + observed.recorded, + observed.held, + observed.elapsed + ); + } + } + } + + fn indexed_gauges() -> (i64, i64) { + ( + METRIC_OBJECT_STORE_WAL_INDEXED_OBJECTS.get(), + METRIC_OBJECT_STORE_WAL_INDEXED_BYTES.get(), + ) + } + + /// Asserts that both indexed gauges hold exactly the objects `indexed`, + /// with their sizes as listed. Objects a failed create left behind are + /// listed but not counted. + async fn assert_indexed(io: &dyn WalObjectIo, indexed: &[u64]) { + let bytes = io + .list() + .await + .unwrap() + .into_iter() + .filter(|object| indexed.contains(&object.object_seq)) + .map(|object| object.size) + .sum::(); + assert_eq!( + (indexed.len() as i64, bytes as i64), + indexed_gauges(), + "expected the gauges to hold the objects {indexed:?}" + ); + } + + /// Asserts every append failed with an error `recognise` accepts. + async fn assert_appends_failed( + appends: Vec>>, + recognise: impl Fn(&Error) -> bool, + ) { + for append in appends { + let error = timeout(WAIT, append).await.unwrap().unwrap().unwrap_err(); + assert!( + recognise(unwrap_shared(&error)), + "unexpected error: {error:?}" + ); + } + } + + /// Asserts that a read of the poisoned store fails with the terminal error. + async fn assert_read_is_terminal(store: &ObjectStoreLogStore, region_id: RegionId) { + let error = store + .read(&provider(region_id), 1, None) + .await + .err() + .expect("a read of a poisoned store must fail"); + let terminal = terminal(&store.terminal_error).expect("the store is not poisoned"); + assert!( + matches!(&error, Error::ObjectStoreWal { source, .. } if Arc::ptr_eq(source, &terminal)), + "unexpected error: {error:?}" + ); + } + + /// Polls until `settled` holds. + async fn wait_until(settled: impl Fn() -> bool) { + let deadline = Instant::now() + WAIT; + while !settled() { + assert!(Instant::now() < deadline, "the actor never settled"); + tokio::time::sleep(Duration::from_millis(1)).await; + } + } + + #[tokio::test] + async fn test_store_fault_matrix() { + let _serialized = SERIALIZED.lock().await; + wait_until(|| LIVE_ACTORS.load(Ordering::SeqCst) == 0).await; + let cases = [ + Case { + fault: "a create that succeeds", + row: Row::Created, + run: || Box::pin(create_succeeds()), + }, + Case { + fault: "a create that finds an object with identical content", + row: Row::Created, + run: || Box::pin(create_finds_an_identical_object()), + }, + Case { + fault: "a create whose response is lost", + row: Row::Created, + run: || Box::pin(create_response_is_lost()), + }, + Case { + fault: "a create that fails transiently with later objects created and in flight, `durable` mode", + row: Row::TransientDurable, + run: || Box::pin(create_fails_transiently_durable()), + }, + Case { + fault: "a create that fails transiently, `enqueued` mode", + row: Row::TransientEnqueued, + run: || Box::pin(create_fails_transiently_enqueued()), + }, + Case { + fault: "a create that fails transiently after `stop` began, `enqueued` mode", + row: Row::TransientEnqueuedAfterStop, + run: || Box::pin(create_fails_transiently_after_stop_began()), + }, + Case { + fault: "a create that finds an object of an earlier epoch, `durable` mode", + row: Row::EarlierEpochDurable, + run: || Box::pin(create_finds_an_earlier_epoch()), + }, + Case { + fault: "a create that finds an object of an earlier epoch, `enqueued` mode", + row: Row::Conflict, + run: || Box::pin(create_finds_an_earlier_epoch_enqueued()), + }, + Case { + fault: "a create that finds an object of the same or a later epoch", + row: Row::Conflict, + run: || Box::pin(create_finds_another_writer()), + }, + Case { + fault: "a create that finds an object whose header cannot be read", + row: Row::Conflict, + run: || Box::pin(create_finds_an_unreadable_object()), + }, + Case { + fault: "a create that completes after a conflict poisoned the store", + row: Row::Conflict, + run: || Box::pin(create_completes_after_the_store_was_poisoned()), + }, + Case { + fault: "a created object the catalog rejects", + row: Row::CatalogError, + run: || Box::pin(catalog_rejects_the_created_object()), + }, + Case { + fault: "a create at the last representable object sequence", + row: Row::LastSequence, + run: || Box::pin(create_takes_the_last_object_sequence()), + }, + Case { + fault: "a start object that meets an object of an earlier epoch", + row: Row::StartMovesPast, + run: || Box::pin(start_object_meets_an_earlier_epoch()), + }, + Case { + fault: "a start object that meets an object of the same or a later epoch", + row: Row::StartFails, + run: || Box::pin(start_object_meets_another_writer()), + }, + Case { + fault: "a start object that meets an object with identical content", + row: Row::StartFails, + run: || Box::pin(start_object_meets_an_identical_object()), + }, + Case { + fault: "a start object whose create has an unknown outcome", + row: Row::StartFails, + run: || Box::pin(start_object_create_has_an_unknown_outcome()), + }, + ]; + + common_telemetry::init_default_ut_logging(); + for case in &cases { + common_telemetry::info!("Fault matrix case: {}", case.fault); + (case.run)().await; + } + for row in ROWS { + assert!( + cases.iter().any(|case| case.row == row), + "no case covers the failure matrix row {row:?}" + ); + } + } + + /// The batches are acknowledged in order and the sequence advances. + /// Each create is held, the second past the whole first observation, + /// so both write path timers are observed around disjoint waits. + async fn create_succeeds() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let (region_id, other) = (region(1), region(2)); + let mut seals = Vec::new(); + let mut acknowledgements = Vec::new(); + let mut hold = Duration::from_millis(20); + // The second batch holds entries of two regions, so entries and + // objects differ. + for batch in [ + vec![(region_id, "a1")], + vec![(region_id, "a2"), (region_id, "a3"), (other, "b1")], + ] { + let seal = timer_before(&METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS); + let acknowledgement = timer_before(&METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS); + let started_at = Instant::now(); + let entries = batch + .iter() + .map(|(region_id, data)| entry(&store, *region_id, data)) + .collect(); + let appended = spawn_append_batch(&store, entries); + let (_, release) = next_parked_operation(&mut parked).await; + let parked_at = Instant::now(); + tokio::time::sleep(hold).await; + let held = parked_at.elapsed(); + release.send(true).unwrap(); + timeout(WAIT, appended).await.unwrap().unwrap().unwrap(); + let elapsed = started_at.elapsed(); + let timed = |timer, before| Timed::since(timer, before, held, elapsed); + seals.push(timed( + &METRIC_OBJECT_STORE_WAL_SEAL_TO_DURABLE_SECONDS, + seal, + )); + acknowledgements.push(timed( + &METRIC_OBJECT_STORE_WAL_APPEND_ACK_SECONDS, + acknowledgement, + )); + hold = elapsed + Duration::from_millis(20); + } + Timed::assert_pair("seal to durable", seals[0], seals[1]); + Timed::assert_pair( + "append acknowledgement", + acknowledgements[0], + acknowledgements[1], + ); + + // Object 1 carries the boundary past the start object. + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!( + expected_entries( + region_id, + &[(id(1, 1), "a1"), (id(2, 1), "a2"), (id(2, 2), "a3")] + ), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1, 2]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: one_entry_object() + + object_bytes_of(&[(region_id, "a2"), (region_id, "a3"), (other, "b1")]), + object_entries: 4, + acknowledged_appends: 2, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Appends over `io` after `arm` made the create meet the very object + /// it writes, and asserts it was accepted as an identical retry. + async fn assert_identical_retry(io: Arc, arm: impl FnOnce()) { + let store = open_over(io.clone(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + arm(); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1"), (id(2, 1), "a2")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1, 2]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + acknowledged_appends: 2, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_finds_an_identical_object() { + let (io, _) = RecordingIo::over(memory_store()); + let armed = io.clone(); + assert_identical_retry(io.clone(), move || { + armed.duplicate_next_put.store(true, Ordering::Relaxed) + }) + .await; + assert!(!io.duplicate_next_put.load(Ordering::Relaxed)); + } + + /// The object store writes the object and loses the response, and the + /// read-back that reconciles the create finds the same bytes. + async fn create_response_is_lost() { + let lose_next = Arc::new(AtomicBool::new(false)); + let losing = lose_next.clone(); + let object_store = memory_store().layer( + MockLayerBuilder::default() + .writer_factory(Arc::new(move |_, _, inner| { + if losing.swap(false, Ordering::SeqCst) { + Box::new(LosingWriter(inner)) + } else { + inner + } + })) + .build() + .unwrap(), + ); + let io = Arc::new(ObjectStoreIo::new(object_store, PREFIX).unwrap()); + let armed = lose_next.clone(); + assert_identical_retry(io, move || armed.store(true, Ordering::SeqCst)).await; + assert!(!lose_next.load(Ordering::SeqCst)); + } + + /// A writer that writes the object and then reports a failure, as a + /// create whose response is lost. + struct LosingWriter(oio::Writer); + + impl oio::Write for LosingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.0.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + self.0.close().await?; + Err(mock::Error::new(mock::ErrorKind::Unexpected, "lost response").set_temporary()) + } + + async fn abort(&mut self) -> mock::Result<()> { + self.0.abort().await + } + } + + /// Object 2 is created and object 3 in flight when object 1 fails. + async fn create_fails_transiently_durable() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 3).await; + let mut releases = parked_operations(&mut parked, 3).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(false).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectStore { .. }) + }) + .await; + assert_indexed(io.as_ref(), &[0]).await; + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 3, + object_bytes: 3 * one_entry_object(), + object_entries: 3, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + releases.remove(&3).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 2).await; + let appended = spawn_append_batch(&store, vec![entry(&store, region_id, "a4")]); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(4, object_seq); + release.send(true).unwrap(); + timeout(WAIT, appended).await.unwrap().unwrap().unwrap(); + assert_eq!(vec![2, 3, 4], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(4, 1), "a4")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[4]).await; + assert_eq!( + Counters { + created_objects: 3, + seal_to_durable: 3, + create_failures: 1, + sealed_objects: 4, + object_bytes: 4 * one_entry_object(), + object_entries: 4, + acknowledged_appends: 1, + reads: 2, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_fails_transiently_enqueued() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(eager())).await; + let before = Counters::sample(); + let region_id = region(1); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let (object_seq, release) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + release.send(false).unwrap(); + + // While the repeat is parked nothing is written or indexed. + let (object_seq, repeat) = next_parked_operation(&mut parked).await; + assert_eq!(1, object_seq); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert_eq!(0, store.durable_entry_id(&provider(region_id)).unwrap()); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + repeat.send(true).unwrap(); + timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap(); + assert_eq!(vec![1], collected_object_seqs(&store).await); + assert_eq!( + expected_entries(region_id, &[(id(1, 1), "a1")]), + read_entries(&store, region_id, 1).await + ); + assert_indexed(io.as_ref(), &[1]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 2, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_fails_transiently_after_stop_began() { + let (store, io, mut parked) = + open_parking_creates(memory_store(), &enqueued(manual())).await; + let before = Counters::sample(); + let region_id = region(1); + append(&store, region_id, "a1").await.unwrap(); + let seal = { + let store = store.clone(); + tokio::spawn(async move { store.seal_open_batch().await }) + }; + let (_, release) = next_parked_operation(&mut parked).await; + store.begin_stop(); + release.send(false).unwrap(); + timeout(WAIT, seal).await.unwrap().unwrap().unwrap_err(); + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectStore { .. }), + "unexpected error: {error:?}" + ); + assert!(terminal(&store.terminal_error).is_none()); + assert_eq!(vec![0], object_seqs(io.as_ref()).await); + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_failures: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + async fn create_finds_an_earlier_epoch() { + let object_store = memory_store(); + let store = open(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, 0).await; + let before = Counters::sample(); + let region_id = region(1); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::StaleWalObject { .. }), + "unexpected error: {error:?}" + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + assert_indexed(store.io.as_ref(), &[0]).await; + assert!(read_entries(&store, region_id, 1).await.is_empty()); + assert_eq!( + Counters { + create_conflicts: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + + let response = append(&store, region_id, "a2").await.unwrap(); + assert_eq!(Some(&id(2, 1)), response.last_entry_ids.get(®ion_id)); + assert_eq!(vec![1, 2], collected_object_seqs(&store).await); + assert_indexed(store.io.as_ref(), &[2]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + acknowledged_appends: 1, + reads: 1, + deleted_objects: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_finds_an_earlier_epoch_enqueued() { + let object_store = memory_store(); + let store = open(object_store.clone(), &enqueued(eager())).await; + put_foreign(&object_store, 1, 0).await; + let before = Counters::sample(); + let region_id = region(1); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!(Some(&id(1, 1)), response.last_entry_ids.get(®ion_id)); + let error = timeout(WAIT, store.wait_durable(&provider(region_id), id(1, 1))) + .await + .unwrap() + .unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(store.io.as_ref(), &[0]).await; + let error = store.stop().await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectConflict { .. }), + "unexpected error: {error:?}" + ); + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + ..Counters::default() + }, + Counters::since(before) + ); + } + + async fn create_finds_another_writer() { + for epoch in [1, 2] { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, epoch).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + releases.remove(&1).unwrap().send(true).unwrap(); + let path = io.object_path(1); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectConflict { path: actual, .. } if *actual == path) + }) + .await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + } + + /// The epoch of the object at the sequence cannot be read: the conflict + /// is counted once and the store poisons itself. + async fn create_finds_an_unreadable_object() { + let object_store = memory_store(); + let (io, _) = RecordingIo::over(object_store.clone()); + let store = open_over(io.clone(), &eager()).await; + put_foreign(&object_store, 1, 1).await; + let before = Counters::sample(); + let region_id = region(1); + io.damage_next_range_read.store(true, Ordering::Relaxed); + let error = append(&store, region_id, "a1").await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::InvalidWalObject { .. }), + "unexpected error: {error:?}" + ); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// The create of object 2 completes after the conflict at object 1 + /// poisoned the store; it is counted and timed all the same. + async fn create_completes_after_the_store_was_poisoned() { + let object_store = memory_store(); + let (store, io, mut parked) = + open_parking_creates(object_store.clone(), &eager()).await; + put_foreign(&object_store, 1, 1).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&1).unwrap().send(true).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::WalObjectConflict { .. }) + }) + .await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert_read_is_terminal(&store, region_id).await; + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + create_conflicts: 1, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn catalog_rejects_the_created_object() { + let (store, io, mut parked) = open_parking_creates(memory_store(), &eager()).await; + let before = Counters::sample(); + let region_id = region(1); + let appends = spawn_appends(&store, region_id, 2).await; + let mut releases = parked_operations(&mut parked, 2).await; + releases.remove(&2).unwrap().send(true).unwrap(); + wait_until(|| Counters::since(before).created_objects == 1).await; + let occupying = FooterEntry { + region_id: region(9), + min_entry_id: id(1, 1), + max_entry_id: id(1, 1), + entry_count: 1, + segment_offset: HEADER_LEN as u64, + segment_len: 1, + segment_crc32: 0, + }; + store + .catalog + .write() + .unwrap() + .insert_object(1, vec![occupying]) + .unwrap(); + releases.remove(&1).unwrap().send(true).unwrap(); + assert_appends_failed(appends, |error| { + matches!(error, Error::CorruptedWalObject { .. }) + }) + .await; + assert_eq!(vec![0, 1, 2], object_seqs(io.as_ref()).await); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + // This store indexed neither object, whatever the case put into + // its catalog. + assert_indexed(io.as_ref(), &[0]).await; + assert_eq!( + Counters { + created_objects: 2, + seal_to_durable: 2, + poisoned: 1, + sealed_objects: 2, + object_bytes: 2 * one_entry_object(), + object_entries: 2, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + async fn create_takes_the_last_object_sequence() { + let object_store = memory_store(); + let last_object_seq = OBJECT_SEQ_LIMIT - 1; + put_records(&object_store, last_object_seq - 2, &[]).await; + let store = open(object_store, &eager()).await; + let region_id = region(1); + // A poisoned store hands out no entry, so the second is built now. + let second = entry(&store, region_id, "a2"); + let before = Counters::sample(); + let response = append(&store, region_id, "a1").await.unwrap(); + assert_eq!( + Some(&id(last_object_seq, 1)), + response.last_entry_ids.get(®ion_id) + ); + let error = store.append_batch(vec![second]).await.unwrap_err(); + assert!( + matches!(unwrap_shared(&error), Error::WalObjectSequenceExhausted { last_object_seq: actual, .. } + if *actual == last_object_seq), + "unexpected error: {error:?}" + ); + assert!(store.latest_entry_id(&provider(region_id)).is_err()); + assert_read_is_terminal(&store, region_id).await; + assert_indexed( + store.io.as_ref(), + &[last_object_seq - 2, last_object_seq - 1, last_object_seq], + ) + .await; + assert_eq!( + Counters { + created_objects: 1, + seal_to_durable: 1, + poisoned: 1, + sealed_objects: 1, + object_bytes: one_entry_object(), + object_entries: 1, + acknowledged_appends: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Writes object 0 of epoch 1 with one entry of region 1. + async fn put_chain_start(object_store: &ObjectStore) -> ObjectStoreIo { + let fixtures = ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(); + let record = Record { + region_id: region(1), + entry_id: id(0, 1), + payload: Bytes::from_static(b"a0"), + }; + put_fixture(&fixtures, 0, &[record]).await; + fixtures + } + + /// Opens a store whose start object meets `late` at sequence 1. + async fn open_meeting( + object_store: &ObjectStore, + late: Bytes, + ) -> Result> { + let io = Arc::new(RacingIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + late: Mutex::new(Some((1, late))), + }); + ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()).await + } + + async fn start_object_meets_an_earlier_epoch() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let late = encode_object(chain_header(1, 1, None), &[]).unwrap().bytes; + let before = Counters::sample(); + let store = open_meeting(&object_store, late).await.unwrap(); + assert_eq!(vec![0, 1, 2], object_seqs(&fixtures).await); + assert_eq!( + 3, + decode_header(&fixtures.get(2).await.unwrap()) + .unwrap() + .epoch + ); + assert_indexed(&fixtures, &[0, 2]).await; + assert_eq!( + expected_entries(region(1), &[(id(0, 1), "a0")]), + read_entries(&store, region(1), 0).await + ); + assert_eq!( + Counters { + created_objects: 1, + create_conflicts: 1, + recovered_objects: 1, + recoveries: 1, + reads: 1, + ..Counters::default() + }, + Counters::since(before) + ); + store.stop().await.unwrap(); + } + + /// Asserts a construction failure at the start object: no other + /// sequence was tried, and only the recovery and `create` are counted. + async fn assert_start_fails( + fixtures: &ObjectStoreIo, + before: Counters, + create: Counters, + error: Error, + recognise: impl Fn(&Error) -> bool, + retry_hint: RetryHint, + ) { + assert!(recognise(&error), "unexpected error: {error:?}"); + assert_eq!(retry_hint, error.retry_hint()); + assert_eq!(vec![0, 1], object_seqs(fixtures).await); + assert_indexed(fixtures, &[0]).await; + assert_eq!( + Counters { + recovered_objects: 1, + recoveries: 1, + ..create + }, + Counters::since(before) + ); + } + + async fn start_object_meets_another_writer() { + for epoch in [2, 3] { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let late = encode_object(chain_header(1, epoch, None), &[]) + .unwrap() + .bytes; + let before = Counters::sample(); + let error = open_meeting(&object_store, late).await.unwrap_err(); + let path = fixtures.object_path(1); + assert_start_fails( + &fixtures, + before, + Counters { + create_conflicts: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::WalObjectConflict { path: actual, .. } if *actual == path), + RetryHint::NonRetryable, + ) + .await; + } + } + + async fn start_object_meets_an_identical_object() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let same = encode_object(chain_header(1, 2, Some((0, 1))), &[]) + .unwrap() + .bytes; + let before = Counters::sample(); + let error = open_meeting(&object_store, same).await.unwrap_err(); + assert_start_fails( + &fixtures, + before, + Counters { + created_objects: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::UnconfirmedWalEpochStart { epoch: 2, .. }), + RetryHint::Retryable, + ) + .await; + } + + async fn start_object_create_has_an_unknown_outcome() { + let object_store = memory_store(); + let fixtures = put_chain_start(&object_store).await; + let io = Arc::new(LostResponseIo { + inner: ObjectStoreIo::new(object_store.clone(), PREFIX).unwrap(), + lose_next: AtomicBool::new(true), + }); + let before = Counters::sample(); + let error = ObjectStoreLogStore::open(io, &eager(), PREFIX.to_string()) + .await + .unwrap_err(); + assert_start_fails( + &fixtures, + before, + Counters { + create_failures: 1, + ..Counters::default() + }, + error, + |error| matches!(error, Error::WalObjectStore { .. }), + RetryHint::Retryable, + ) + .await; + } + } } From be91fb5ad2e990b12721f6a9b8286e3503548e5e Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 9 Oct 2026 07:01:51 +0000 Subject: [PATCH 114/133] test: support metasrv plugins in integration cluster builder (#9423) * test: support plugins and flownodes in integration cluster builder Allow the meta-srv mock and the tests-integration distributed cluster builder to install metasrv Plugins and to start flownodes that register with the mock metasrv, so integration tests can exercise plugin-provided DDL handlers and flows against a real multi-datanode cluster. The default path is unchanged: mock()/mock_with_client_channel_config() keep their signatures, and GreptimeDbClusterBuilder builds the metasrv plugin-free and starts no flownode unless with_plugins/with_flownodes is called. Signed-off-by: discord9 * test: narrow cluster plugins to metasrv injection Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: remove unnecessary plugin container mutability Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/meta-srv/src/mocks.rs | 29 +++++++++++++++++++ tests-integration/src/cluster.rs | 49 +++++++++++++++++++++++++++++++- 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/src/meta-srv/src/mocks.rs b/src/meta-srv/src/mocks.rs index 1a6fa306743..1afc5dbd334 100644 --- a/src/meta-srv/src/mocks.rs +++ b/src/meta-srv/src/mocks.rs @@ -21,6 +21,7 @@ use api::v1::meta::heartbeat_server::HeartbeatServer; use api::v1::meta::procedure_service_server::ProcedureServiceServer; use api::v1::meta::store_server::StoreServer; use client::client_manager::NodeClients; +use common_base::Plugins; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::etcd::EtcdStore; @@ -95,6 +96,7 @@ pub async fn mock( datanode_clients, in_memory, None, + Plugins::default(), ) .await } @@ -114,6 +116,30 @@ pub async fn mock_with_client_channel_config( datanode_clients, in_memory, Some(client_channel_config), + Plugins::default(), + ) + .await +} + +/// Builds a mock metasrv with the given [`Plugins`]. +/// +/// Same as [`mock`], except that the plugins are passed to the metasrv builder. +pub async fn mock_with_plugins( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + plugins: Plugins, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + None, + plugins, ) .await } @@ -125,6 +151,7 @@ async fn mock_inner( datanode_clients: Option>, in_memory: Option, client_channel_config: Option, + plugins: Plugins, ) -> MockInfo { let server_addr = opts.grpc.server_addr.clone(); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); @@ -152,6 +179,8 @@ async fn mock_inner( None => builder, }; + let builder = builder.plugins(plugins); + let metasrv = builder.build().await.unwrap(); metasrv.try_start().await.unwrap(); diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 21a0d5eba2d..76f198c9707 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -174,6 +174,7 @@ pub struct GreptimeDbClusterBuilder { frontend_auto_create_table: bool, shared_home_dir: Option>, meta_selector: Option, + plugins: Plugins, local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, } @@ -210,6 +211,7 @@ impl GreptimeDbClusterBuilder { frontend_auto_create_table: true, shared_home_dir: None, meta_selector: None, + plugins: Plugins::default(), local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), } @@ -284,6 +286,13 @@ impl GreptimeDbClusterBuilder { self } + /// Sets the [Plugins] used by the metasrv. + #[must_use] + pub fn with_plugins(mut self, plugins: Plugins) -> Self { + self.plugins = plugins; + self + } + /// Configure the frontend COPY sandbox for filesystem integration tests. pub fn with_local_file_access(mut self, access: LocalFileAccess) -> Self { self.local_file_access = access; @@ -297,6 +306,7 @@ impl GreptimeDbClusterBuilder { guards: Vec, ) -> GreptimeDbCluster { let datanodes = datanode_options.len(); + let channel_config = ChannelConfig::new().timeout(Some(Duration::from_secs(20))); let datanode_clients = Arc::new(NodeClients::new(channel_config)); @@ -322,12 +332,13 @@ impl GreptimeDbClusterBuilder { test_util::prepare_another_catalog_and_schema_with_kv_backend(self.kv_backend.clone()) .await; - let metasrv = meta_srv::mocks::mock( + let metasrv = meta_srv::mocks::mock_with_plugins( opt, self.kv_backend.clone(), self.meta_selector.clone(), Some(datanode_clients.clone()), None, + self.plugins.clone(), ) .await; @@ -692,6 +703,8 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { #[cfg(test)] #[allow(deprecated)] mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use api::v1::flow::FlowRequest; use api::v1::region::{ ListMetadataRequest, RegionRequest, RegionRequestHeader, region_request, @@ -700,12 +713,14 @@ mod tests { use client::Client; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; + use common_meta::ddl_manager::{DdlManager, DdlManagerConfigurator, DdlManagerConfiguratorRef}; use common_meta::key::TableMetadataManager; use common_meta::key::flow::FlowMetadataManager; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::{DatanodeManager, FlownodeManager}; use common_meta::peer::Peer; use flow::{FlownodeBuilder, FlownodeOptions, FlownodeServiceBuilder, FrontendClient}; + use meta_srv::metasrv::builder::DdlManagerConfigureContext; use super::*; @@ -882,6 +897,38 @@ mod tests { flownode } + struct TestDdlManagerConfigurator(Arc); + + #[async_trait::async_trait] + impl DdlManagerConfigurator for TestDdlManagerConfigurator { + async fn configure( + &self, + ddl_manager: DdlManager, + _ctx: DdlManagerConfigureContext, + ) -> std::result::Result { + self.0.store(true, Ordering::SeqCst); + Ok(ddl_manager) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_build_cluster_with_plugins() { + let configured = Arc::new(AtomicBool::new(false)); + let plugins = Plugins::default(); + plugins.insert::>(Arc::new( + TestDdlManagerConfigurator(configured.clone()), + )); + + let _cluster = GreptimeDbClusterBuilder::new("test_cluster_with_plugins") + .await + .with_datanodes(1) + .with_plugins(plugins) + .build(false) + .await; + + assert!(configured.load(Ordering::SeqCst)); + } + #[tokio::test(flavor = "multi_thread")] async fn test_region_requester_health_check_and_list_metadata() { // Arrange: configure active and inactive synthetic peers over the duplex harness. From 21536a092bc666ec758ef24c7e770f8a95304dde Mon Sep 17 00:00:00 2001 From: zhaiAohan <1573612521@qq.com> Date: Fri, 9 Oct 2026 07:05:01 +0000 Subject: [PATCH 115/133] build(deps): update opensrv-mysql to fix MariaDB 12.3 USE information_schema hang (#9357) Update opensrv-mysql from f3135e5 to 1f3d88c, which replies a proper EOF packet to COM_FIELD_LIST instead of a malformed OK packet that MariaDB 12.3 misparses and hangs on USE information_schema. Fixes #9158 Signed-off-by: tian1220A <1573612521@qq.com> --- Cargo.lock | 2 +- src/servers/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c08ac6f1ca8..c02c7fdcc60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9966,7 +9966,7 @@ dependencies = [ [[package]] name = "opensrv-mysql" version = "0.8.0" -source = "git+https://github.com/GreptimeTeam/opensrv?rev=f3135e5f00401e06075cf241c0cc26489e45b20a#f3135e5f00401e06075cf241c0cc26489e45b20a" +source = "git+https://github.com/GreptimeTeam/opensrv?rev=1f3d88ccaafb5721be6cfcc3ccbca48642426119#1f3d88ccaafb5721be6cfcc3ccbca48642426119" dependencies = [ "async-trait", "byteorder", diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index f9c3b45832f..8a2b6ab2149 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -90,7 +90,7 @@ notify.workspace = true object-pool = "0.5" once_cell.workspace = true # Wait for https://github.com/databendlabs/opensrv/pull/81 -opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "f3135e5f00401e06075cf241c0cc26489e45b20a" } +opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "1f3d88ccaafb5721be6cfcc3ccbca48642426119" } opentelemetry-proto.workspace = true operator.workspace = true otel-arrow-rust.workspace = true From 6da5352c53e3987c2dcae85140961cacf37b6857 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 9 Oct 2026 07:16:42 +0000 Subject: [PATCH 116/133] fix(promql): keep vector labels when a computed scalar is on the right (#9478) * fix(promql): keep vector labels when a computed scalar is on the right `metric + time()` or `metric * scalar(other)` planned the join with the scalar side's context, which has no tags, so every series collapsed into one unlabeled series. Use the vector side's context when the right operand is scalar-typed. Signed-off-by: Dennis Zhuang * test(promql): update results for labels kept with a computed scalar on the right v * pi() and v + time() now keep the vector's labels and time index name, as in Prometheus. Signed-off-by: Dennis Zhuang * fix(promql): qualify columns of an unnamed operand in vector joins A comparison filter's projection has no table name, so after the join its time index resolved unqualified and was ambiguous when the other operand had the same column. Also keep the empty-metric branch unchanged for scalar left operands. Signed-off-by: Dennis Zhuang * test(promql): cover comparison filters joined with a scalar Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/planner.rs | 12 +++ .../common/promql/computed_scalar_rhs.result | 97 +++++++++++++++++++ .../common/promql/computed_scalar_rhs.sql | 43 ++++++++ .../standalone/common/promql/math.result | 42 ++++---- .../standalone/common/promql/time_fn.result | 4 +- .../common/promql/timestamp_fn.result | 2 +- 6 files changed, 176 insertions(+), 24 deletions(-) create mode 100644 tests/cases/standalone/common/promql/computed_scalar_rhs.result create mode 100644 tests/cases/standalone/common/promql/computed_scalar_rhs.sql diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 2d52cfb83f7..a0ace9b9db6 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1257,6 +1257,18 @@ impl PromPlanner { } } else if right_is_empty_metric && !left_is_empty_metric { self.ctx = left_context.clone(); + } else if rhs.value_type() == ValueType::Scalar + && lhs.value_type() == ValueType::Vector + { + // A computed scalar on the right (`time()`, `scalar(...)`) has no labels; + // the result keeps the labels of the vector on the left. + self.ctx = left_context.clone(); + } + // An operand without a table name, such as a comparison filter's projection, is + // joined under the `""` qualifier; resolve the projected columns against it. + // Otherwise a time index present on both sides is ambiguous. + if self.ctx.table_name.is_none() { + self.ctx.table_name = Some(String::new()); } // Computed scalars reach this join path instead of the literal projection paths. // Broadcast them for arithmetic in the same way as literal scalars. diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.result b/tests/cases/standalone/common/promql/computed_scalar_rhs.result new file mode 100644 index 00000000000..96b925d1cf5 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.result @@ -0,0 +1,97 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +Affected Rows: 2 + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +Affected Rows: 1 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + ++------+---------------------+-------------------------------------------------+ +| host | ts | computed_scalar_lhs.val + .time / Float64(1000) | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + ++------+---------------------+-------------------------------------------------+ +| host | ts | .time / Float64(1000) + computed_scalar_lhs.val | ++------+---------------------+-------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 61.0 | +| b | 1970-01-01T00:01:00 | 62.0 | ++------+---------------------+-------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + ++------+---------------------+-----------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val * computed_scalar_rhs.scalar(val) | ++------+---------------------+-----------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+-----------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + ++------+---------------------+--------------------------------------------------------------+ +| host | ts | computed_scalar_lhs.val < .time / Float64(1000) * Float64(2) | ++------+---------------------+--------------------------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1.0 | ++------+---------------------+--------------------------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + ++------+---------------------+---------------------------------------------+ +| host | ts | power(.val,computed_scalar_rhs.scalar(val)) | ++------+---------------------+---------------------------------------------+ +| a | 1970-01-01T00:01:00 | 1.0 | +| b | 1970-01-01T00:01:00 | 1024.0 | ++------+---------------------+---------------------------------------------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + ++------+---------------------+----------------------------------------+ +| host | ts | computed_scalar_rhs.scalar(val) * .val | ++------+---------------------+----------------------------------------+ +| a | 1970-01-01T00:01:00 | 10.0 | +| b | 1970-01-01T00:01:00 | 20.0 | ++------+---------------------+----------------------------------------+ + +drop table computed_scalar_lhs; + +Affected Rows: 0 + +drop table computed_scalar_rhs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/computed_scalar_rhs.sql b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql new file mode 100644 index 00000000000..57cc431c1f9 --- /dev/null +++ b/tests/cases/standalone/common/promql/computed_scalar_rhs.sql @@ -0,0 +1,43 @@ +-- A vector combined with a computed scalar on the right keeps the vector's +-- labels, like a scalar on the left does. Expected results are from Prometheus 3. + +create table computed_scalar_lhs ( + ts timestamp time index, + val double, + host string primary key +); + +create table computed_scalar_rhs ( + ts timestamp time index, + val double, + host string primary key +); + +insert into computed_scalar_lhs values + (0, 1.0, 'a'), + (0, 2.0, 'b'); + +insert into computed_scalar_rhs values + (0, 10.0, 'x'); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs + time(); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') time() + computed_scalar_lhs; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs * scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') computed_scalar_lhs < bool (time() * 2); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') (time() >= computed_scalar_lhs) ^ scalar(computed_scalar_rhs); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (60, 60, '1s') scalar(computed_scalar_rhs) * (time() >= computed_scalar_lhs); + +drop table computed_scalar_lhs; + +drop table computed_scalar_rhs; diff --git a/tests/cases/standalone/common/promql/math.result b/tests/cases/standalone/common/promql/math.result index a6271c158cb..e1b884c50f1 100644 --- a/tests/cases/standalone/common/promql/math.result +++ b/tests/cases/standalone/common/promql/math.result @@ -55,15 +55,15 @@ TQL EVAL (0, 20, '5s') pi() + 1; -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 20, '5s') angles{unit="radians"} * pi(); -+---------------------+--------------------+ -| time | angles.val * .pi() | -+---------------------+--------------------+ -| 1970-01-01T00:00:00 | 0.0 | -| 1970-01-01T00:00:05 | 4.934813740258847 | -| 1970-01-01T00:00:10 | 9.869627480517694 | -| 1970-01-01T00:00:15 | 14.80444122077654 | -| 1970-01-01T00:00:20 | 19.73925496103539 | -+---------------------+--------------------+ ++---------+---------------------+--------------------+ +| unit | ts | angles.val * .pi() | ++---------+---------------------+--------------------+ +| radians | 1970-01-01T00:00:00 | 0.0 | +| radians | 1970-01-01T00:00:05 | 4.934813740258847 | +| radians | 1970-01-01T00:00:10 | 9.869627480517694 | +| radians | 1970-01-01T00:00:15 | 14.80444122077654 | +| radians | 1970-01-01T00:00:20 | 19.73925496103539 | ++---------+---------------------+--------------------+ -- Test rad() function (converts degrees to radians) -- SQLNESS SORT_RESULT 3 1 @@ -310,18 +310,18 @@ TQL EVAL (25, 35, '5s') sgn(angles - 42.5); -- SQLNESS SORT_RESULT 3 1 TQL EVAL (0, 35, '5s') sgn(angles{unit="radians"} - pi()); -+---------------------+----------------------------+ -| time | signum(angles.val - .pi()) | -+---------------------+----------------------------+ -| 1970-01-01T00:00:00 | -1.0 | -| 1970-01-01T00:00:05 | -1.0 | -| 1970-01-01T00:00:10 | 1.0 | -| 1970-01-01T00:00:15 | 1.0 | -| 1970-01-01T00:00:20 | 1.0 | -| 1970-01-01T00:00:25 | 1.0 | -| 1970-01-01T00:00:30 | 1.0 | -| 1970-01-01T00:00:35 | 1.0 | -+---------------------+----------------------------+ ++---------------------+----------------------------+---------+ +| ts | signum(angles.val - .pi()) | unit | ++---------------------+----------------------------+---------+ +| 1970-01-01T00:00:00 | -1.0 | radians | +| 1970-01-01T00:00:05 | -1.0 | radians | +| 1970-01-01T00:00:10 | 1.0 | radians | +| 1970-01-01T00:00:15 | 1.0 | radians | +| 1970-01-01T00:00:20 | 1.0 | radians | +| 1970-01-01T00:00:25 | 1.0 | radians | +| 1970-01-01T00:00:30 | 1.0 | radians | +| 1970-01-01T00:00:35 | 1.0 | radians | ++---------------------+----------------------------+---------+ -- SQLNESS SORT_RESULT 3 1 TQL EVAL (25, 35, '5s') sgn(angles) * angles; diff --git a/tests/cases/standalone/common/promql/time_fn.result b/tests/cases/standalone/common/promql/time_fn.result index 25e0cd11030..f947f2fb726 100644 --- a/tests/cases/standalone/common/promql/time_fn.result +++ b/tests/cases/standalone/common/promql/time_fn.result @@ -150,7 +150,7 @@ tql eval (1, 2, '1s') time() == bool metrics; tql eval (1, 2, '1s') metrics + time(); +---------------------+-------------------------------------+ -| time | metrics.val + .time / Float64(1000) | +| ts | metrics.val + .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 2.0 | | 1970-01-01T00:00:02 | 4.0 | @@ -170,7 +170,7 @@ tql eval (1, 2, '1s') metrics == time(); tql eval (1, 2, '1s') metrics == bool time(); +---------------------+-------------------------------------+ -| time | metrics.val = .time / Float64(1000) | +| ts | metrics.val = .time / Float64(1000) | +---------------------+-------------------------------------+ | 1970-01-01T00:00:01 | 1.0 | | 1970-01-01T00:00:02 | 1.0 | diff --git a/tests/cases/standalone/common/promql/timestamp_fn.result b/tests/cases/standalone/common/promql/timestamp_fn.result index fc4e9c46fa6..32f5ba11cfb 100644 --- a/tests/cases/standalone/common/promql/timestamp_fn.result +++ b/tests/cases/standalone/common/promql/timestamp_fn.result @@ -95,7 +95,7 @@ tql eval (0, 60, '30s') timestamp(timestamp_test) > bool 30; tql eval (0, 60, '30s') timestamp(timestamp_test) - time(); +---------------------+----------------------------------------------+ -| time | timestamp_test.value - .time / Float64(1000) | +| ts | timestamp_test.value - .time / Float64(1000) | +---------------------+----------------------------------------------+ | 1970-01-01T00:00:00 | 0.0 | | 1970-01-01T00:00:30 | -29.0 | From 2ebb03fbf6be2fb3135c60d9f95b3c099e48537a Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Fri, 9 Oct 2026 07:26:52 +0000 Subject: [PATCH 117/133] fix(otlp): convert boolean metric attributes into label values (#9468) Signed-off-by: fengyue-xve <1759519137@qq.com> --- src/servers/src/otlp/metrics.rs | 70 ++++++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index fd496db5a98..7a485545bdd 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -349,12 +349,13 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } -/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +/// Non-scalar values (arrays, maps, bytes) are not representable as tags. fn scalar_value_string(value: Option<&AnyValue>) -> Option { match value.and_then(|v| v.value.as_ref())? { any_value::Value::StringValue(s) => Some(s.clone()), any_value::Value::IntValue(v) => Some(v.to_string()), any_value::Value::DoubleValue(v) => Some(v.to_string()), + any_value::Value::BoolValue(v) => Some(v.to_string()), _ => None, } } @@ -1571,6 +1572,15 @@ mod tests { } } + fn bool_keyvalue(key: &str, value: bool) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(Val::BoolValue(value)), + }), + } + } + fn descriptor_ctx() -> OtlpMetricCtx { OtlpMetricCtx { resource_info: true, @@ -1734,6 +1744,64 @@ mod tests { assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); } + #[test] + fn test_boolean_attributes_keep_series_distinct() { + set_default_prefix(None).unwrap(); + let request = metrics_request(vec![Metric { + name: "bool_repro".to_string(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![ + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", false)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(100.0)), + ..Default::default() + }, + NumberDataPoint { + attributes: vec![bool_keyvalue("failed", true)], + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + value: Some(Value::AsDouble(7.0)), + ..Default::default() + }, + ], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + is_monotonic: true, + })), + ..Default::default() + }]); + let conversion = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + let rows = conversion + .requests + .inserts + .iter() + .find(|insert| insert.table_name == "bool_repro_total") + .expect("missing bool_repro_total table") + .rows + .as_ref() + .unwrap(); + assert_eq!(2, rows.rows.len()); + let failed = rows + .schema + .iter() + .position(|column| column.column_name == "failed") + .expect("boolean attribute `failed` was dropped from the label set"); + let values = rows + .rows + .iter() + .map(|row| row.values[failed].value_data.clone()) + .collect::>(); + assert_eq!( + vec![ + Some(ValueData::StringValue("false".to_string())), + Some(ValueData::StringValue("true".to_string())), + ], + values + ); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); From 4ade13d9b664670fb20951bdb45a8298d43eb854 Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Fri, 9 Oct 2026 07:39:03 +0000 Subject: [PATCH 118/133] perf(catalog): optimize single-table lookups in information_schema.partitions (#9473) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(catalog): optimize single-table lookups in information_schema.partitions Signed-off-by: 风月 <1759519137@qq.com> * style: apply rustfmt formatting to partitions tests Signed-off-by: 风月 <1759519137@qq.com> --------- Signed-off-by: 风月 <1759519137@qq.com> --- .../information_schema/partitions.rs | 295 ++++++++++++++++++ src/catalog/src/system_schema/predicate.rs | 100 ++++++ 2 files changed, 395 insertions(+) diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index d2cce93150a..1bfe3f634fd 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -254,6 +254,51 @@ impl InformationSchemaPartitionsBuilder { let predicates = Predicates::from_scan_request(&request); + // Fast path: when the predicates pin both the schema and the table name + // with equality, resolve the table directly instead of enumerating all + // tables in all schemas and fetching partition metadata for each of them. + if let (Some(schema_name), Some(table_name)) = ( + predicates + .find_eq_value(TABLE_SCHEMA) + .and_then(Value::as_string), + predicates + .find_eq_value(TABLE_NAME) + .and_then(Value::as_string), + ) { + if let Some(table) = catalog_manager + .table(&catalog_name, &schema_name, &table_name, None) + .await? + { + let table_info = table.table_info(); + + // The enumeration path below skips temporary tables, keep the + // behavior identical. + if table_info.table_type != TableType::Temporary { + let mut table_partitions = partition_manager + .batch_find_table_partitions(&[table_info.table_id()]) + .await + .context(FindPartitionsSnafu)?; + + let partitions = table_partitions + .remove(&table_info.table_id()) + .unwrap_or_default(); + + self.add_partitions( + &predicates, + &table_info, + &catalog_name, + &schema_name, + &table_info.name, + &partitions, + ); + } + } + + // The pinned table fully determines the result: the table above was + // found (or not) unambiguously, so no other table can contribute rows. + return self.finish(); + } + for schema_name in catalog_manager.schema_names(&catalog_name, None).await? { let table_info_stream = catalog_manager .tables(&catalog_name, &schema_name, None) @@ -421,3 +466,253 @@ impl DfPartitionStream for InformationSchemaPartitions { )) } } + +#[cfg(test)] +mod tests { + use api::v1::meta::Peer; + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::key::schema_name::SchemaNameKey; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::wal_provider::RegionWalOptions; + use datafusion::logical_expr::{BinaryExpr, Expr, Operator, col, lit}; + use futures::TryStreamExt; + use store_api::storage::RegionId; + + use super::*; + use crate::information_schema::NoopInformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + + const SCHEMA_1: &str = "partitions_schema_1"; + const SCHEMA_2: &str = "partitions_schema_2"; + const PARTITIONED_TABLE: &str = "partitioned_table"; + const MISSING_TABLE: &str = "missing_table"; + const TABLE_ID_1: TableId = 1; + const TABLE_ID_2: TableId = 2; + const REGION_NUMBER_1: u32 = 1; + const REGION_NUMBER_2: u32 = 2; + const REGION_NUMBER_7: u32 = 7; + + fn kv_catalog_manager() -> Arc { + let backend = Arc::new(MemoryKvBackend::default()); + let layered_cache_builder = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()); + let fundamental_cache_registry = build_fundamental_cache_registry(backend.clone()); + let layered_cache_registry = Arc::new( + with_default_composite_cache_registry( + layered_cache_builder.add_cache_registry(fundamental_cache_registry), + ) + .unwrap() + .build(), + ); + + KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + backend, + layered_cache_registry, + ) + .build() + } + + fn region_route(table_id: TableId, region_number: u32) -> RegionRoute { + RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + leader_peer: Some(Peer { + id: 1, + addr: "127.0.0.1:3001".to_string(), + }), + ..Default::default() + } + } + + fn table_info(schema: &str, table: &str, table_id: TableId) -> TableInfo { + // The catalog defaults to the default catalog name. + let mut table_info = + common_meta::key::test_utils::new_test_table_info_with_name(table_id, table); + table_info.schema_name = schema.to_string(); + table_info + } + + /// Creates a catalog with two partitioned tables named `PARTITIONED_TABLE` + /// in two different schemas: + /// + /// - `SCHEMA_1.PARTITIONED_TABLE` has two regions. + /// - `SCHEMA_2.PARTITIONED_TABLE` has one region. + async fn setup_catalog() -> Arc { + let manager = kv_catalog_manager(); + let metadata = manager.table_metadata_manager_ref().clone(); + + for schema in [SCHEMA_1, SCHEMA_2] { + metadata + .schema_manager() + .create(SchemaNameKey::new(DEFAULT_CATALOG_NAME, schema), None, true) + .await + .unwrap(); + } + + for (table_id, schema, region_numbers) in [ + (TABLE_ID_1, SCHEMA_1, vec![REGION_NUMBER_1, REGION_NUMBER_2]), + (TABLE_ID_2, SCHEMA_2, vec![REGION_NUMBER_7]), + ] { + let routes = region_numbers + .into_iter() + .map(|region_number| region_route(table_id, region_number)) + .collect(); + metadata + .create_table_metadata( + table_info(schema, PARTITIONED_TABLE, table_id), + TableRouteValue::physical(routes), + RegionWalOptions::default(), + ) + .await + .unwrap(); + } + + manager + } + + /// Scans `information_schema.partitions` and returns the total row count + /// and the pretty printed batches. + async fn scan_partitions( + manager: &Arc, + filters: Vec, + ) -> (usize, String) { + let table = manager + .table( + DEFAULT_CATALOG_NAME, + INFORMATION_SCHEMA_NAME, + PARTITIONS, + None, + ) + .await + .unwrap() + .unwrap(); + let stream = table + .scan_to_stream(ScanRequest { + filters, + ..Default::default() + }) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let rows = batches.iter().map(|batch| batch.num_rows()).sum(); + let output = batches + .iter() + .map(|batch| batch.pretty_print()) + .collect::>() + .join("\n"); + (rows, output) + } + + fn eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::Eq, + Box::new(lit(value)), + )) + } + + fn not_eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::NotEq, + Box::new(lit(value)), + )) + } + + #[tokio::test] + async fn test_scan_pinned_table_skips_catalog_enumeration() { + let manager: Arc = setup_catalog().await; + + // No predicates: every partitioned table is listed. + let (rows, output) = scan_partitions(&manager, vec![]).await; + assert_eq!(3, rows, "{output}"); + + // Pins both the schema and the table name, so the scan resolves the + // table directly instead of enumerating all catalog entries. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, PARTITIONED_TABLE), + ], + ) + .await; + assert_eq!(2, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(!output.contains(SCHEMA_2), "{output}"); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_1) + .as_u64() + .to_string() + ), + "{output}" + ); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_2) + .as_u64() + .to_string() + ), + "{output}" + ); + + // The fallback enumeration path must return the identical result. The + // inequality on `table_name` doesn't pin a single name, so no direct + // lookup happens here. + let (fallback_rows, fallback_output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + not_eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(rows, fallback_rows, "{fallback_output}"); + assert_eq!(output, fallback_output); + + // Only the table name is pinned: the schema can't be resolved and both + // tables of that name are kept. + let (rows, output) = + scan_partitions(&manager, vec![eq_filter(TABLE_NAME, PARTITIONED_TABLE)]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + + // An unknown table in a pinned schema yields no rows. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(0, rows, "{output}"); + } + + #[tokio::test] + async fn test_scan_equality_under_disjunction_falls_back_to_enumeration() { + let manager: Arc = setup_catalog().await; + + // `table_schema = SCHEMA_1 OR table_name = PARTITIONED_TABLE` matches + // the tables in both schemas: equality nested under a disjunction must + // not pin a single table. + let filter = Expr::BinaryExpr(BinaryExpr::new( + Box::new(eq_filter(TABLE_SCHEMA, SCHEMA_1)), + Operator::Or, + Box::new(eq_filter(TABLE_NAME, PARTITIONED_TABLE)), + )); + let (rows, output) = scan_partitions(&manager, vec![filter]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + } +} diff --git a/src/catalog/src/system_schema/predicate.rs b/src/catalog/src/system_schema/predicate.rs index 7957ba454d4..cb71e51ed0e 100644 --- a/src/catalog/src/system_schema/predicate.rs +++ b/src/catalog/src/system_schema/predicate.rs @@ -214,6 +214,22 @@ impl Predicate { _ => None, } } + + /// Collects the values that constrain `column` to be equal to them at this + /// predicate's top-level conjunctions. + /// + /// Constraints under `Or`/`Not` are ignored because they don't guarantee + /// that `column` equals a single value. + fn collect_eq_values<'a>(&'a self, column: &str, values: &mut Vec<&'a Value>) { + match self { + Predicate::Eq(c, v) if c == column => values.push(v), + Predicate::And(left, right) => { + left.collect_eq_values(column, values); + right.collect_eq_values(column, values); + } + _ => {} + } + } } /// Perform SQL left LIKE right, return `None` if fail to evaluate. @@ -269,6 +285,25 @@ impl Predicates { } } + /// Returns the single equality value that constrains `column` in the + /// top-level conjunctions of the predicates. + /// + /// Returns `None` if there is no equality constraint on `column`, or if + /// the predicates constrain it with multiple distinct values. + pub fn find_eq_value(&self, column: &str) -> Option<&Value> { + let mut values = Vec::new(); + for predicate in &self.predicates { + predicate.collect_eq_values(column, &mut values); + } + + let (first, rest) = values.split_first()?; + if rest.iter().all(|v| v == first) { + Some(*first) + } else { + None + } + } + /// Evaluate the predicates with the row. /// returns true when all the predicates are satisfied or can't be evaluated. pub fn eval(&self, row: &[(&str, &Value)]) -> bool { @@ -556,6 +591,71 @@ mod tests { ); } + #[test] + fn test_find_eq_value() { + fn eq(name: &str, value: &str) -> Predicate { + Predicate::Eq(name.to_string(), Value::from(value)) + } + + fn and(left: Predicate, right: Predicate) -> Predicate { + Predicate::And(Box::new(left), Box::new(right)) + } + + fn or(left: Predicate, right: Predicate) -> Predicate { + Predicate::Or(Box::new(left), Box::new(right)) + } + + fn predicates(predicates: Vec) -> Predicates { + Predicates { predicates } + } + + // No predicates, no equality value. + assert!(predicates(vec![]).find_eq_value("a").is_none()); + + // No equality on the column. + let p = predicates(vec![Predicate::NotEq("a".to_string(), Value::from("x"))]); + assert!(p.find_eq_value("a").is_none()); + assert!(p.find_eq_value("b").is_none()); + + // Simple equality. + let p = predicates(vec![eq("a", "x")]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + assert!(p.find_eq_value("b").is_none()); + + // Duplicated equality on the same value. + let p = predicates(vec![and(eq("a", "x"), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Conflicting equalities reject the lookup. + let p = predicates(vec![and(eq("a", "x"), eq("a", "y"))]); + assert!(p.find_eq_value("a").is_none()); + + // Conflicting equalities in separate predicates reject the lookup too. + let p = predicates(vec![eq("a", "x"), eq("a", "y")]); + assert!(p.find_eq_value("a").is_none()); + + // Equality nested in deeper conjunctions. + let p = predicates(vec![and(and(eq("b", "y"), eq("a", "x")), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Equality under a disjunction or a negation doesn't constrain the column. + let p = predicates(vec![or(eq("a", "x"), eq("a", "x"))]); + assert!(p.find_eq_value("a").is_none()); + let p = predicates(vec![Predicate::Not(Box::new(eq("a", "x")))]); + assert!(p.find_eq_value("a").is_none()); + + // Equality value in a top-level conjunction composed of separate predicates. + let p = predicates(vec![ + eq("a", "x"), + Predicate::Like("b".to_string(), "%y".to_string(), false), + ]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + } + #[test] fn test_predicates_eval_row() { let wrong_row = [ From f1a27ab599a615362aceaaea913ea3ce91879bfe Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Fri, 9 Oct 2026 09:58:41 +0000 Subject: [PATCH 119/133] refactor(function): rename uadf_impl to udf_impl in aggregate functions (#9482) refactor: rename uadf_impl to udf_impl in aggregate functions Signed-off-by: fengyue-xve <1759519137@qq.com> --- src/common/function/src/aggrs/geo.rs | 4 ++-- src/common/function/src/aggrs/geo/encoding.rs | 2 +- src/common/function/src/aggrs/geo/geo_path.rs | 2 +- src/common/function/src/aggrs/vector.rs | 6 +++--- src/common/function/src/aggrs/vector/avg.rs | 2 +- src/common/function/src/aggrs/vector/product.rs | 2 +- src/common/function/src/aggrs/vector/sum.rs | 2 +- 7 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/common/function/src/aggrs/geo.rs b/src/common/function/src/aggrs/geo.rs index c76fe7348f2..abe9070a0a0 100644 --- a/src/common/function/src/aggrs/geo.rs +++ b/src/common/function/src/aggrs/geo.rs @@ -21,7 +21,7 @@ pub(crate) struct GeoFunction; impl GeoFunction { pub fn register(registry: &FunctionRegistry) { - registry.register_aggr(encoding::JsonEncodePathAccumulator::uadf_impl()); - registry.register_aggr(geo_path::GeoPathAccumulator::uadf_impl()); + registry.register_aggr(encoding::JsonEncodePathAccumulator::udf_impl()); + registry.register_aggr(geo_path::GeoPathAccumulator::udf_impl()); } } diff --git a/src/common/function/src/aggrs/geo/encoding.rs b/src/common/function/src/aggrs/geo/encoding.rs index 67dfd493a77..5b4591bb6d8 100644 --- a/src/common/function/src/aggrs/geo/encoding.rs +++ b/src/common/function/src/aggrs/geo/encoding.rs @@ -47,7 +47,7 @@ impl JsonEncodePathAccumulator { Self::default() } - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { create_udaf( JSON_ENCODE_PATH_NAME, // Input types: lat, lng, timestamp diff --git a/src/common/function/src/aggrs/geo/geo_path.rs b/src/common/function/src/aggrs/geo/geo_path.rs index c4e366c83b0..fe8d8bae21c 100644 --- a/src/common/function/src/aggrs/geo/geo_path.rs +++ b/src/common/function/src/aggrs/geo/geo_path.rs @@ -47,7 +47,7 @@ impl GeoPathAccumulator { Self::default() } - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { create_udaf( GEO_PATH_NAME, // Input types: lat, lng, timestamp diff --git a/src/common/function/src/aggrs/vector.rs b/src/common/function/src/aggrs/vector.rs index 03489a51d44..0c0c70f7360 100644 --- a/src/common/function/src/aggrs/vector.rs +++ b/src/common/function/src/aggrs/vector.rs @@ -25,8 +25,8 @@ pub(crate) struct VectorFunction; impl VectorFunction { pub fn register(registry: &FunctionRegistry) { - registry.register_aggr(VectorSum::uadf_impl()); - registry.register_aggr(VectorProduct::uadf_impl()); - registry.register_aggr(VectorAvg::uadf_impl()); + registry.register_aggr(VectorSum::udf_impl()); + registry.register_aggr(VectorProduct::udf_impl()); + registry.register_aggr(VectorAvg::udf_impl()); } } diff --git a/src/common/function/src/aggrs/vector/avg.rs b/src/common/function/src/aggrs/vector/avg.rs index 82e6a28059f..a8af8cf6d0c 100644 --- a/src/common/function/src/aggrs/vector/avg.rs +++ b/src/common/function/src/aggrs/vector/avg.rs @@ -39,7 +39,7 @@ pub struct VectorAvg { impl VectorAvg { /// Create a new `AggregateUDF` for the `vec_avg` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), diff --git a/src/common/function/src/aggrs/vector/product.rs b/src/common/function/src/aggrs/vector/product.rs index 8178b74f369..4f0e9bb6217 100644 --- a/src/common/function/src/aggrs/vector/product.rs +++ b/src/common/function/src/aggrs/vector/product.rs @@ -39,7 +39,7 @@ pub struct VectorProduct { impl VectorProduct { /// Create a new `AggregateUDF` for the `vec_product` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), diff --git a/src/common/function/src/aggrs/vector/sum.rs b/src/common/function/src/aggrs/vector/sum.rs index b1b68906ded..466fc42cccb 100644 --- a/src/common/function/src/aggrs/vector/sum.rs +++ b/src/common/function/src/aggrs/vector/sum.rs @@ -39,7 +39,7 @@ pub struct VectorSum { impl VectorSum { /// Create a new `AggregateUDF` for the `vec_sum` aggregate function. - pub fn uadf_impl() -> AggregateUDF { + pub fn udf_impl() -> AggregateUDF { let signature = Signature::one_of( vec![ TypeSignature::Exact(vec![DataType::Utf8]), From 764ab8bc1a1e8d9e579576109deb7d554e523fb9 Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Fri, 9 Oct 2026 09:59:20 +0000 Subject: [PATCH 120/133] test(promql): cover label_replace normalization with distinct series (#9485) The duplicate-label-set rejection must not fire when a label rewrite keeps every series distinct through another label. Normalizing mixed histogram boundary representations (`10.0` -> `10`) while `instance` still differs is the false positive reported in #9444. The check already accounts for it since it became a per-timestamp row count instead of a planning-time rejection, but no fixture covered the positive case. Add one, together with the aggregation variant that produces a genuine duplicate and must still fail. Signed-off-by: fengyue-xve <1759519137@qq.com> --- .../standalone/common/promql/label.result | 49 +++++++++++++++++++ .../cases/standalone/common/promql/label.sql | 29 +++++++++++ 2 files changed, 78 insertions(+) diff --git a/tests/cases/standalone/common/promql/label.result b/tests/cases/standalone/common/promql/label.result index 3fe8d559221..f6e169eebe3 100644 --- a/tests/cases/standalone/common/promql/label.result +++ b/tests/cases/standalone/common/promql/label.result @@ -484,3 +484,52 @@ DROP TABLE test; Affected Rows: 0 +-- Issue 9444 -- +-- Normalizing mixed histogram boundary representations must not be rejected while +-- another label still tells the series apart. +CREATE TABLE demo_bucket ( + ts timestamp(3) time index, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY(instance, le), +); + +Affected Rows: 0 + +INSERT INTO TABLE demo_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +Affected Rows: 7 + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (3000, 3000, '5s') label_replace(demo_bucket, "le", "$1", "le", "([0-9]+)[.]0+"); + ++---------------------+------+------+----------+ +| ts | val | le | instance | ++---------------------+------+------+----------+ +| 1970-01-01T00:50:00 | 40.0 | 10 | a | +| 1970-01-01T00:50:00 | 40.0 | 10 | b | +| 1970-01-01T00:50:00 | 50.0 | +Inf | a | +| 1970-01-01T00:50:00 | 50.0 | +Inf | b | +| 1970-01-01T00:50:00 | 50.0 | 100 | a | +| 1970-01-01T00:50:00 | 50.0 | 100 | b | +| 1970-01-01T00:50:00 | 50.0 | 1000 | b | ++---------------------+------+------+----------+ + +-- Aggregating `instance` away first makes the same rewrite collapse two series into +-- one label set, which is a genuine duplicate and must still fail. +TQL EVAL (3000, 3000, '5s') label_replace(sum by (le) (demo_bucket), "le", "$1", "le", "([0-9]+)[.]0+"); + +Error: 3001(EngineExecuteQuery), Execution error: vector cannot contain metrics with the same labelset + +DROP TABLE demo_bucket; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/label.sql b/tests/cases/standalone/common/promql/label.sql index 6c2d64d42ac..e269b6a1783 100644 --- a/tests/cases/standalone/common/promql/label.sql +++ b/tests/cases/standalone/common/promql/label.sql @@ -144,3 +144,32 @@ TQL EVAL (0, 1, '5s') test{job=""}; TQL EVAL (0, 1, '5s') test{job!=""}; DROP TABLE test; + +-- Issue 9444 -- +-- Normalizing mixed histogram boundary representations must not be rejected while +-- another label still tells the series apart. +CREATE TABLE demo_bucket ( + ts timestamp(3) time index, + instance STRING, + le STRING, + val DOUBLE, + PRIMARY KEY(instance, le), +); + +INSERT INTO TABLE demo_bucket VALUES + (3000000, 'a', '10', 40), + (3000000, 'a', '100', 50), + (3000000, 'a', '+Inf', 50), + (3000000, 'b', '10.0', 40), + (3000000, 'b', '100.0', 50), + (3000000, 'b', '1000.0', 50), + (3000000, 'b', '+Inf', 50); + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (3000, 3000, '5s') label_replace(demo_bucket, "le", "$1", "le", "([0-9]+)[.]0+"); + +-- Aggregating `instance` away first makes the same rewrite collapse two series into +-- one label set, which is a genuine duplicate and must still fail. +TQL EVAL (3000, 3000, '5s') label_replace(sum by (le) (demo_bucket), "le", "$1", "le", "([0-9]+)[.]0+"); + +DROP TABLE demo_bucket; From 0dde3a7b7091381856d9e974cfc7aaf033891aab Mon Sep 17 00:00:00 2001 From: Niladri Adhikary <91966855+niladrix719@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:03:36 +0000 Subject: [PATCH 121/133] fix(promql): return no point for zero-interval irate/idelta (#9455) fix(promql): irate/idelta return no point for a zero-ms sample gap Signed-off-by: Niladri Adhikary --- src/promql/src/functions/idelta.rs | 35 ++++++++- .../promql/sub_millisecond_interval.result | 77 +++++++++++++++++++ .../promql/sub_millisecond_interval.sql | 30 ++++++++ 3 files changed, 140 insertions(+), 2 deletions(-) create mode 100644 tests/cases/standalone/common/promql/sub_millisecond_interval.result create mode 100644 tests/cases/standalone/common/promql/sub_millisecond_interval.sql diff --git a/src/promql/src/functions/idelta.rs b/src/promql/src/functions/idelta.rs index 35ce44ba270..5734733234d 100644 --- a/src/promql/src/functions/idelta.rs +++ b/src/promql/src/functions/idelta.rs @@ -140,8 +140,13 @@ impl IDelta { let last_offset = ts_offset + last_position; let prev_offset = ts_offset + prev_position; - let sampled_interval = - (ts_values[last_offset] - ts_values[prev_offset]) as f64 / 1000.0; + let sampled_interval_ms = ts_values[last_offset] - ts_values[prev_offset]; + // Sub-millisecond samples round to the same ms; skip instead of dividing by zero. + if sampled_interval_ms == 0 { + result_builder.append_null(); + continue; + } + let sampled_interval = sampled_interval_ms as f64 / 1000.0; let last_value = value_values[value_offset + last_position]; let prev_value = value_values[value_offset + prev_position]; @@ -263,4 +268,30 @@ mod test { vec![Some(3.0), None, None], ); } + + #[test] + fn idelta_skips_samples_in_same_millisecond() { + // The last two samples of the first range share 999ms; the second range is 1ms apart. + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([ + 999i64, 999, 999, 1000, + ])); + let values_array = Arc::new(Float64Array::from_iter([10.0, 13.0, 999.0, 1000.0])); + let ranges = [(0, 3), (2, 2)]; + + simple_range_udf_runner( + IDelta::::scalar_udf(), + RangeArray::from_ranges(ts_array.clone(), ranges).unwrap(), + RangeArray::from_ranges(values_array.clone(), ranges).unwrap(), + vec![], + vec![None, Some(1.0)], + ); + + simple_range_udf_runner( + IDelta::::scalar_udf(), + RangeArray::from_ranges(ts_array, ranges).unwrap(), + RangeArray::from_ranges(values_array, ranges).unwrap(), + vec![], + vec![None, Some(1000.0)], + ); + } } diff --git a/tests/cases/standalone/common/promql/sub_millisecond_interval.result b/tests/cases/standalone/common/promql/sub_millisecond_interval.result new file mode 100644 index 00000000000..6461c8e0b0d --- /dev/null +++ b/tests/cases/standalone/common/promql/sub_millisecond_interval.result @@ -0,0 +1,77 @@ +-- The last two 'sub_ms' samples are 400ns apart and share one millisecond, which +-- irate and idelta cannot resolve, so they return no point instead of inf. +-- The 'one_ms' samples are 1ms apart and still produce a rate. +CREATE TABLE sub_ms_interval ( + job STRING, + `value` DOUBLE, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY(job) +); + +Affected Rows: 0 + +INSERT INTO sub_ms_interval VALUES + ('sub_ms', 999, 999999900), + ('sub_ms', 10, 999999100), + ('sub_ms', 13, 999999500), + ('one_ms', 10, 998000000), + ('one_ms', 13, 999000000); + +Affected Rows: 5 + +SELECT * FROM sub_ms_interval ORDER BY job, ts; + ++--------+-------+-------------------------------+ +| job | value | ts | ++--------+-------+-------------------------------+ +| one_ms | 10.0 | 1970-01-01T00:00:00.998 | +| one_ms | 13.0 | 1970-01-01T00:00:00.999 | +| sub_ms | 10.0 | 1970-01-01T00:00:00.999999100 | +| sub_ms | 13.0 | 1970-01-01T00:00:00.999999500 | +| sub_ms | 999.0 | 1970-01-01T00:00:00.999999900 | ++--------+-------+-------------------------------+ + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + ++---------------------+----------------------------+--------+ +| ts | prom_irate(ts_range,value) | job | ++---------------------+----------------------------+--------+ +| 1970-01-01T00:00:01 | 3000.0 | one_ms | ++---------------------+----------------------------+--------+ + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + ++---------------------+-----------------------------+--------+ +| ts | prom_idelta(ts_range,value) | job | ++---------------------+-----------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | one_ms | ++---------------------+-----------------------------+--------+ + +ADMIN FLUSH_TABLE('sub_ms_interval'); + ++--------------------------------------+ +| ADMIN FLUSH_TABLE('sub_ms_interval') | ++--------------------------------------+ +| 0 | ++--------------------------------------+ + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + ++---------------------+----------------------------+--------+ +| ts | prom_irate(ts_range,value) | job | ++---------------------+----------------------------+--------+ +| 1970-01-01T00:00:01 | 3000.0 | one_ms | ++---------------------+----------------------------+--------+ + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + ++---------------------+-----------------------------+--------+ +| ts | prom_idelta(ts_range,value) | job | ++---------------------+-----------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | one_ms | ++---------------------+-----------------------------+--------+ + +DROP TABLE sub_ms_interval; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/sub_millisecond_interval.sql b/tests/cases/standalone/common/promql/sub_millisecond_interval.sql new file mode 100644 index 00000000000..d643479390d --- /dev/null +++ b/tests/cases/standalone/common/promql/sub_millisecond_interval.sql @@ -0,0 +1,30 @@ +-- The last two 'sub_ms' samples are 400ns apart and share one millisecond, which +-- irate and idelta cannot resolve, so they return no point instead of inf. +-- The 'one_ms' samples are 1ms apart and still produce a rate. +CREATE TABLE sub_ms_interval ( + job STRING, + `value` DOUBLE, + ts TIMESTAMP(9) TIME INDEX, + PRIMARY KEY(job) +); + +INSERT INTO sub_ms_interval VALUES + ('sub_ms', 999, 999999900), + ('sub_ms', 10, 999999100), + ('sub_ms', 13, 999999500), + ('one_ms', 10, 998000000), + ('one_ms', 13, 999000000); + +SELECT * FROM sub_ms_interval ORDER BY job, ts; + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + +ADMIN FLUSH_TABLE('sub_ms_interval'); + +TQL EVAL (1, 1, '1s') irate(sub_ms_interval[5s]); + +TQL EVAL (1, 1, '1s') idelta(sub_ms_interval[5s]); + +DROP TABLE sub_ms_interval; From d1151b0ed6c9034963a939c22dd92fcf64ecc16f Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Fri, 9 Oct 2026 10:26:29 +0000 Subject: [PATCH 122/133] fix(metasrv): stop parking the heartbeat response on the region supervisor queue (#9474) * fix(meta-srv): stop parking the heartbeat response on the region supervisor queue Signed-off-by: lczllx <2181719471@qq.com> * test(meta-srv): cover the heartbeat response path at the handler level Signed-off-by: lczllx <2181719471@qq.com> --------- Signed-off-by: lczllx <2181719471@qq.com> --- src/meta-srv/src/handler/failure_handler.rs | 117 +++++++++++++------- src/meta-srv/src/metrics.rs | 3 + src/meta-srv/src/region/supervisor.rs | 77 ++++++++++++- 3 files changed, 154 insertions(+), 43 deletions(-) diff --git a/src/meta-srv/src/handler/failure_handler.rs b/src/meta-srv/src/handler/failure_handler.rs index 5f1448d3ab6..d1d287d13ff 100644 --- a/src/meta-srv/src/handler/failure_handler.rs +++ b/src/meta-srv/src/handler/failure_handler.rs @@ -52,9 +52,10 @@ impl HeartbeatHandler for RegionFailureHandler { return Ok(HandleControl::Continue); }; + // Never blocks: the region lease has already been decided earlier in the + // chain, so a busy supervisor must not withhold the heartbeat response. self.heartbeat_acceptor - .accept(DatanodeHeartbeat::from(stat)) - .await; + .accept(DatanodeHeartbeat::from(stat)); Ok(HandleControl::Continue) } @@ -62,6 +63,8 @@ impl HeartbeatHandler for RegionFailureHandler { #[cfg(test)] mod tests { + use std::time::Duration; + use api::v1::meta::HeartbeatRequest; use common_catalog::consts::default_engine; use common_meta::datanode::{RegionManifestInfo, RegionStat, Stat}; @@ -73,7 +76,45 @@ mod tests { use crate::handler::{HeartbeatAccumulator, HeartbeatHandler}; use crate::metasrv::builder::MetasrvBuilder; use crate::region::supervisor::tests::new_test_supervisor; - use crate::region::supervisor::{Event, HeartbeatAcceptor}; + use crate::region::supervisor::{DatanodeHeartbeat, Event, HeartbeatAcceptor}; + + fn new_region_stat(region_id: u64) -> RegionStat { + RegionStat { + id: RegionId::from_u64(region_id), + rcus: 0, + wcus: 0, + approximate_bytes: 0, + engine: default_engine().to_string(), + role: RegionRole::Follower, + num_rows: 0, + memtable_size: 0, + manifest_size: 0, + sst_size: 0, + sst_num: 0, + index_size: 0, + region_manifest: RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + min_timestamp: None, + max_timestamp: None, + } + } + + fn new_stat() -> Stat { + Stat { + id: 42, + region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], + timestamp_millis: 1000, + ..Default::default() + } + } #[tokio::test] async fn test_handle_heartbeat() { @@ -85,40 +126,7 @@ mod tests { let metasrv = builder.build().await.unwrap(); let mut ctx = metasrv.new_ctx(); let acc = &mut HeartbeatAccumulator::default(); - fn new_region_stat(region_id: u64) -> RegionStat { - RegionStat { - id: RegionId::from_u64(region_id), - rcus: 0, - wcus: 0, - approximate_bytes: 0, - engine: default_engine().to_string(), - role: RegionRole::Follower, - num_rows: 0, - memtable_size: 0, - manifest_size: 0, - sst_size: 0, - sst_num: 0, - index_size: 0, - region_manifest: RegionManifestInfo::Mito { - manifest_version: 0, - flushed_entry_id: 0, - file_removed_cnt: 0, - }, - data_topic_latest_entry_id: 0, - metadata_topic_latest_entry_id: 0, - written_bytes: 0, - query_cpu_time: 0, - query_scanned_bytes: 0, - min_timestamp: None, - max_timestamp: None, - } - } - acc.stat = Some(Stat { - id: 42, - region_stats: vec![new_region_stat(1), new_region_stat(2), new_region_stat(3)], - timestamp_millis: 1000, - ..Default::default() - }); + acc.stat = Some(new_stat()); handler.handle(req, &mut ctx, acc).await.unwrap(); let (tx, rx) = oneshot::channel(); @@ -126,4 +134,39 @@ mod tests { let detector = rx.await.unwrap(); assert_eq!(detector.iter().collect::>().len(), 3); } + + #[tokio::test] + async fn test_handle_heartbeat_does_not_block_on_full_supervisor_queue() { + // Regression test for #9419: a full queue used to park the handler. + // Keep the supervisor alive, or the queue reports Closed. + let (_supervisor, sender) = new_test_supervisor(); + // Built directly, so the supervisor stays undrained. + let handler = RegionFailureHandler { + heartbeat_acceptor: HeartbeatAcceptor::new(sender.clone()), + }; + + let stat = new_stat(); + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(DatanodeHeartbeat::from(&stat))) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + + let builder = MetasrvBuilder::new(); + let metasrv = builder.build().await.unwrap(); + let mut ctx = metasrv.new_ctx(); + let acc = &mut HeartbeatAccumulator::default(); + // Without a stat the handler returns before the acceptor. + acc.stat = Some(stat); + + let result = tokio::time::timeout( + Duration::from_secs(5), + handler.handle(&HeartbeatRequest::default(), &mut ctx, acc), + ) + .await; + assert!(result.is_ok()); + } } diff --git a/src/meta-srv/src/metrics.rs b/src/meta-srv/src/metrics.rs index 1cf32194302..c92e6f31fff 100644 --- a/src/meta-srv/src/metrics.rs +++ b/src/meta-srv/src/metrics.rs @@ -48,6 +48,9 @@ lazy_static! { /// Heartbeat received by metasrv. pub static ref METRIC_META_HEARTBEAT_RECV: IntCounterVec = register_int_counter_vec!("greptime_meta_heartbeat_recv", "heartbeats received by metasrv", &["pusher_key"]).unwrap(); + /// Heartbeats dropped by metasrv. + pub static ref METRIC_META_HEARTBEAT_DROPPED: IntCounterVec = + register_int_counter_vec!("greptime_meta_heartbeat_dropped", "heartbeats dropped by metasrv", &["reason"]).unwrap(); /// The migration execute histogram. pub static ref METRIC_META_REGION_MIGRATION_EXECUTE: HistogramVec = register_histogram_vec!("greptime_meta_region_migration_execute", "meta region migration execute", &["state"]).unwrap(); diff --git a/src/meta-srv/src/region/supervisor.rs b/src/meta-srv/src/region/supervisor.rs index 56a9fe4ee57..99c68ac23c8 100644 --- a/src/meta-srv/src/region/supervisor.rs +++ b/src/meta-srv/src/region/supervisor.rs @@ -35,6 +35,7 @@ use common_time::util::current_time_millis; use futures::{StreamExt, TryStreamExt}; use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; +use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::oneshot; use tokio::time::{MissedTickBehavior, interval, interval_at}; @@ -43,6 +44,7 @@ use crate::discovery::utils::accept_ingest_workload; use crate::error::{self, Result}; use crate::failure_detector::PhiAccrualFailureDetectorOptions; use crate::metasrv::{RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef}; +use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::manager::{ RegionMigrationManagerRef, RegionMigrationTriggerReason, SubmitRegionMigrationTaskResult, }; @@ -356,10 +358,29 @@ impl HeartbeatAcceptor { Self { sender } } - /// Accepts heartbeats from datanodes. - pub(crate) async fn accept(&self, heartbeat: DatanodeHeartbeat) { - if let Err(err) = self.sender.send(Event::HeartbeatArrived(heartbeat)).await { - error!(err; "RegionSupervisor has stop receiving heartbeat."); + /// Accepts heartbeats from datanodes without waiting for queue capacity: a full queue + /// drops the heartbeat, counted in [`METRIC_META_HEARTBEAT_DROPPED`]. A dropped + /// heartbeat costs the failure detector one sample. + pub(crate) fn accept(&self, heartbeat: DatanodeHeartbeat) { + let datanode_id = heartbeat.datanode_id; + match self.sender.try_send(Event::HeartbeatArrived(heartbeat)) { + Ok(()) => {} + Err(TrySendError::Full(_)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .inc(); + warn!( + "Dropping heartbeat because the region supervisor event queue is full, datanode_id: {}, queue_capacity: {}", + datanode_id, + self.sender.max_capacity() + ); + } + Err(TrySendError::Closed(err)) => { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["closed"]) + .inc(); + error!(err; "RegionSupervisor has stop receiving heartbeat."); + } } } } @@ -899,14 +920,15 @@ pub(crate) mod tests { use tokio::time::sleep; use super::RegionSupervisorSelector; + use crate::metrics::METRIC_META_HEARTBEAT_DROPPED; use crate::procedure::region_migration::RegionMigrationTriggerReason; use crate::procedure::region_migration::manager::{ RegionMigrationManager, SubmitRegionMigrationTaskResult, }; use crate::procedure::region_migration::test_util::TestingEnv; use crate::region::supervisor::{ - DatanodeHeartbeat, Event, RegionFailureDetectorControl, RegionSupervisor, - RegionSupervisorTicker, + DatanodeHeartbeat, Event, HeartbeatAcceptor, RegionFailureDetectorControl, + RegionSupervisor, RegionSupervisorTicker, }; use crate::selector::test_utils::{RandomNodeSelector, new_test_selector_context}; @@ -940,6 +962,14 @@ pub(crate) mod tests { ) } + fn filler_heartbeat() -> DatanodeHeartbeat { + DatanodeHeartbeat { + datanode_id: 1, + regions: vec![], + timestamp: 0, + } + } + #[tokio::test] async fn test_heartbeat() { let (mut supervisor, sender) = new_test_supervisor(); @@ -1372,4 +1402,39 @@ pub(crate) mod tests { assert!(supervisor.failure_detector.contains(&detecting_region)); assert!(supervisor.failover_counts.contains_key(&detecting_region)); } + + #[tokio::test] + async fn test_heartbeat_accept_does_not_block_on_full_queue() { + // Regression test for #9419: a full queue used to park the caller. + // The supervisor is never run, so the queue stays full. + let (mut supervisor, sender) = new_test_supervisor(); + + let mut queued = 0; + while sender + .try_send(Event::HeartbeatArrived(filler_heartbeat())) + .is_ok() + { + queued += 1; + } + assert_eq!(queued, sender.max_capacity()); + assert_eq!(sender.capacity(), 0); + + let acceptor = HeartbeatAcceptor::new(sender.clone()); + let dropped = || { + METRIC_META_HEARTBEAT_DROPPED + .with_label_values(&["full"]) + .get() + }; + let dropped_before = dropped(); + // The full queue drops the heartbeat and counts it. + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + + // Free one slot: the next heartbeat is accepted. + supervisor.receiver.try_recv().unwrap(); + acceptor.accept(filler_heartbeat()); + assert_eq!(dropped(), dropped_before + 1); + assert_eq!(sender.capacity(), 0); + } } From c33a6cfa56e8802773f1d4be1fcac7364f544799 Mon Sep 17 00:00:00 2001 From: Aum Date: Fri, 9 Oct 2026 10:41:40 +0000 Subject: [PATCH 123/133] fix(promql): keep metric name for comparison operators (#9332) * fix(promql): keep metric name for comparison operators `collect_metric_names` dropped the response `__name__` for every binary operator except `and` / `or` / `unless`, so comparisons such as `node_load1 > 5` came back without the metric name. Grafana rendered them as `{instance="..."}` and alert templates reading `$labels.__name__` got nothing. Follow Prometheus' `shouldDropMetricName` / `resultMetric`: arithmetic operators drop the name, comparisons keep the left-hand name, except when the comparison returns a bool or is a one-to-one `on(...)` match. A `group_left` / `group_right` modifier does not reduce the labels, so the name is kept there. Both the HTTP PromQL API and the gRPC gateway share `retrieve_metric_name_and_result_type`, so one fix covers both. Signed-off-by: Aumnertic * docs: document binary_keeps_metric_name Signed-off-by: Aumnertic * fix(promql): handle on(__name__), scalar lhs and group_right Signed-off-by: Aumnertic * chore(promql): correct result-side docs and drop unreachable String arm Signed-off-by: Aumnertic * fix(promql): drop metric name when explicitly ignored Handle ignoring(__name__) for one-to-one comparisons and cover comparison, group_left, group_right and unless behavior. Signed-off-by: Aumnertic --------- Signed-off-by: Aumnertic --- src/servers/src/http/prometheus.rs | 218 +++++++++++++++++++++++++++-- 1 file changed, 207 insertions(+), 11 deletions(-) diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 420a0c8f4b4..da89d2a7b6c 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -49,11 +49,11 @@ use futures::future::join_all; use futures::{StreamExt, TryStreamExt}; use itertools::Itertools; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; -use promql_parser::parser::token::{self}; +use promql_parser::parser::token::TokenType; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ - AggregateExpr, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, MatrixSelector, ParenExpr, - SubqueryExpr, UnaryExpr, VectorSelector, + AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromqlExpr, LabelModifier, + MatrixSelector, ParenExpr, SubqueryExpr, UnaryExpr, VectorMatchCardinality, VectorSelector, }; use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery, QueryLanguageParser, QueryStatement}; use serde::de::{self, MapAccess, Visitor}; @@ -1373,6 +1373,72 @@ fn promql_expr_to_metric_name(expr: &PromqlExpr) -> Option { } } +/// Whether the metric name survives a binary operation; see +/// [`binary_metric_name_side`] for the operand that supplies it. +/// +/// Follows Prometheus' `shouldDropMetricName` and `resultMetric`: set operators +/// and comparisons keep the name, while arithmetic operators drop it. A +/// comparison also drops it when it returns a bool or is a one-to-one +/// `on(...)` match that does not keep `__name__` or an `ignoring(...)` match +/// that excludes it. Prometheus keeps only the `on(...)` labels or deletes the +/// `ignoring(...)` labels for one-to-one matches. +fn binary_keeps_metric_name(op: &TokenType, modifier: Option<&BinModifier>) -> bool { + if op.is_set_operator() { + return true; + } + + if !op.is_comparison_operator() { + return false; + } + + let Some(m) = modifier else { + return true; + }; + + if m.return_bool { + return false; + } + + if m.card == VectorMatchCardinality::OneToOne { + match &m.matching { + Some(LabelModifier::Include(labels)) + if !labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + Some(LabelModifier::Exclude(labels)) + if labels.labels.contains(&METRIC_NAME.to_string()) => + { + return false; + } + _ => {} + } + } + + true +} + +/// The operand whose labels supply the result metric name of a binary +/// operation. +/// +/// Prometheus builds the result from the side whose labels it carries, i.e. the +/// "many" side of the match: `group_right` makes the right-hand operand that +/// side, and a scalar left-hand operand (for example `0.5 < a`) leaves the +/// vector on the right. Every other case takes the left-hand operand. +fn binary_metric_name_side<'a>( + lhs: &'a PromqlExpr, + rhs: &'a PromqlExpr, + modifier: Option<&BinModifier>, +) -> &'a PromqlExpr { + let group_right = + modifier.is_some_and(|m| matches!(m.card, VectorMatchCardinality::OneToMany(_))); + if group_right || matches!(lhs.value_type(), ValueType::Scalar) { + rhs + } else { + lhs + } +} + /// Recursively collect all metric names from a PromQL expression fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { match expr { @@ -1395,14 +1461,16 @@ fn collect_metric_names(expr: &PromqlExpr, metric_names: &mut HashSet) { collect_metric_names(expr, metric_names) } PromqlExpr::Unary(UnaryExpr { .. }) => metric_names.clear(), - PromqlExpr::Binary(BinaryExpr { lhs, op, .. }) => { - if matches!( - op.id(), - token::T_LAND // INTERSECT - | token::T_LOR // UNION - | token::T_LUNLESS // EXCEPT - ) { - collect_metric_names(lhs, metric_names) + PromqlExpr::Binary(BinaryExpr { + lhs, + op, + rhs, + modifier, + .. + }) => { + if binary_keeps_metric_name(op, modifier.as_ref()) { + let side = binary_metric_name_side(lhs, rhs, modifier.as_ref()); + collect_metric_names(side, metric_names) } else { metric_names.clear() } @@ -3211,6 +3279,134 @@ mod tests { expected_type: ValueType::Vector, should_error: false, }, + // Arithmetic operators always drop the metric name. + TestCase { + name: "arithmetic between metrics", + promql: "a + b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "arithmetic with scalar", + promql: "a * 2", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + // Comparisons keep the metric name from the side whose labels the + // result carries, i.e. the "many" side of the match: normally the + // left operand, but `group_right` takes it from the right operand and + // a scalar left operand leaves the vector on the right. The name is + // dropped when the comparison returns a bool or is a one-to-one + // `on(...)` match that does not keep `__name__` or an `ignoring(...)` + // match that excludes it. A `group_left`/`group_right` modifier does + // not reduce the labels, so the name is kept. + TestCase { + name: "bool comparison between metrics", + promql: "a > bool b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "bool comparison with scalar", + promql: "a > bool 0.5", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar", + promql: "a > 0.5", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with scalar on the left", + promql: "0.5 < a", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "equality comparison with scalar", + promql: "up == 0", + expected_metric: Some("up"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison between metrics", + promql: "a > b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring", + promql: "a > ignoring(x) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__", + promql: "a > ignoring(__name__) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_left", + promql: "a > ignoring(__name__) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with ignoring __name__ and group_right", + promql: "a > ignoring(__name__) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "unless with ignoring __name__", + promql: "a unless ignoring(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on", + promql: "a > on(x) b", + expected_metric: None, + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on keeping __name__", + promql: "a > on(__name__) b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_left", + promql: "a > on(x) group_left b", + expected_metric: Some("a"), + expected_type: ValueType::Vector, + should_error: false, + }, + TestCase { + name: "comparison with on and group_right", + promql: "a > on(x) group_right b", + expected_metric: Some("b"), + expected_type: ValueType::Vector, + should_error: false, + }, // Subquery cases TestCase { name: "basic subquery", From b3c756a7dd05cd7e1b6685489714c4a5795a33df Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 9 Oct 2026 10:45:15 +0000 Subject: [PATCH 124/133] perf(promql): reduce timestamp payload allocations and use narrow cursor scans (#9417) * perf(promql): zero-copy timestamp payload for millisecond zero-offset ranges RangeManipulate rebuilt the timestamp range payload per sample through wide-precision i128 conversion even for millisecond input with zero offset, where the input is already millisecond and nothing changes. That turned a previously zero-copy path into a per-sample loop plus a Vec allocation, regressing range-function queries. Reuse the input timestamp array directly when the time index is already millisecond and the offset is zero. Other units and non-zero offsets keep the wide-precision conversion (required so ns/us data near the native type limits stays correct under large offsets). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize timestamp payload with plain i64 RangeManipulate rebuilt the timestamp range payload per sample through wide-precision i128 conversion, turning a hot path into a scalar loop with a Vec allocation. Realistic timestamps and offsets stay far inside i64, so the wide path is unnecessary. - millisecond + zero offset reuses the input array with zero copy; - other units / offsets shift in one vectorized i64 pass (no i128, no per-sample Option, no overflow branch). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): check timestamp payload overflow via arrow try_unary Replace the silent wrapping shift with try_unary + checked_mul/checked_add. checked arithmetic lowers to a per-lane overflow flag so the single pass still vectorizes over the dense i64 buffer; any overflow short-circuits to an error instead of silently wrapping to a wrong timestamp. No i128, no per-sample Option. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize cursor scans with i64 fast path RangeManipulate::calculate_range and InstantManipulate::manipulate compared every sample through i128 wide-precision arithmetic. For realistic timestamps and offsets the shifted values stay far inside i64, so the wide path was pure overhead. Split both scans into an i64 fast path (checked arithmetic precomputed once per batch) and an i128 exact fallback. Any input that would overflow i64 falls back to the i128 scan, keeping window membership semantics identical. Verified: 47/47 range_manipulate + instant_manipulate tests, 275/275 promql crate tests, 1612 randomized cases byte-identical to i128 scan. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * refactor(promql): generalize cursor scans and share endpoint precheck Address review comments on PR #9417: - Generalize scan_ranges/scan_instants over Nanoseconds (i64/i128), eliminating the duplicated scan bodies. - Extract shared narrow_scan precheck used by both RangeManipulate and InstantManipulate, with identical checked endpoint logic. - Trim doc comments to one line; fix factual errors about fallback trigger (2262 for ms/s, not native type limits). Semantics unchanged: same cursor algorithm, same overflow fallback, same i64-vs-i128 selection, same stale/null/tie handling. Verified: 47/47 range_manipulate + instant_manipulate tests, clippy clean. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 * perf(promql): vectorize payload with try_unary and add scan carrier bench Replace the per-sample i128 payload loop with a single arrow try_unary pass: i128 intermediate holds any native timestamp shift, only the final millisecond result must fit i64. This is semantically equivalent to the original i128 implementation (no silent wrap, no per-sample branching). Add cursor_scan_carriers A/B benchmark: i64 fast path vs i128 exact fallback on four workloads. Add five deterministic differential tests pinning scan_ranges:: and scan_ranges:: to be pointwise equivalent on boundary shapes (sub-ms samples, fallback triggers, nanosecond extremes, seeded precision matrix). Bench (N=4096, 15s cadence): - payload wall_time 20x15s: 74.94us -> 48.98us (-35%) - payload wall_time 240x15s: 102.44us - scan i64 fast path: 10.97us vs i128 19.84us (1.81x) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): restore zero-copy payload for ms+offset=0 and fix comment Restore the zero-copy fast path for millisecond timestamps with zero offset, guarded by the input column's exact data type (Timestamp (Millisecond, None)) rather than self.time_unit. The old time_unit guard would accept a tz-carrying millisecond column and produce a dictionary value type disagreeing with the output schema; the datatype guard is the correct check. Also fix the comment: arrow's try_unary is a fallible kernel that LLVM cannot auto-vectorize; the win over the hand-written loop comes from eliminating the per-sample Option/Vec and branches, not SIMD. Add a regression test with a timezone-carrying millisecond index to prove it falls back to the try_unary path. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): replace try_unary with unary + Cell overflow flag Arrow's try_unary is a fallible kernel that LLVM cannot auto-vectorize. Replace it with the infallible unary kernel: the closure is a plain i64 -> i64 mapping that sets a Cell on overflow instead of returning Err per sample. The flag is checked once after the loop. This removes the per-lane Result construction and lets LLVM vectorize the shift+truncate sequence. Bench (N=4096, 15s cadence): - window4x15s: 42.12us -> 28.77us (-32%) - window20x15s: 48.98us -> 27.68us (-43%) - window240x15s: 102.44us -> 27.61us (-73%) - regular_sparse: 90.04us -> 29.20us (-68%) All cases now run at ~27-29us regardless of window size; the zero-copy path dominates for ms+offset=0. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): dispatch payload conversion on null_count; drop wall-time bench Null-free batches (dominant PromQL case) take arrow's vectorizable unary kernel with a Cell overflow flag; nullable batches take try_unary so a NULL slot's arbitrary underlying value cannot raise a false batch-wide overflow error. Add a regression test exercising both paths. Remove bench_range_manipulate_wall_time (all-ms zero-offset cases only exercised the zero-copy path, measuring nothing about the conversion) and its orphaned helpers; restore the cursor-scan bench helpers it shared. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): keep nullable timestamp cursor scans on i128 NULL slots may contain unordered arbitrary values, so the endpoint-only narrow-scan proof cannot bound their interior timestamp arithmetic. Both range and instant scans now attempt i64 only for null-free timestamp arrays and otherwise retain the existing exact i128 path. Add focused interior-NULL regressions through both active entrypoints, checking concrete selections and output against the existing wide scan. The null-free fast path and timestamp payload conversion are unchanged. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Signed-off-by: discord9 --- src/promql/benches/bench_range_fn.rs | 468 ++++--- .../src/extension_plan/instant_manipulate.rs | 156 ++- .../src/extension_plan/range_manipulate.rs | 1203 ++++++++++++++++- .../promql_long_range_rate/case.toml | 119 ++ 4 files changed, 1686 insertions(+), 260 deletions(-) create mode 100644 tests/perf/query_cases/promql_long_range_rate/case.toml diff --git a/src/promql/benches/bench_range_fn.rs b/src/promql/benches/bench_range_fn.rs index b4e026bae3d..51e0aefdc79 100644 --- a/src/promql/benches/bench_range_fn.rs +++ b/src/promql/benches/bench_range_fn.rs @@ -17,25 +17,14 @@ use std::sync::Arc; use criterion::{BenchmarkId, Criterion, criterion_group}; -use datafusion::arrow::array::{ - Array, ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray, -}; -use datafusion::arrow::datatypes::{ - DataType as ArrowDataType, Field as ArrowField, Int64Type, Schema, -}; -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::common::ToDFSchema; -use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; -use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; +use datafusion::arrow::array::{Array, Float64Array, TimestampMillisecondArray}; +use datafusion::arrow::datatypes::TimeUnit; +use datafusion::physical_plan::ColumnarValue; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::ScalarFunctionArgs; use datatypes::arrow::datatypes::{DataType, Field}; -use futures::StreamExt; -use promql::extension_plan::RangeManipulate; +use promql::extension_plan::RangeManipulateStream; use promql::functions::{ AbsentOverTime, Changes, CountOverTime, Delta, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, @@ -964,223 +953,294 @@ fn bench_edge_count_functions(c: &mut Criterion) { group.finish(); } -const RANGE_MANIPULATE_CADENCE_MS: i64 = 15_000; +/// One cursor-scan workload: the timeline, the window and the aligned instants to visit. +#[derive(Clone, Debug)] +struct CursorScanWorkload { + time_unit: TimeUnit, + offset_ms: i64, + timestamps: Vec, + window_ms: i64, + interval_ms: i64, + start_ms: i64, + end_ms: i64, +} -fn make_range_manipulate_batch(timestamps: Vec, field_count: usize) -> RecordBatch { - let mut fields = Vec::with_capacity(field_count + 1); - let mut columns: Vec = Vec::with_capacity(field_count + 1); - fields.push(ArrowField::new( - "timestamp", - ArrowDataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), - false, - )); - columns.push(Arc::new(TimestampMillisecondArray::from(timestamps)) as _); - - for field_index in 0..field_count { - fields.push(ArrowField::new( - format!("value_{field_index}"), - ArrowDataType::Float64, - false, - )); - let values = (0..columns[0].len()) - .map(|row| row as f64 + field_index as f64 * 0.01) - .collect::>(); - columns.push(Arc::new(Float64Array::from(values)) as _); +impl CursorScanWorkload { + /// The workload through the narrow `i64` carrier, or `None` when production would fall + /// back to `i128` for it. + /// + /// For the fallback cases the fast path is timed on a twin workload: the same shape + /// translated onto a timeline the narrow carrier holds, which keeps the number of sample + /// and window comparisons identical to the fallback run. + fn bench_scan_narrow(&self) -> Option> { + RangeManipulateStream::bench_scan_narrow( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) } - RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() -} + /// The workload through the exact `i128` carrier, the fallback production takes + /// whenever [`Self::bench_scan_narrow`] returns `None`. + fn bench_scan_wide(&self) -> Vec<(u32, u32)> { + RangeManipulateStream::bench_scan_wide( + self.time_unit, + self.offset_ms, + self.window_ms, + self.interval_ms, + &self.timestamps, + self.start_ms, + self.end_ms, + ) + } -async fn execute_and_drain( - plan: Arc, - task_ctx: Arc, -) -> Vec { - let mut stream = plan.execute(0, task_ctx).unwrap(); - let mut output = Vec::new(); - while let Some(batch) = stream.next().await { - output.push(batch.unwrap()); + fn evaluations(&self) -> usize { + ((self.end_ms - self.start_ms) / self.interval_ms) as usize + 1 + } + + fn label(&self) -> String { + format!( + "N{}_eval{}_window{}x{}ms", + self.timestamps.len(), + self.evaluations(), + self.window_ms / self.interval_ms, + self.interval_ms + ) } - output } -fn range_keys(batch: &RecordBatch, index: usize) -> Vec<(u32, u32)> { - let ranges = batch - .column(index) - .as_any() - .downcast_ref::>() - .unwrap() - .clone(); - RangeArray::try_new(ranges) - .unwrap() - .ranges() - .map(Option::unwrap) - .collect() +/// Cadence of the cursor-scan benchmarks: PromQL's usual 15s step. +const CURSOR_SCAN_CADENCE_MS: i64 = 15_000; + +/// Milliseconds of a one-hour window, the widest window benchmarked here. +const CURSOR_SCAN_HOUR_MS: i64 = 3_600_000; + +/// Builds a workload from a first timestamp and a sample count. +/// +/// `sample_spacing_native` is the distance between samples in native ticks: 15s worth, so +/// every native-precision workload has the same shape (one sample per cadence, 20 samples +/// per 300s window). The aligned instants start at the first shifted sample, so every +/// window overlaps the input the same way in both arms of the A/B comparison. +fn cursor_scan_workload( + time_unit: TimeUnit, + offset_ms: i64, + first_timestamp: i64, + samples: usize, + window_ms: i64, + sample_spacing_native: i64, +) -> CursorScanWorkload { + let interval_ms = CURSOR_SCAN_CADENCE_MS; + let timestamps = (0..samples as i64) + .map(|index| first_timestamp + index * sample_spacing_native) + .collect::>(); + // `Millisecond` instants of the first shifted sample, flooring like `calculate_range`. + let start_ms = shifted_floor_millis(time_unit, offset_ms, first_timestamp); + let evaluations = samples - 1; + CursorScanWorkload { + time_unit, + offset_ms, + timestamps, + window_ms, + interval_ms, + start_ms, + end_ms: start_ms + evaluations as i64 * interval_ms, + } } -fn brute_force_range_keys( - timestamps: &[i64], - evaluations: usize, - window_points: usize, -) -> Vec<(u32, u32)> { - let range = window_points as i64 * RANGE_MANIPULATE_CADENCE_MS; - (0..evaluations) - .map(|evaluation| { - let current = evaluation as i64 * RANGE_MANIPULATE_CADENCE_MS; - let mut offset = None; - let mut length = 0; - for (index, timestamp) in timestamps.iter().enumerate() { - if *timestamp > current - range && *timestamp <= current { - offset.get_or_insert(index as u32); - length += 1; - } - } - (offset.unwrap_or(0), length) - }) - .collect() +/// Native ticks of one benchmark cadence, i.e. 15s in the given precision. +fn cadence_in_native_ticks(time_unit: TimeUnit) -> i64 { + match time_unit { + TimeUnit::Second => CURSOR_SCAN_CADENCE_MS / 1_000, + TimeUnit::Millisecond => CURSOR_SCAN_CADENCE_MS, + TimeUnit::Microsecond => CURSOR_SCAN_CADENCE_MS * 1_000, + TimeUnit::Nanosecond => CURSOR_SCAN_CADENCE_MS * 1_000_000, + } } -fn validate_range_manipulate_output( - output: &[RecordBatch], - field_count: usize, - expected_range_keys: Option<&[(u32, u32)]>, -) { - let Some(expected_range_keys) = expected_range_keys else { - assert!(output.is_empty()); - return; +/// Floor of a shifted sample in milliseconds, in exact `i128` native precision. +fn shifted_floor_millis(time_unit: TimeUnit, offset_ms: i64, timestamp: i64) -> i64 { + let scale = match time_unit { + TimeUnit::Second => 1_000_000_000i128, + TimeUnit::Millisecond => 1_000_000i128, + TimeUnit::Microsecond => 1_000i128, + TimeUnit::Nanosecond => 1i128, + }; + i64::try_from( + ((timestamp as i128) * scale + (offset_ms as i128) * 1_000_000).div_euclid(1_000_000), + ) + .expect("the shifted millisecond must stay legal") +} + +/// A/B benchmark of the cursor-scan carriers behind `RangeManipulate::scan_ranges`. +/// +/// `RangeManipulate` scans the aligned instants in `i64` nanoseconds whenever +/// `narrow_scan` proves the batch and the bounds fit, and in exact `i128` nanoseconds +/// otherwise. Both arms below drive the production scan through the carrier the operator +/// would pick, so the numbers isolate the carrier and not the surrounding operator work: +/// +/// - `ms_zero_offset`: millisecond batch with no offset, the `i64` fast path. The same +/// workload is also timed through `i128` to price the fallback on this input. +/// - `ms_zero_offset_long_window`: same, with a one-hour window, where the scan advances +/// both cursors much further per evaluation. +/// - `ms_large_offset`: an offset whose nanosecond product overflows `i64`, so production +/// falls back to `i128`. The `i64` arm times the identical workload shape translated onto +/// a narrow timeline, which keeps the selections and thus the work identical. +/// - `ns_native_floor`: nanosecond batch at the `i64::MIN` floor, where the shifted start +/// bound underflows `i64` nanoseconds, so production falls back to `i128`. The `i64` arm +/// again times the translated narrow shape. +/// +/// Every arm asserts that both carriers selected the same ranges, so the comparison cannot +/// be won by doing less work. The fallback triggers themselves are pinned by the unit tests +/// in `range_manipulate.rs` (`scan_falls_back_to_i128_when_only_the_sample_tail_overflows` +/// and `scan_carriers_split_at_the_nanosecond_boundaries`). +fn bench_cursor_scan_carriers(c: &mut Criterion) { + const SAMPLES: usize = 4_096; + const SHORT_WINDOW_MS: i64 = 20 * CURSOR_SCAN_CADENCE_MS; + // `i64::MAX` milliseconds is ~9.22e18 ms; one microsecond past it in the nanosecond + // product makes `narrow_scan` reject the offset, while the shifted millisecond stays legal. + const LARGE_OFFSET_MS: i64 = 9_223_372_036_855; + const LARGE_OFFSET_BASE_MS: i64 = 1_000_000_000_000_000; + const NANOSECOND_FLOOR_SHIFT_NS: i64 = 9_223_372_036_000_000_000; + const NANOSECOND_FLOOR_SAMPLES: usize = 4_096; + + let ms_zero_offset = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + let ms_zero_offset_long_window = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + CURSOR_SCAN_HOUR_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The large offset is shifted inside one millisecond of the nanosecond ceiling, so the + // payload stays legal: the fallback is driven by the product, not by an illegal instant. + let ms_large_offset = cursor_scan_workload( + TimeUnit::Millisecond, + LARGE_OFFSET_MS, + LARGE_OFFSET_BASE_MS, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + // The narrow twin of `ms_large_offset`: the same batch shape without the offset. + let ms_large_offset_narrow = cursor_scan_workload( + TimeUnit::Millisecond, + 0, + 0, + SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Millisecond), + ); + + // Nanosecond batch on the native floor: samples at `i64::MIN + i * 15s`, whose shifted + // start bound underflows `i64` nanoseconds. + let nanosecond_floor_first = i64::MIN; + let nanosecond_floor = cursor_scan_workload( + TimeUnit::Nanosecond, + 0, + nanosecond_floor_first, + NANOSECOND_FLOOR_SAMPLES, + SHORT_WINDOW_MS, + cadence_in_native_ticks(TimeUnit::Nanosecond), + ); + // The narrow twin: the same shape translated by a multiple of one millisecond, so both + // the samples and the aligned instants keep their relative geometry. + let nanosecond_floor_narrow = CursorScanWorkload { + time_unit: TimeUnit::Nanosecond, + offset_ms: 0, + timestamps: nanosecond_floor + .timestamps + .iter() + .map(|timestamp| timestamp + NANOSECOND_FLOOR_SHIFT_NS) + .collect(), + window_ms: nanosecond_floor.window_ms, + interval_ms: nanosecond_floor.interval_ms, + start_ms: nanosecond_floor.start_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, + end_ms: nanosecond_floor.end_ms + NANOSECOND_FLOOR_SHIFT_NS / 1_000_000, }; - assert_eq!(output.len(), 1); - let batch = &output[0]; - assert_eq!(batch.num_rows(), expected_range_keys.len()); - - let timestamp_range_keys = range_keys(batch, field_count + 1); - assert_eq!(timestamp_range_keys, expected_range_keys); - - for field_index in 0..field_count { - assert_eq!(range_keys(batch, field_index + 1), expected_range_keys); - } -} - -fn bench_range_manipulate_wall_time(c: &mut Criterion) { - let mut group = c.benchmark_group("range_manipulate_wall_time"); - let primary_points = 4_096; - let sparse_evaluations = primary_points - 1; - let mut cases = vec![ + let cases = [ ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 4, - 1, + "ms_zero_offset", + ms_zero_offset.clone(), + ms_zero_offset, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 1, + "ms_zero_offset_long_window", + ms_zero_offset_long_window.clone(), + ms_zero_offset_long_window, false, ), ( - "one_field", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 240, - 1, - false, - ), - ( - "regular_sparse", - (0..sparse_evaluations) - .step_by(2) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - sparse_evaluations, - 20, - 1, - false, - ), - ( - "all_empty_pathological", - (0..primary_points) - .map(|point| { - (primary_points as i64 + 21 + point as i64) * RANGE_MANIPULATE_CADENCE_MS - }) - .collect::>(), - primary_points, - 20, - 1, + "ms_large_offset", + ms_large_offset_narrow, + ms_large_offset, true, ), ( - "multi_field_control", - (0..primary_points) - .map(|point| point as i64 * RANGE_MANIPULATE_CADENCE_MS) - .collect::>(), - primary_points, - 20, - 4, - false, + "ns_native_floor", + nanosecond_floor_narrow, + nanosecond_floor, + true, ), ]; - for (case_name, timestamps, evaluations, window_points, field_count, expect_empty) in - cases.drain(..) - { - let expected_range_keys = (!expect_empty) - .then(|| brute_force_range_keys(×tamps, evaluations, window_points)); - let input_batch = make_range_manipulate_batch(timestamps, field_count); - let input_rows = input_batch.num_rows(); - let input_schema = input_batch.schema(); - let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: input_schema.clone().to_dfschema_ref().unwrap(), - }); - let field_columns = (0..field_count) - .map(|field_index| format!("value_{field_index}")) - .collect::>(); - let logical_plan = RangeManipulate::new( - 0, - (evaluations as i64 - 1) * RANGE_MANIPULATE_CADENCE_MS, - RANGE_MANIPULATE_CADENCE_MS, - 0, - window_points as i64 * RANGE_MANIPULATE_CADENCE_MS, - "timestamp".to_string(), - field_columns, - logical_input, - ) - .unwrap(); - let physical_input: Arc = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![input_batch]], input_schema, None).unwrap(), - ))); - let execution_plan = logical_plan.to_execution_plan(physical_input); - let runtime = tokio::runtime::Runtime::new().unwrap(); - let task_ctx = datafusion::prelude::SessionContext::new().task_ctx(); + let mut group = c.benchmark_group("cursor_scan_carriers"); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2)); + group.sample_size(50); - // Validate the public operator output and RangeArray keys before timing. - let output = runtime.block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - validate_range_manipulate_output(&output, field_count, expected_range_keys.as_deref()); + for (case_name, narrow, wide, wide_falls_back) in cases { + // Both arms must be timed on work the operator would really do: the narrow arm in + // `i64`, and the wide arm in `i128` exactly when the narrow baselines overflow. + let narrow_expected = narrow + .bench_scan_narrow() + .unwrap_or_else(|| panic!("{case_name}: the i64 arm must fit the narrow carrier")); + assert_eq!( + wide.bench_scan_narrow().is_none(), + wide_falls_back, + "{case_name}: the i128 arm's fallback trigger" + ); + // Fairness check before timing: both carriers have to select the same ranges, which + // makes the sample and window comparisons of the two arms identical. + assert_eq!( + wide.bench_scan_wide(), + narrow_expected, + "{case_name}: the two carriers must select the same ranges" + ); + assert_eq!( + narrow_expected.len(), + narrow.evaluations(), + "{case_name}: one range per aligned evaluation instant" + ); - // This measures public RangeManipulate execution wall time, including stream draining. + group.throughput(criterion::Throughput::Elements( + narrow.timestamps.len() as u64 + )); group.bench_with_input( - BenchmarkId::new( - case_name, - format!("N{input_rows}_eval{evaluations}_window{window_points}x15s"), - ), - &(), - |b, _| { - b.iter(|| { - let output = runtime - .block_on(execute_and_drain(execution_plan.clone(), task_ctx.clone())); - std::hint::black_box(output); - }) + BenchmarkId::new(format!("{case_name}/i64"), narrow.label()), + &narrow, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_narrow())); + }, + ); + group.bench_with_input( + BenchmarkId::new(format!("{case_name}/i128"), wide.label()), + &wide, + |b, workload| { + b.iter(|| std::hint::black_box(workload.bench_scan_wide())); }, ); } @@ -1196,5 +1256,5 @@ criterion_group!( bench_rate_window_steps, bench_edge_count_functions, bench_extrema_functions, - bench_range_manipulate_wall_time + bench_cursor_scan_carriers ); diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 4085d382c37..4a9a3fe5fb8 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -46,6 +46,7 @@ use prost::Message; use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; +use crate::extension_plan::range_manipulate::{Nanoseconds, ScanNanoseconds, narrow_scan}; use crate::extension_plan::series_divide::SeriesDivide; use crate::extension_plan::{ METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset, @@ -691,6 +692,7 @@ impl InstantManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let len = timestamps.len(); // Shift the native-tick timeline in i128 before comparing samples. Doing @@ -735,28 +737,77 @@ impl InstantManipulateStream { let estimated_points = estimated_points as usize; let aligned_start = aligned_start as i64; let aligned_end = aligned_end as i64; - let mut take_indices = Vec::with_capacity(estimated_points); - let mut aligned_ts = Vec::with_capacity(estimated_points); - let mut cursor = 0; - for expected_ms in (aligned_start..=aligned_end).step_by(self.interval as usize) { - let expected = (expected_ms as i128) * 1_000_000; + // Narrow `i64` nanoseconds whenever the batch fits them, exact `i128` otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.lookback_delta, + timestamps, + aligned_start, + aligned_end, + ) + }) + .flatten(); + if let Some(scan) = &narrow { + debug_assert_eq!(i64::try_from(first_ns), Ok(scan.shift(timestamps[0]))); + debug_assert_eq!(i64::try_from(last_ns), Ok(scan.shift(timestamps[len - 1]))); + } + let (take_indices, aligned_ts) = match narrow { + Some(scan) => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &scan, + estimated_points, + &is_stale, + ), + None => self.scan_instants::( + timestamps, + aligned_start, + aligned_end, + &ScanNanoseconds::::wide(self.time_unit, self.offset, self.lookback_delta), + estimated_points, + &is_stale, + ), + }; + self.take_record_batch_optional(input, take_indices, aligned_ts) + } + + /// Chooses the sample of every aligned evaluation instant, in `Ns` nanoseconds. + /// + /// Keeps the first row among exact timestamp ties and the latest preceding row + /// otherwise; zero lookback admits only exact samples. Staleness is tested after + /// choosing: a selected stale marker suppresses this evaluation rather than + /// falling back to an older finite value. + fn scan_instants( + &self, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, + capacity: usize, + is_stale: impl Fn(usize) -> bool, + ) -> (Vec, Vec) { + let len = timestamps.len(); + let mut cursor = 0usize; + let mut take_indices = Vec::with_capacity(capacity); + let mut aligned_ts = Vec::with_capacity(capacity); + for expected_ms in (start..=end).step_by(self.interval as usize) { + let expected = scan.instant(expected_ms); let mut exact_candidate = None; - while cursor < len && to_nanoseconds(timestamps[cursor]) <= expected { - if to_nanoseconds(timestamps[cursor]) == expected && exact_candidate.is_none() { + while cursor < len && scan.shift(timestamps[cursor]) <= expected { + if scan.shift(timestamps[cursor]) == expected && exact_candidate.is_none() { exact_candidate = Some(cursor); } cursor += 1; } - // Keep the first row among exact timestamp ties; otherwise use the - // latest preceding row. Zero lookback admits only exact samples. - // Test staleness after choosing: a selected stale marker suppresses - // this evaluation rather than falling back to an older finite value. let Some(candidate) = exact_candidate.or_else(|| cursor.checked_sub(1)) else { continue; }; - let candidate_ts = to_nanoseconds(timestamps[candidate]); - let lower = expected - (self.lookback_delta as i128) * 1_000_000; - if (candidate_ts == expected || candidate_ts > lower) + let candidate_ts = scan.shift(timestamps[candidate]); + if (candidate_ts == expected || candidate_ts > scan.window_start(expected)) && candidate_ts <= expected && !is_stale(candidate) { @@ -764,7 +815,7 @@ impl InstantManipulateStream { aligned_ts.push(expected_ms); } } - self.take_record_batch_optional(input, take_indices, aligned_ts) + (take_indices, aligned_ts) } /// Helper function to apply "take" on record batch. @@ -1775,6 +1826,81 @@ mod test { .await; } + #[test] + fn nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + ), + Field::new("value", DataType::Float64, false), + ])); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(array), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 0, + start: 0, + end: 1_000, + lookback_delta: 1_000, + interval: 1_000, + time_index: 0, + time_unit: TimeUnit::Second, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, false), + ])), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + + let output = stream.manipulate(input).unwrap(); + let output_timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let output_values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let wide = stream.scan_instants::( + ×tamps, + 0, + 1_000, + &ScanNanoseconds::::wide(TimeUnit::Second, 0, 1_000), + 2, + |_| false, + ); + assert_eq!(output_timestamps.values(), &[0]); + assert_eq!(output_values.values(), &[1.0]); + assert_eq!(wide, (vec![0], vec![0])); + assert_eq!(wide.1, output_timestamps.values().to_vec()); + } + #[test] fn exact_ties_select_first_and_lookback_uses_latest() { for (values, expected_timestamp) in [ diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 8ea7ba6b191..38ecb7aa4db 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashSet; +use std::ops::{Add, Mul, Sub}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -20,7 +21,7 @@ use std::task::{Context, Poll}; use common_telemetry::{debug, warn}; use datafusion::arrow::array::{Array, ArrayRef, Int64Array, TimestampMillisecondArray}; use datafusion::arrow::compute; -use datafusion::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit}; +use datafusion::arrow::datatypes::{DataType, Field, Int64Type, SchemaRef, TimeUnit}; use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -702,33 +703,68 @@ impl RangeManipulateStream { new_columns[*index] = new_column; } - // The timestamp range payload is always millisecond ABI. Shift in wide - // native precision before truncating toward zero, preserving null validity. - let scale = nanoseconds_per_native_tick(self.time_unit); - let (timestamps, _) = timestamp_array_to_primitive(input.column(self.time_index)) - .ok_or_else(|| { + // The timestamp range payload is always millisecond ABI. A column is reusable as-is + // only when its own data type already is that ABI: the output schema fixes the + // dictionary value type, so a millisecond column carrying a timezone must be + // converted like any other native unit. + let ts_column = input.column(self.time_index); + let timestamp_values: ArrayRef = if self.offset == 0 + && matches!( + ts_column.data_type(), + DataType::Timestamp(TimeUnit::Millisecond, None) + ) { + ts_column.clone() + } else { + // Shift in wide native precision before truncating toward zero, preserving + // null validity. + let scale = nanoseconds_per_native_tick(self.time_unit); + let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; - let timestamp_values = timestamps - .values() - .iter() - .enumerate() - .map(|(index, timestamp)| { - if !input.column(self.time_index).is_valid(index) { - return Ok(None); + // Single pass over the timestamps: an i128 intermediate holds any native + // timestamp shift; only the final millisecond result must fit i64. Null-free + // batches (the dominant PromQL case) take arrow's infallible `unary` kernel, + // which LLVM can vectorize, and report overflow through a shared cell. But + // `unary` visits NULL slots' underlying values, so a NULL slot holding an + // arbitrary overflowing value would raise a false batch-wide error; nullable + // batches therefore take `try_unary`, which applies the closure to valid rows + // only. + let offset_ns = (self.offset as i128) * 1_000_000; + let timestamp_values = if timestamps.null_count() == 0 { + let overflow = std::cell::Cell::new(false); + let shifted = compute::unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + let shifted_ms = shifted_ns / 1_000_000; + match i64::try_from(shifted_ms) { + Ok(ms) => ms, + Err(_) => { + overflow.set(true); + 0 + } + } + }); + if overflow.get() { + return Err(DataFusionError::Execution( + "RangeManipulate timestamp payload overflow".into(), + )); } - let shifted_ns = (*timestamp as i128) * scale + (self.offset as i128) * 1_000_000; - i64::try_from(shifted_ns / 1_000_000) - .map(Some) - .map_err(|_| { + shifted + } else { + compute::try_unary::(×tamps, |v| { + let shifted_ns = (v as i128) * scale + offset_ns; + i64::try_from(shifted_ns / 1_000_000).map_err(|_| { ArrowError::ComputeError( "RangeManipulate timestamp payload overflow".into(), ) }) - }) - .collect::, _>>()?; - let timestamp_values = TimestampMillisecondArray::from(timestamp_values); - let ts_range_column = RangeArray::from_ranges(Arc::new(timestamp_values), ranges.clone()) + })? + }; + Arc::new(TimestampMillisecondArray::new( + timestamp_values.values().clone(), + timestamp_values.nulls().cloned(), + )) + }; + let ts_range_column = RangeArray::from_ranges(timestamp_values, ranges.clone()) .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? .into_dict(); new_columns.push(Arc::new(ts_range_column)); @@ -770,6 +806,7 @@ impl RangeManipulateStream { let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { DataFusionError::Execution("Time index column is not a timestamp".into()) })?; + let timestamp_null_count = timestamps.null_count(); let timestamps = timestamps.values(); let timestamp = |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000; @@ -799,36 +836,193 @@ impl RangeManipulateStream { // The intersection is within the declared i64 query bounds. let start = start as i64; let end = end as i64; - let mut ranges = Vec::new(); // Range membership is decided on shifted native ticks, before the // timestamp-range payload is converted to its millisecond ABI. This // keeps sub-millisecond samples distinct in a range; equal millisecond // payload values are not a reason to deduplicate input samples. // - // Calculate for every aligned timestamp (`curr_ts`), assuming ordered timestamps. - let mut left = 0usize; - let mut right = 0usize; - for curr_ts in (start..=end).step_by(self.interval as _) { - let start_ts = (curr_ts as i128) * 1_000_000 - (self.range as i128) * 1_000_000; - - while left < len && timestamp(left) <= start_ts { - left += 1; + // Calculate for every aligned timestamp (`curr_ts`), assuming ordered + // timestamps: the scan compares shifted samples in `i64` nanoseconds + // whenever they fit the batch, and in exact `i128` nanoseconds otherwise. + let narrow = (timestamp_null_count == 0) + .then(|| { + narrow_scan( + self.time_unit, + self.offset, + self.range, + timestamps, + start, + end, + ) + }) + .flatten(); + let ranges = match narrow { + Some(scan) => scan_ranges::(self.interval, timestamps, start, end, &scan), + None => { + let scan = ScanNanoseconds::::wide(self.time_unit, self.offset, self.range); + scan_ranges::(self.interval, timestamps, start, end, &scan) } - right = right.max(left); - while right < len && timestamp(right) <= (curr_ts as i128) * 1_000_000 { - right += 1; - } - - if left == right { - ranges.push((0, 0)); - } else { - ranges.push((left as _, (right - left) as _)); - } - } + }; Ok((ranges, (start, end))) } + + /// Cursor scan of `timestamps` through the narrow `i64` carrier, for benchmarks. + /// + /// `bench_range_fn` times the fast path directly, next to the exact `i128` fallback in + /// [`RangeManipulateStream::bench_scan_wide`]; production selects the carrier through + /// [`narrow_scan`] inside `calculate_range`. Returns `None` when that carrier is + /// rejected, which is exactly the condition that makes the operator fall back to `i128`. + #[doc(hidden)] + pub fn bench_scan_narrow( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option> { + let scan = narrow_scan(time_unit, offset, window, timestamps, start, end)?; + Some(scan_ranges::(interval, timestamps, start, end, &scan)) + } + + /// Cursor scan of `timestamps` through the exact `i128` carrier, for benchmarks. + /// + /// The carrier is exact for every input, so this times the fallback that + /// [`narrow_scan`] selects whenever the narrow baselines would overflow. + #[doc(hidden)] + pub fn bench_scan_wide( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Vec<(u32, u32)> { + scan_ranges::( + interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(time_unit, offset, window), + ) + } +} + +/// Cursor scan over the aligned evaluation points, in `Ns` nanoseconds. +fn scan_ranges( + interval: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + scan: &ScanNanoseconds, +) -> Vec<(u32, u32)> { + let len = timestamps.len(); + let mut ranges = Vec::new(); + let mut left = 0usize; + let mut right = 0usize; + for curr_ts in (start..=end).step_by(interval as _) { + let curr_ts_ns = scan.instant(curr_ts); + let start_ts = scan.window_start(curr_ts_ns); + + while left < len && scan.shift(timestamps[left]) <= start_ts { + left += 1; + } + right = right.max(left); + while right < len && scan.shift(timestamps[right]) <= curr_ts_ns { + right += 1; + } + + if left == right { + ranges.push((0, 0)); + } else { + ranges.push((left as _, (right - left) as _)); + } + } + ranges +} + +/// Cursor scan baselines in nanoseconds, shared with `InstantManipulate`. +/// +/// The narrow `i64` baselines come from [`narrow_scan`], the exact ones from +/// [`ScanNanoseconds::wide`]. +#[derive(Debug, Clone, Copy)] +pub(super) struct ScanNanoseconds { + /// Nanoseconds per timestamp tick. + scale: Ns, + /// The offset shift in nanoseconds. + offset_ns: Ns, + /// The range or lookback window in nanoseconds. + window_ns: Ns, +} + +/// Nanoseconds carrier of the cursor scans: `i64` where [`narrow_scan`] proved that +/// the batch fits, `i128` for the exact fallback. +pub(super) trait Nanoseconds: + Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl Nanoseconds for T where + T: Copy + PartialOrd + From + Add + Sub + Mul +{ +} + +impl ScanNanoseconds { + /// The `i128` baselines, exact for every input. + pub(super) fn wide(time_unit: TimeUnit, offset: Millisecond, window: Millisecond) -> Self { + Self { + scale: nanoseconds_per_native_tick(time_unit), + offset_ns: i128::from(offset) * 1_000_000, + window_ns: i128::from(window) * 1_000_000, + } + } +} + +impl ScanNanoseconds { + /// Shifted nanoseconds of an input timestamp. + pub(super) fn shift(&self, timestamp: i64) -> Ns { + Ns::from(timestamp) * self.scale + self.offset_ns + } + + /// Nanoseconds of an evaluation instant given in milliseconds. + pub(super) fn instant(&self, millis: Millisecond) -> Ns { + Ns::from(millis) * Ns::from(1_000_000) + } + + /// The exclusive lower boundary of the window ending at `instant`. + pub(super) fn window_start(&self, instant: Ns) -> Ns { + instant - self.window_ns + } +} + +/// The narrow `i64` baselines of a sorted batch, or `None` when they would overflow. +pub(super) fn narrow_scan( + time_unit: TimeUnit, + offset: Millisecond, + window: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, +) -> Option> { + let scale = i64::try_from(nanoseconds_per_native_tick(time_unit)).ok()?; + let offset_ns = offset.checked_mul(1_000_000)?; + let window_ns = window.checked_mul(1_000_000)?; + // Sorted samples, so the checked endpoints bound every interior shift and product. + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + shift(*timestamps.first()?)?; + shift(*timestamps.last()?)?; + // The checked `[start, end]` bounds every instant and every `instant - window`. + start.checked_mul(1_000_000)?.checked_sub(window_ns)?; + end.checked_mul(1_000_000)?.checked_sub(window_ns)?; + Some(ScanNanoseconds { + scale, + offset_ns, + window_ns, + }) } #[cfg(test)] @@ -859,6 +1053,9 @@ mod test { const TIME_INDEX_COLUMN: &str = "timestamp"; + /// The `(offset, length)` selections of one cursor scan. + type RangeSelections = Vec<(u32, u32)>; + fn project_batch(batch: &RecordBatch, indices: &[usize]) -> RecordBatch { let fields = indices .iter() @@ -1711,6 +1908,36 @@ mod test { stream.calculate_range(&batch).unwrap() } + #[test] + fn calculate_range_nullable_interior_timestamp_matches_wide_scan() { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let stream = scan_stream_for_test(TimeUnit::Second, 0, 0, 1_000, 1_000, 1_000); + let timestamps = [0, i64::MAX, 1]; + let array = TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(NullBuffer::from(vec![false, false, true])), + ); + let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + + let (ranges, bounds) = stream.calculate_range(&batch).unwrap(); + let wide = RangeManipulateStream::bench_scan_wide( + TimeUnit::Second, + 0, + 1_000, + 1_000, + ×tamps, + bounds.0, + bounds.1, + ); + assert_eq!(bounds, (0, 1_000)); + assert_eq!(ranges, vec![(0, 1), (0, 0)]); + assert_eq!(ranges, wide); + } + #[test] fn calculate_range_keeps_query_aligned_tail() { let (ranges, bounds) = calculate_range_for_test(4, 94, 30, 15, &[20, 50, 80]); @@ -1954,4 +2181,898 @@ mod test { ); } } + + /// Nanoseconds per native tick, as kept by the narrow carrier. + fn native_scale(time_unit: TimeUnit) -> i64 { + i64::try_from(nanoseconds_per_native_tick(time_unit)).unwrap() + } + + /// Shifted nanoseconds of a sample, computed in the exact `i128` carrier. + fn shifted_nanoseconds(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> i128 { + (timestamp as i128) * nanoseconds_per_native_tick(time_unit) + (offset as i128) * 1_000_000 + } + + /// Shifted milliseconds of the instant grid of a sample, flooring like `calculate_range`. + fn shifted_floor_millis( + time_unit: TimeUnit, + offset: Millisecond, + timestamp: i64, + ) -> Millisecond { + Millisecond::try_from( + shifted_nanoseconds(time_unit, offset, timestamp).div_euclid(1_000_000), + ) + .unwrap() + } + + /// Shifted milliseconds kept by the range payload of a sample. + /// + /// The payload conversion truncates toward zero, so samples below the epoch may keep + /// the millisecond above their own native tick. + fn shifted_millis(time_unit: TimeUnit, offset: Millisecond, timestamp: i64) -> Millisecond { + Millisecond::try_from(shifted_nanoseconds(time_unit, offset, timestamp) / 1_000_000) + .unwrap() + } + + /// Stream whose cursor scan the differential tests below drive directly. + fn scan_stream_for_test( + time_unit: TimeUnit, + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); + RangeManipulateStream { + offset, + start, + end, + interval, + range, + time_index: 0, + time_unit, + field_columns: vec![], + aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])), + output_schema: schema, + input: Box::pin(empty_stream), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + } + } + + /// Exact expectations of the cursor scan, formulated outside the production loops. + /// + /// Every evaluation instant keeps the samples of `(instant - range, instant]` in + /// shifted nanoseconds: the lower boundary is exclusive, the upper one inclusive, and + /// native ticks below a millisecond stay distinct. + fn scan_oracle( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> Vec<(u32, u32)> { + let window_ns = (range as i128) * 1_000_000; + let shifted = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + (start..=end) + .step_by(interval as usize) + .map(|instant| { + let instant_ns = (instant as i128) * 1_000_000; + let mut offset = None; + let mut length = 0u32; + for (index, timestamp) in shifted.iter().enumerate() { + if *timestamp > instant_ns - window_ns && *timestamp <= instant_ns { + offset.get_or_insert(index as u32); + length += 1; + } + } + (offset.unwrap_or(0), length) + }) + .collect() + } + + /// Runs the cursor scan of one input through both carriers. + fn scan_both_carriers( + stream: &RangeManipulateStream, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> (Option, RangeSelections) { + let exact = scan_ranges::( + stream.interval, + timestamps, + start, + end, + &ScanNanoseconds::::wide(stream.time_unit, stream.offset, stream.range), + ); + let narrow = narrow_scan( + stream.time_unit, + stream.offset, + stream.range, + timestamps, + start, + end, + ) + .map(|scan| scan_ranges::(stream.interval, timestamps, start, end, &scan)); + // The benchmark hooks must measure the very scans the operator runs. + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + exact, + "the i128 benchmark hook diverged from the cursor scan" + ); + if let Some(narrow) = &narrow { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + Some(narrow.clone()), + "the i64 benchmark hook diverged from the cursor scan" + ); + } else { + assert_eq!( + RangeManipulateStream::bench_scan_narrow( + stream.time_unit, + stream.offset, + stream.range, + stream.interval, + timestamps, + start, + end, + ), + None, + "the i64 benchmark hook must report the fallback" + ); + } + (narrow, exact) + } + + /// Asserts the narrow carrier was taken and matches the exact scan evaluation by evaluation. + fn assert_narrow_carrier_agrees( + narrow: Option>, + exact: &[(u32, u32)], + context: &str, + ) { + let Some(narrow) = narrow else { + panic!("the i64 carrier was expected, but narrow_scan rejected the input: {context}"); + }; + assert_eq!( + narrow.len(), + exact.len(), + "carriers disagree on the number of evaluations: {context}" + ); + for (evaluation, (narrow_range, exact_range)) in narrow.iter().zip(exact).enumerate() { + assert_eq!( + narrow_range, exact_range, + "carriers disagree on evaluation {evaluation}: {context}" + ); + } + } + + /// Replays the cursor scan in checked `i64` arithmetic. + /// + /// Reports whether the narrow carrier overflows on an operation the scan evaluates, + /// which is what makes the `i128` fallback unavoidable for an input. + fn narrow_carrier_overflows( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + interval: Millisecond, + ) -> bool { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return true; + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return true; + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return true; + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let mut left = 0usize; + let mut right = 0usize; + for instant in (start..=end).step_by(interval as usize) { + let Some(instant_ns) = instant.checked_mul(1_000_000) else { + return true; + }; + let Some(window_start) = instant_ns.checked_sub(window_ns) else { + return true; + }; + while left < timestamps.len() { + let Some(shifted) = shift(timestamps[left]) else { + return true; + }; + if shifted <= window_start { + left += 1; + } else { + break; + } + } + right = right.max(left); + while right < timestamps.len() { + let Some(shifted) = shift(timestamps[right]) else { + return true; + }; + if shifted <= instant_ns { + right += 1; + } else { + break; + } + } + } + false + } + + /// The checked `i64` operation that makes [`narrow_scan`] reject an input, if any. + /// + /// Mirrors the rejection criteria of [`narrow_scan`] so the deterministic cases below + /// can pin which boundary drives the `i128` fallback. + fn narrow_carrier_rejection( + time_unit: TimeUnit, + offset: Millisecond, + range: Millisecond, + timestamps: &[i64], + start: Millisecond, + end: Millisecond, + ) -> Option<&'static str> { + let Ok(scale) = i64::try_from(nanoseconds_per_native_tick(time_unit)) else { + return Some("nanoseconds per native tick do not fit i64"); + }; + let Some(offset_ns) = offset.checked_mul(1_000_000) else { + return Some("offset in nanoseconds does not fit i64"); + }; + let Some(window_ns) = range.checked_mul(1_000_000) else { + return Some("range in nanoseconds does not fit i64"); + }; + let shift = |timestamp: i64| timestamp.checked_mul(scale)?.checked_add(offset_ns); + let (Some(first), Some(last)) = (timestamps.first(), timestamps.last()) else { + return Some("empty batch has no endpoints"); + }; + if shift(*first).is_none() || shift(*last).is_none() { + return Some("sample nanoseconds do not fit i64"); + } + if start + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("start bound in nanoseconds does not fit i64"); + } + if end + .checked_mul(1_000_000) + .and_then(|nanoseconds| nanoseconds.checked_sub(window_ns)) + .is_none() + { + return Some("end bound in nanoseconds do not fit i64"); + } + None + } + + #[test] + fn scan_carriers_agree_on_range_boundaries_within_one_millisecond() { + // Membership is `(lower, upper]` on native ticks: the sample exactly on the lower + // boundary is dropped, the sample exactly on the upper boundary is kept, and the + // three samples sharing one millisecond payload stay distinct and survive together. + let cases = [ + ( + TimeUnit::Microsecond, + vec![1_000_000, 1_000_001, 1_000_999, 1_001_000, 1_001_001], + ), + ( + TimeUnit::Nanosecond, + vec![ + 1_000_000_000, + 1_000_000_001, + 1_000_000_999, + 1_001_000_000, + 1_001_000_001, + ], + ), + ]; + + for (time_unit, timestamps) in cases { + let (start, end, interval, range) = (1_001, 1_001, 1, 1); + let context = format!("unit={time_unit:?}, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!( + oracle, + vec![(1, 3)], + "{context}: (lower, upper] expectations" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + + // Two distinct native ticks below the upper boundary share the millisecond + // payload `1_000`; they must not be deduplicated by going through that ABI. + let payloads = timestamps[1..4] + .iter() + .map(|timestamp| shifted_millis(time_unit, 0, *timestamp)) + .collect::>(); + assert_eq!( + payloads, + vec![1_000, 1_000, 1_001], + "{context}: kept samples in one millisecond keep both native ticks" + ); + assert_ne!( + timestamps[1], timestamps[2], + "{context}: the two ticks sharing a payload must stay distinct" + ); + } + } + + #[test] + fn scan_falls_back_to_i128_when_only_the_sample_tail_overflows() { + // `narrow_scan` inspects the sorted endpoints, so a batch whose prefix fits the + // `i64` carrier but whose tail does not has to take the exact carrier. The tail + // sample stays outside every window, so the selections keep the prefix ones. + struct Case { + time_unit: TimeUnit, + narrow_safe: Vec, + overflowing: i64, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + } + + let cases = [ + Case { + time_unit: TimeUnit::Second, + narrow_safe: vec![9_223_372_034, 9_223_372_035, 9_223_372_036], + overflowing: 9_223_372_037, + start: 9_223_372_034_000, + end: 9_223_372_036_000, + interval: 1_000, + range: 1_000, + }, + Case { + time_unit: TimeUnit::Millisecond, + narrow_safe: vec![9_223_372_036_853, 9_223_372_036_854], + overflowing: 9_223_372_036_855, + start: 9_223_372_036_853, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + Case { + time_unit: TimeUnit::Microsecond, + narrow_safe: vec![9_223_372_036_853_999, 9_223_372_036_854_000], + overflowing: 9_223_372_036_854_776, + start: 9_223_372_036_854, + end: 9_223_372_036_854, + interval: 1, + range: 1, + }, + ]; + + for case in cases { + let Case { + time_unit, + narrow_safe, + overflowing, + start, + end, + interval, + range, + } = case; + let context = format!("unit={time_unit:?}, timestamps={narrow_safe:?} + {overflowing}"); + let scale = native_scale(time_unit); + // The prefix fits the narrow carrier, the tail sample does not... + assert!( + narrow_safe + .iter() + .all(|timestamp| timestamp.checked_mul(scale).is_some()), + "{context}: prefix samples must fit i64 nanoseconds" + ); + assert!( + overflowing.checked_mul(scale).is_none(), + "{context}: the tail sample must overflow i64 nanoseconds" + ); + // ...even though the millisecond payload of the tail stays legal. + let payload = shifted_millis(time_unit, 0, overflowing); + assert!( + payload.abs() < i64::MAX / 2, + "{context}: the shifted payload {payload} must stay legal" + ); + + let stream = scan_stream_for_test(time_unit, 0, start, end, interval, range); + + let (narrow, prefix_exact) = scan_both_carriers(&stream, &narrow_safe, start, end); + let prefix_oracle = + scan_oracle(time_unit, 0, range, &narrow_safe, start, end, interval); + assert_eq!( + prefix_exact, prefix_oracle, + "{context}: prefix i128 carrier" + ); + assert_narrow_carrier_agrees(narrow, &prefix_exact, &format!("{context} prefix")); + + let mut timestamps = narrow_safe.clone(); + timestamps.push(overflowing); + assert_eq!( + narrow_carrier_rejection(time_unit, 0, range, ×tamps, start, end), + Some("sample nanoseconds do not fit i64"), + "{context}: the overflowing endpoint must reject the narrow carrier" + ); + assert!( + narrow_carrier_overflows(time_unit, 0, range, ×tamps, start, end, interval), + "{context}: the scan evaluates the overflowing tail" + ); + + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle(time_unit, 0, range, ×tamps, start, end, interval); + assert_eq!(exact, oracle, "{context}: exact carrier on the full batch"); + assert_eq!( + exact, prefix_exact, + "{context}: the out-of-window tail keeps the prefix selections" + ); + } + } + + #[test] + fn scan_carriers_split_at_the_nanosecond_boundaries() { + // Nanoseconds are not scaled, so the `i64` carrier is rejected by the query bounds: + // below the millisecond of `i64::MIN` the start bound underflows, while a batch near + // `i64::MAX` still fits. Both keep a legal millisecond payload. + let low = i64::MIN; + let low_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, low); + assert_eq!(low_floor, -9_223_372_036_855); + assert!(low_floor.checked_mul(1_000_000).is_none()); + assert_eq!( + shifted_millis(TimeUnit::Nanosecond, 0, low), + -9_223_372_036_854 + ); + + let timestamps = vec![low, low + 1_000_000]; + let (start, end, interval, range) = (low_floor, low_floor + 1, 1, 1); + let context = format!("nanosecond floor, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + Some("start bound in nanoseconds does not fit i64"), + "{context}" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert!( + narrow.is_none(), + "{context}: the i128 carrier must be taken" + ); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + // The ceiling side still fits the narrow carrier, so both carriers have to agree. + let high = i64::MAX - 2_000_000; + let high_floor = shifted_floor_millis(TimeUnit::Nanosecond, 0, high); + assert!(high_floor.checked_mul(1_000_000).is_some()); + let timestamps = vec![high, high + 1_000_000, high + 2_000_000]; + let (start, end, interval, range) = (high_floor, high_floor + 2, 1, 1); + let context = format!("nanosecond ceiling, timestamps={timestamps:?}"); + let stream = scan_stream_for_test(TimeUnit::Nanosecond, 0, start, end, interval, range); + assert_eq!( + narrow_carrier_rejection(TimeUnit::Nanosecond, 0, range, ×tamps, start, end), + None, + "{context}: the ceiling batch still fits the narrow carrier" + ); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + let oracle = scan_oracle( + TimeUnit::Nanosecond, + 0, + range, + ×tamps, + start, + end, + interval, + ); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + + #[test] + fn scan_carriers_agree_for_seeded_native_precision_matrix() { + // Seeded matrix over all four native precisions: positive and negative offsets, + // negative times, sub-millisecond samples sharing a millisecond payload, and native + // extremes whose intermediate product overflows while the shifted millisecond stays + // legal. + let units = [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ]; + // `narrow_scan` rejects the two large offsets: `offset * 1ms` overflows `i64`. + let offsets = [0, 1, -1, 250, -250, 9_223_372_036_855, -9_223_372_036_855]; + let extremes = |time_unit: TimeUnit| match time_unit { + TimeUnit::Second => [-9_223_372_037, 9_223_372_037], + TimeUnit::Millisecond => [-9_223_372_036_855, 9_223_372_036_855], + TimeUnit::Microsecond => [-9_223_372_036_854_776, 9_223_372_036_854_776], + TimeUnit::Nanosecond => [i64::MIN, i64::MAX], + }; + let mut prng = TinyPrng(0x5ca1_ab1e_d1ff_0001); + let mut narrow_cases = 0usize; + let mut wide_cases = 0usize; + let mut skipped = 0usize; + + for case in 0..256usize { + let time_unit = units[case % units.len()]; + let offset = offsets[(case / units.len()) % offsets.len()]; + let [low, high] = extremes(time_unit); + let anchor = match prng.next_i64(0, 3) { + 0 => prng.next_i64(-2_000, 2_000), + 1 => -prng.next_i64(1, 2_000), + 2 => low, + _ => high, + }; + // Samples stay within a few milliseconds, with a native granularity small + // enough that sub-millisecond samples share a millisecond payload. + let step = match time_unit { + TimeUnit::Second => 1, + TimeUnit::Millisecond => 1, + TimeUnit::Microsecond => 1 + prng.next_i64(0, 9), + TimeUnit::Nanosecond => 1 + prng.next_i64(0, 999_999), + }; + let count = 1 + prng.next_i64(0, 4); + let ascending = anchor != high; + let mut timestamps = (0..count) + .map(|index| { + let delta = index * step; + if ascending { + anchor + delta + } else { + anchor - delta + } + }) + .collect::>(); + if !ascending { + timestamps.reverse(); + } + // The review shape keeps a legal millisecond payload; skip the extremes that the + // large offsets push out of the millisecond space. + let nanoseconds = timestamps + .iter() + .map(|timestamp| shifted_nanoseconds(time_unit, offset, *timestamp)) + .collect::>(); + if nanoseconds + .iter() + .any(|value| i64::try_from(value / 1_000_000).is_err()) + { + skipped += 1; + continue; + } + + let first_millis = Millisecond::try_from(nanoseconds[0].div_euclid(1_000_000)).unwrap(); + let interval = prng.next_i64(1, 4); + let points = 1 + prng.next_i64(0, 4); + let start = first_millis - interval; + let end = start + interval * (points - 1); + let range = prng.next_i64(0, 3) * interval; + let context = format!( + "case={case}, unit={time_unit:?}, offset={offset}, start={start}, end={end}, \ + interval={interval}, range={range}, timestamps={timestamps:?}" + ); + + let stream = scan_stream_for_test(time_unit, offset, start, end, interval, range); + let oracle = scan_oracle(time_unit, offset, range, ×tamps, start, end, interval); + let (narrow, exact) = scan_both_carriers(&stream, ×tamps, start, end); + assert_eq!(exact, oracle, "{context}: i128 carrier must stay exact"); + + match ( + &narrow, + narrow_carrier_rejection(time_unit, offset, range, ×tamps, start, end), + ) { + (Some(_), None) => { + narrow_cases += 1; + assert!( + !narrow_carrier_overflows( + time_unit, + offset, + range, + ×tamps, + start, + end, + interval + ), + "{context}: every narrow operation must fit i64" + ); + assert_narrow_carrier_agrees(narrow, &exact, &context); + } + (None, Some(_reason)) => { + wide_cases += 1; + } + (narrow, rejection) => panic!( + "{context}: narrow_scan and its mirror disagree, narrow={}, rejection={rejection:?}", + narrow.is_some() + ), + } + } + + assert!( + narrow_cases > 0 && wide_cases > 0, + "both carriers must be covered: {narrow_cases} narrow, {wide_cases} wide \ + ({skipped} skipped)" + ); + assert!( + skipped < 64, + "the generator skipped too many cases: {skipped}" + ); + } + + /// Builds the single-column timestamp batch the cursor scans consume. + fn timestamp_batch(time_unit: TimeUnit, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + )])); + let array: ArrayRef = match time_unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from(timestamps.to_vec())), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps.to_vec())), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps.to_vec())), + }; + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + #[test] + fn calculate_range_dispatches_carriers_and_agrees_pointwise() { + // Production dispatch, not the hooks: `calculate_range` must scan through the narrow + // carrier while the batch fits and through the exact one once the shifted tail does + // not, and both must select the same ranges for the bounds it derives itself. + let narrow_timestamps = [0i64, 15_000, 30_000]; + let wide_timestamps = [9_223_372_036_853, 9_223_372_036_854, 9_223_372_036_855]; + let cases: [(&str, &[i64], Millisecond, Millisecond, bool); 2] = [ + ("narrow carrier", &narrow_timestamps, 0, 30_000, true), + ( + "i128 fallback", + &wide_timestamps, + wide_timestamps[0], + wide_timestamps[1], + false, + ), + ]; + + for (name, timestamps, query_start, query_end, expect_narrow) in cases { + let (interval, range) = if expect_narrow { + (15_000, 15_000) + } else { + (1, 1) + }; + // `scan_stream_for_test` panics on an empty batch when `narrow_scan` is asked for + // the endpoints, so the batch below is never empty. + let stream = scan_stream_for_test( + TimeUnit::Millisecond, + 0, + query_start, + query_end, + interval, + range, + ); + let batch = timestamp_batch(TimeUnit::Millisecond, timestamps); + let (ranges, (start, end)) = stream.calculate_range(&batch).unwrap(); + + let oracle = scan_oracle( + TimeUnit::Millisecond, + 0, + range, + timestamps, + start, + end, + interval, + ); + assert_eq!(ranges, oracle, "{name}: calculate_range pointwise"); + + let rejection = + narrow_carrier_rejection(TimeUnit::Millisecond, 0, range, timestamps, start, end); + assert_eq!( + rejection.is_none(), + expect_narrow, + "{name}: carrier choice by narrow_scan, rejection={rejection:?}" + ); + let (narrow, exact) = scan_both_carriers(&stream, timestamps, start, end); + assert_eq!(exact, ranges, "{name}: i128 carrier is the oracle"); + assert_eq!(narrow.is_some(), expect_narrow, "{name}: narrow carrier"); + assert_eq!( + RangeManipulateStream::bench_scan_wide( + stream.time_unit, + stream.offset, + stream.range, + interval, + timestamps, + start, + end, + ), + ranges, + "{name}: the benchmark hook matches calculate_range" + ); + } + } + + /// Stream whose single millisecond timestamp column is folded into the range payload. + fn payload_stream( + offset: Millisecond, + start: Millisecond, + end: Millisecond, + interval: Millisecond, + range: Millisecond, + ) -> RangeManipulateStream { + let payload_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, payload_type.clone(), false), + Field::new( + RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN), + RangeArray::convert_data_type(payload_type), + false, + ), + ])); + let mut stream = + scan_stream_for_test(TimeUnit::Millisecond, offset, start, end, interval, range); + stream.output_schema = schema; + stream + } + + /// Builds the single-column millisecond timestamp batch the payload tests consume. + fn timestamp_batch_with_timezone(timezone: Option<&str>, timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, timezone.map(Arc::from)), + false, + )])); + let array: ArrayRef = Arc::new( + TimestampMillisecondArray::from(timestamps.to_vec()).with_timezone_opt(timezone), + ); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + /// Millisecond payload values of the range timestamp column of a manipulated batch. + fn payload_millis(batch: &RecordBatch) -> Vec { + let dict = batch + .column(batch.num_columns() - 1) + .as_any() + .downcast_ref::>() + .expect("the timestamp range payload must be a folded dictionary") + .clone(); + let payload = RangeArray::try_new(dict).unwrap(); + assert_eq!( + payload.values().data_type(), + &DataType::Timestamp(TimeUnit::Millisecond, None), + "the payload value type must agree with the output schema" + ); + payload + .values() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[test] + fn manipulate_timezone_time_index_skips_zero_copy_payload() { + // Oracle-verified regression: `Timestamp(Millisecond, Some(tz))` passes a + // `time_unit == Millisecond` check but is not the payload ABI. Cloning such a column + // would stamp a dictionary value type disagreeing with the output schema and + // `RecordBatch::try_new` would fail, so the reuse must be guarded on the data type. + let timestamps = [0i64, 1_000, 2_000, 3_000]; + let stream = payload_stream(0, 0, 3_000, 1_000, 1_000); + + // With offset 0 the data type is the only reason to convert, and the `try_unary` + // path is the identity for a millisecond index. + let zoned = stream + .manipulate(timestamp_batch_with_timezone(Some("+08:00"), ×tamps)) + .expect("a timezone-carrying index must still build the output batch") + .expect("the ranges are not all empty"); + assert_eq!( + payload_millis(&zoned), + timestamps, + "the converted payload must keep the try_unary values" + ); + + // The same logical samples without a timezone do take the zero-copy path; both paths + // must agree on the payload. + let plain = stream + .manipulate(timestamp_batch(TimeUnit::Millisecond, ×tamps)) + .unwrap() + .unwrap(); + assert_eq!( + payload_millis(&plain), + payload_millis(&zoned), + "the zero-copy and the converted path must agree" + ); + } + + /// A NULL slot's underlying value must never trigger the payload overflow + /// error: a nullable batch is converted with arrow's `try_unary`, which applies + /// the closure to valid rows only. The same value as a *valid* sample must still + /// error, through the null-free `unary` + overflow-cell path. + #[test] + fn manipulate_null_slot_underlying_value_does_not_overflow() { + // Second-slot underlying value i64::MAX seconds overflows on conversion, + // but the slot is NULL so the conversion must not see it. offset=1ms + // keeps the batch off the zero-copy path. The query range must be bounded + // and contain the first sample: with `end == 0` no range is selected and + // the conversion (and so the regression) is skipped, while an unbounded + // `end` makes the range scan visit `end - start` evaluation points. + let timestamps = [0i64, i64::MAX]; + let validity = NullBuffer::from(vec![true, false]); + let array: ArrayRef = Arc::new(TimestampSecondArray::new( + timestamps.to_vec().into(), + Some(validity), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let output = stream + .manipulate(batch) + .expect("a NULL slot's underlying value must not overflow the payload") + .expect("the first sample (0s + 1ms) is inside the query range"); + assert_eq!( + payload_millis(&output)[0], + 1, + "the valid sample must be converted to its shifted millisecond payload" + ); + + // Same underlying value, but now valid: the conversion must reject it. + let array: ArrayRef = Arc::new(TimestampSecondArray::from(timestamps.to_vec())); + let schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + let mut stream = payload_stream(1, 0, 1_000, 1, 1_000); + stream.time_unit = TimeUnit::Second; + let error = stream + .manipulate(batch) + .expect_err("a valid overflowing timestamp must still error"); + assert!( + error.to_string().contains("timestamp payload overflow"), + "unexpected error: {error}" + ); + } } diff --git a/tests/perf/query_cases/promql_long_range_rate/case.toml b/tests/perf/query_cases/promql_long_range_rate/case.toml new file mode 100644 index 00000000000..51dfe1f1a86 --- /dev/null +++ b/tests/perf/query_cases/promql_long_range_rate/case.toml @@ -0,0 +1,119 @@ +# PromQL long-window rate benchmark for issue #9070. +# +# 2048 series × 1440 samples at 15-second cadence = 2,949,120 rows spanning the +# six hours from 2024-01-01T00:00:00Z to 2024-01-01T06:00:00Z. Evaluation covers +# the last two hours of that data (04:00:00Z–06:00:00Z at a 15-second step, 481 +# evaluations), so every evaluation holds a full 240-sample one-hour window. +# `rate_1h` is the primary regressed shape: it drives the RangeManipulate +# timestamp payload and the cursor scan across a long, heavily overlapping +# window set. `rate_5m` is the small-window control, `offset 1h` exercises the +# shifted payload path, and the bare selector controls unrelated-path variance. +# +# The generator derives each tag value as `series % cardinality`, so the number +# of distinct (host, instance) series is the least common multiple of the two +# cardinalities. `instance` declares 2048 values so the host(16) × instance tag +# shape still yields 2048 distinct series; a 128-value instance tag would cap +# this fixture at 128 distinct series. +# +# The deterministic wave is not a true monotonic counter, so the rate queries +# measure latency on the regular rate path rather than a numeric counter-rate +# oracle. + +[case] +name = "promql_long_range_rate" +description = "PromQL long-window rate latency regression for #9070 over six hours of 15-second samples" +issue = "https://github.com/GreptimeTeam/greptimedb/issues/9070" + +[scenario] +kind = "direct_readable_sst" +seed = 9070 + +[[scenario.tables]] +database = "public" +name = "promql_long_range_rate" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 2048, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +# Eight 45-minute SSTs (180 timestamps each) tile the six-hour range. +# A (03:00:00Z, 06:00:00Z] scan prunes to four of them; the offset query scans +# the five SSTs covering [02:00:00Z, 05:00:00Z). +[scenario.layout] +regions = 1 +sst_count = 8 +rows_per_sst = 368640 +row_group_size = 36864 +series_count = 2048 +start_unix_nanos = 1704067200000000000 +step_nanos = 15000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +# 240 samples per full window: the primary regressed shape. +[[scenario.queries]] +name = "rate_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Twenty samples per full window: small-window control. +[[scenario.queries]] +name = "rate_5m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[5m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# The same long window shifted back an hour: non-zero-offset payload path. +[[scenario.queries]] +name = "rate_1h_offset_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') rate(promql_long_range_rate{host=~'host.*'}[1h] offset 1h)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 + +# Bare selector over the same evaluation range: unrelated-path variance control. +[[scenario.queries]] +name = "plain_selector_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704081600, 1704088800, '15s') promql_long_range_rate{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 20 From 7605f27608716c5b03273906fe9f1ba912379b62 Mon Sep 17 00:00:00 2001 From: Dragon Roar <2181719471@qq.com> Date: Fri, 9 Oct 2026 11:59:28 +0000 Subject: [PATCH 125/133] fix(metasrv): stop holding the pusher registry lock across mailbox sends (#9445) fix(meta-srv): stop holding the pusher registry lock across mailbox sends Signed-off-by: lczllx <2181719471@qq.com> --- src/meta-srv/src/handler.rs | 193 +++++++++++++++++++++++++++++------- 1 file changed, 156 insertions(+), 37 deletions(-) diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 6bf71674a82..24c1764db2a 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -224,18 +224,50 @@ impl Pusher { } #[inline] - pub async fn push(&self, res: HeartbeatResponse) -> Result<()> { - self.sender.send(Ok(res)).await.map_err(|e| { - error::PushMessageSnafu { - err_msg: e.to_string(), - } - .build() - }) + pub fn header(&self) -> ResponseHeader { + self.res_header.clone() + } +} + +/// An owned snapshot of the [`Pusher`] fields the push paths need, so a caller can drop +/// the [`Pushers`] read guard before awaiting the send. It must not hold a [`Pusher`]: +/// dropping one fires the deregister signal, which reads as a closed channel. +struct PusherHandle { + sender: Sender>, + res_header: ResponseHeader, + deregister_signal_receiver: DeregisterSignalReceiver, +} + +impl PusherHandle { + /// Snapshots `pusher`, which callers read under the [`Pushers`] guard. + fn from_pusher(pusher: &Pusher) -> Self { + Self { + sender: pusher.sender.clone(), + res_header: pusher.header(), + deregister_signal_receiver: pusher.deregister_signal_receiver.clone(), + } + } + + /// Sends one mailbox message, consuming the handle. + async fn push(self, mailbox_message: MailboxMessage) -> Result<()> { + self.sender + .send(Ok(HeartbeatResponse { + header: Some(self.res_header), + mailbox_message: Some(mailbox_message), + ..Default::default() + })) + .await + .map_err(|e| { + error::PushMessageSnafu { + err_msg: e.to_string(), + } + .build() + }) } #[inline] - pub fn header(&self) -> ResponseHeader { - self.res_header.clone() + fn deregister_signal_receiver(&self) -> DeregisterSignalReceiver { + self.deregister_signal_receiver.clone() } } @@ -249,41 +281,41 @@ impl Pushers { pusher_id: PusherId, mailbox_message: MailboxMessage, ) -> Result { - let pushers = self.0.read().await; - let pusher = pushers - .get(&pusher_id) - .with_context(|| error::PusherNotFoundSnafu { - pusher_id: pusher_id.to_string(), - })?; + // Releases the guard before the send below, which parks while the target's + // channel is full. `PusherNotFound` is still reported from inside the guard. + let handle = { + let pushers = self.0.read().await; + let pusher = pushers + .get(&pusher_id) + .with_context(|| error::PusherNotFoundSnafu { + pusher_id: pusher_id.to_string(), + })?; + PusherHandle::from_pusher(pusher) + }; - pusher - .push(HeartbeatResponse { - header: Some(pusher.header()), - mailbox_message: Some(mailbox_message), - ..Default::default() - }) - .await?; + let deregister_signal_receiver = handle.deregister_signal_receiver(); + handle.push(mailbox_message).await?; - Ok(pusher.deregister_signal_receiver.clone()) + Ok(deregister_signal_receiver) } async fn broadcast(&self, role: Role, mailbox_message: &MailboxMessage) -> Result<()> { - let pushers = self.0.read().await; - let pushers = pushers - .range(PusherId::role_range(role)) - .map(|(_, value)| value) - .collect::>(); - let mut results = Vec::with_capacity(pushers.len()); + // Collects owned handles, releasing the guard before the sends below. Collecting + // `&Pusher` would hold it for the whole `join_all`. + let targets = { + let pushers = self.0.read().await; + pushers + .range(PusherId::role_range(role)) + .map(|(_, pusher)| PusherHandle::from_pusher(pusher)) + .collect::>() + }; - for pusher in pushers { + let mut results = Vec::with_capacity(targets.len()); + for handle in targets { let mut mailbox_message = mailbox_message.clone(); mailbox_message.id = 0; // one-way message - results.push(pusher.push(HeartbeatResponse { - header: Some(pusher.header()), - mailbox_message: Some(mailbox_message), - ..Default::default() - })) + results.push(handle.push(mailbox_message)); } // Checks the error out of the loop. @@ -924,14 +956,15 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use api::v1::meta::{HeartbeatRequest, MailboxMessage, RequestHeader, Role}; + use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, MailboxMessage, RequestHeader, Role}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; + use futures::future::poll_immediate; use tokio::sync::mpsc; use super::{ HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroupBuilder, - PusherId, Pushers, + PusherHandle, PusherId, Pushers, }; use crate::error; use crate::handler::collect_stats_handler::CollectStatsHandler; @@ -1397,4 +1430,90 @@ mod tests { let extensions = res.extensions; assert_eq!(extensions.get("custom.key").unwrap(), b"custom-value"); } + + async fn full_pusher( + pushers: &Pushers, + id: PusherId, + ) -> mpsc::Receiver> { + let (tx, rx) = mpsc::channel(1); + tx.try_send(Ok(HeartbeatResponse::default())).unwrap(); + pushers.insert(id, Pusher::new(tx)).await; + // Callers keep `rx` alive: dropping it would make the send fail instead of park. + rx + } + + fn test_mailbox_message(subject: &str) -> MailboxMessage { + MailboxMessage { + id: 0, + subject: subject.to_string(), + timestamp_millis: 123, + ..Default::default() + } + } + + #[tokio::test] + async fn test_pushers_push_does_not_fence_the_registry() { + let pushers = Pushers::default(); + + let full_id = PusherId::new(Role::Datanode, 1); + let _full_rx = full_pusher(&pushers, full_id).await; + + let (healthy_tx, mut healthy_rx) = mpsc::channel(1); + let healthy_id = PusherId::new(Role::Datanode, 2); + pushers.insert(healthy_id, Pusher::new(healthy_tx)).await; + + let msg = test_mailbox_message("fence-test"); + + let mut parked = Box::pin(pushers.push(full_id, msg.clone())); + assert!(poll_immediate(parked.as_mut()).await.is_none()); + + let (writer_tx, _writer_rx) = mpsc::channel(1); + let mut writer = + Box::pin(pushers.insert(PusherId::new(Role::Datanode, 3), Pusher::new(writer_tx))); + assert!(poll_immediate(writer.as_mut()).await.is_some()); + + let mut healthy = Box::pin(pushers.push(healthy_id, msg)); + assert!(poll_immediate(healthy.as_mut()).await.is_some()); + assert!(healthy_rx.try_recv().is_ok()); + } + + #[tokio::test] + async fn test_pushers_broadcast_does_not_fence_the_registry() { + let pushers = Pushers::default(); + + let _full_rx = full_pusher(&pushers, PusherId::new(Role::Datanode, 1)).await; + let (healthy_tx, _healthy_rx) = mpsc::channel(1); + pushers + .insert(PusherId::new(Role::Datanode, 2), Pusher::new(healthy_tx)) + .await; + + let msg = test_mailbox_message("broadcast-fence-test"); + let mut parked = Box::pin(pushers.broadcast(Role::Datanode, &msg)); + assert!(poll_immediate(parked.as_mut()).await.is_none()); + + let (writer_tx, _writer_rx) = mpsc::channel(1); + let mut writer = + Box::pin(pushers.insert(PusherId::new(Role::Flownode, 1), Pusher::new(writer_tx))); + assert!(poll_immediate(writer.as_mut()).await.is_some()); + } + + #[tokio::test] + async fn test_pusher_handle_does_not_fire_the_deregister_signal() { + let (tx, _rx) = mpsc::channel(1); + let pusher = Pusher::new(tx); + let mut signal = pusher.deregister_signal_receiver.clone(); + + let handle = PusherHandle::from_pusher(&pusher); + drop(handle); + + assert!( + tokio::time::timeout(Duration::from_millis(100), signal.changed()) + .await + .is_err() + ); + + // The real `Pusher` still owns the signal. + drop(pusher); + assert!(signal.changed().await.is_ok()); + } } From a026b3b280b4013488220e3d12c96fe3d0d02da1 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 9 Oct 2026 12:21:38 +0000 Subject: [PATCH 126/133] fix(mito): flush memtables containing skip-WAL writes on region close (#9440) * fix(mito): flush unlogged writes when closing regions Signed-off-by: WenyXu * fix(mito): align skip-WAL naming and metric recovery tests Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/metric-engine/src/engine/bulk_insert.rs | 48 ++-- src/metric-engine/src/engine/put.rs | 56 +++-- src/metric-engine/src/test_util.rs | 11 + src/mito2/src/engine/basic_test.rs | 56 ++++- src/mito2/src/engine/scan_test.rs | 58 ++++- src/mito2/src/engine/skip_wal_test.rs | 192 ++++++++++++++- src/mito2/src/flush.rs | 2 +- src/mito2/src/memtable/version.rs | 254 +++++++++++++++++++- src/mito2/src/read/scan_region.rs | 86 +++---- src/mito2/src/region/opener.rs | 27 +-- src/mito2/src/region/version.rs | 10 +- src/mito2/src/region_write_ctx.rs | 196 ++++++++++++++- src/mito2/src/sst/index.rs | 4 +- src/mito2/src/sst/primary_key.rs | 4 +- src/mito2/src/test_util/version_util.rs | 7 +- src/mito2/src/worker/handle_close.rs | 18 +- src/mito2/src/worker/handle_manifest.rs | 11 +- src/mito2/src/worker/handle_write.rs | 27 +-- 18 files changed, 888 insertions(+), 179 deletions(-) diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index fa47ce48008..dd20f34adf8 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -354,17 +354,23 @@ mod tests { #[tokio::test] async fn test_bulk_insert_physical_region_passthrough() { - check_bulk_insert_physical_region_passthrough(false).await; - check_bulk_insert_physical_region_passthrough(true).await; + check_bulk_insert_physical_region_passthrough(false, false).await; + check_bulk_insert_physical_region_passthrough(true, false).await; } - async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool) { + #[tokio::test] + async fn test_bulk_insert_physical_region_close_recovery() { + check_bulk_insert_physical_region_passthrough(false, true).await; + check_bulk_insert_physical_region_passthrough(true, true).await; + } + + async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool, close: bool) { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, ..Default::default() }; - let env = TestEnv::with_mito_config("", mito_config, Default::default()).await; + let mut env = TestEnv::with_mito_config("", mito_config, Default::default()).await; env.init_metric_region().await; env.metric().inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); @@ -414,20 +420,23 @@ mod tests { let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 6); - // Closing without a flush must recover only WAL-backed data. Recreate - // the wrapper too, so its metadata cache cannot hide missing metadata. let stat = env.mito().region_statistic(physical_region_id).unwrap(); assert_eq!(stat.sst_num, 0); - env.metric() - .handle_request( - physical_region_id, - RegionRequest::Close(RegionCloseRequest { - flush_on_close: false, - }), - ) - .await - .unwrap(); - let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + let reopened = if close { + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreating the wrapper prevents its metadata cache from hiding data loss. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Stopping workers without Close exercises recovery from WAL alone. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -446,6 +455,11 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); let stream = reopened .scan_to_stream(logical_region_id, ScanRequest::default()) .await @@ -453,7 +467,7 @@ mod tests { let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!( batches.iter().map(|b| b.num_rows()).sum::(), - if skip_wal { 0 } else { 6 }, + if skip_wal && !close { 0 } else { 6 }, ); } diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index fdf9b02a5ff..b709d99d847 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -934,14 +934,22 @@ mod tests { #[tokio::test] async fn test_put_skip_wal_batch_recovery() { - check_put_skip_wal_batch_recovery("sparse", false).await; - check_put_skip_wal_batch_recovery("sparse", true).await; - check_put_skip_wal_batch_recovery("dense", false).await; - check_put_skip_wal_batch_recovery("dense", true).await; + check_put_skip_wal_batch_recovery("sparse", false, false).await; + check_put_skip_wal_batch_recovery("sparse", true, false).await; + check_put_skip_wal_batch_recovery("dense", false, false).await; + check_put_skip_wal_batch_recovery("dense", true, false).await; } - async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool) { - let env = TestEnv::new().await; + #[tokio::test] + async fn test_put_skip_wal_batch_close_recovery() { + check_put_skip_wal_batch_recovery("sparse", false, true).await; + check_put_skip_wal_batch_recovery("sparse", true, true).await; + check_put_skip_wal_batch_recovery("dense", false, true).await; + check_put_skip_wal_batch_recovery("dense", true, true).await; + } + + async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool, close: bool) { + let mut env = TestEnv::new().await; let engine = env.metric(); engine.inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); @@ -1009,18 +1017,21 @@ mod tests { assert!(stat.memtable_size > 0); assert_eq!(stat.sst_num, 0); } - engine - .handle_request( - physical_region_id, - RegionRequest::Close(RegionCloseRequest { - flush_on_close: false, - }), - ) - .await - .unwrap(); - - // Recreate the wrapper as well, discarding its metadata cache. - let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + let reopened = if close { + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + // Recreate the wrapper as well, discarding its metadata cache. + MetricEngine::try_new(env.mito(), Default::default()).unwrap() + } else { + // Bypass Close to retain coverage of rows lost without WAL protection. + env.reopen_engine(Default::default()).await; + env.metric() + }; reopened.inner.flush_task.stop().await.unwrap(); reopened .handle_request( @@ -1042,12 +1053,17 @@ mod tests { ) .await .unwrap(); + let stat = env + .mito() + .region_statistic(crate::utils::to_data_region_id(physical_region_id)) + .unwrap(); + assert_eq!(stat.sst_num > 0, skip_wal && close); let recovered_metadata = reopened.get_metadata(logical_region_id).await.unwrap(); assert_eq!( metadata_before.column_metadatas, recovered_metadata.column_metadatas ); - let expected = if skip_wal { + let expected = if skip_wal && !close { vec![] } else { vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)] @@ -1055,7 +1071,7 @@ mod tests { assert_eq!( scan_timestamp_values(&reopened, logical_region_id).await, expected, - "encoding={encoding}, skip_wal={skip_wal}" + "encoding={encoding}, skip_wal={skip_wal}, close={close}" ); } diff --git a/src/metric-engine/src/test_util.rs b/src/metric-engine/src/test_util.rs index 0671351ddb3..cf64d968182 100644 --- a/src/metric-engine/src/test_util.rs +++ b/src/metric-engine/src/test_util.rs @@ -121,6 +121,17 @@ impl TestEnv { self.metric.clone() } + /// Restarts Mito without closing regions and recreates the metric metadata cache. + /// Callers stop the old metric engine's metadata flush task before restarting. + pub async fn reopen_engine(&mut self, config: EngineConfig) { + let mito_config = self.mito.mito_config().clone(); + self.mito = self + .mito_env + .reopen_engine(self.mito.clone(), mito_config) + .await; + self.metric = MetricEngine::try_new(self.mito.clone(), config).unwrap(); + } + /// Creates a new follower engine with the same config as the leader engine. pub async fn create_follower_engine(&mut self) -> (MitoEngine, MetricEngine) { let mito = self diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index 776e1f78abe..35f15f2393c 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -1500,7 +1500,10 @@ async fn check_request_skip_wal_recovery(flat_format: bool, skip_wal: bool) { let before = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(before.iter().map(|b| b.num_rows()).sum::(), 4); - reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + // Stop the engine directly, without region close, to exercise WAL-only recovery. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -1625,7 +1628,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( skip_wal: bool, flush_before_reopen: bool, ) { - use crate::region_write_ctx::RegionWriteCtx; + use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::test_util::LogStoreImpl; use crate::wal::Wal; @@ -1657,6 +1660,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ®ion.version_control, region.provider.clone(), None, + WriteSource::Request, ); let mut receivers = Vec::with_capacity(4); for (index, (skip, timestamp)) in [(false, 0), (skip_wal, 0), (false, 0), (skip_wal, 1)] @@ -1692,8 +1696,7 @@ async fn check_request_skip_wal_mixed_batch_recovery( ctx.add_wal_entry(&mut writer).unwrap(); let response = writer.write_to_wal().await.unwrap(); assert_eq!(response.last_entry_ids.get(®ion_id), Some(&1)); - ctx.write_memtable().await; - ctx.publish_sequence_and_entry_id(); + ctx.write_memtables().await; drop(ctx); for rx in receivers { assert_eq!(rx.await.unwrap().unwrap(), 2); @@ -1751,9 +1754,13 @@ async fn check_request_skip_wal_mixed_batch_recovery( ); } - // A no-flush close discards memtables. Only WAL-backed rows recover - // unless an explicit flush has already persisted all requests. - reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + // Bypass region close so only WAL-backed rows recover unless the explicit + // flush above persisted all requests. + let engine = + restart_without_region_close(&mut env, engine, flat_format, region_id, table_dir).await; + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); let loses_skipped_rows = skip_wal && !flush_before_reopen; let mut expected_values = if loses_skipped_rows { vec![(0, 20.0), (1000, 1.0), (3000, 21.0)] @@ -1830,3 +1837,38 @@ async fn request_skip_wal_values(engine: &MitoEngine, region_id: RegionId) -> Ve values.sort_unstable_by_key(|(timestamp, _)| *timestamp); values } + +/// Restarts workers without closing regions, preserving WAL-only recovery coverage. +async fn restart_without_region_close( + env: &mut TestEnv, + engine: MitoEngine, + flat_format: bool, + region_id: RegionId, + table_dir: String, +) -> MitoEngine { + let engine = env + .reopen_engine( + engine, + MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + options: HashMap::new(), + skip_wal_replay: false, + path_type: PathType::Bare, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine +} diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 39a1e66f6d8..f86b327735d 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -56,7 +56,7 @@ use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; use store_api::region_request::{ AlterKind, RegionAlterRequest, RegionBulkInsertsRequest, RegionCompactRequest, - RegionPutRequest, RegionRequest, SetRegionOption, + RegionOpenRequest, RegionPutRequest, RegionRequest, SetRegionOption, }; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ @@ -2991,13 +2991,21 @@ fn build_bulk_insert_request( #[tokio::test] async fn test_bulk_skip_wal_recovery() { - check_bulk_skip_wal_recovery(false, false).await; - check_bulk_skip_wal_recovery(false, true).await; - check_bulk_skip_wal_recovery(true, false).await; - check_bulk_skip_wal_recovery(true, true).await; + check_bulk_skip_wal_recovery(false, false, false).await; + check_bulk_skip_wal_recovery(false, true, false).await; + check_bulk_skip_wal_recovery(true, false, false).await; + check_bulk_skip_wal_recovery(true, true, false).await; } -async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { +#[tokio::test] +async fn test_bulk_skip_wal_close_recovery() { + check_bulk_skip_wal_recovery(false, false, true).await; + check_bulk_skip_wal_recovery(false, true, true).await; + check_bulk_skip_wal_recovery(true, false, true).await; + check_bulk_skip_wal_recovery(true, true, true).await; +} + +async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool, close: bool) { let mut env = TestEnv::new().await; let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -3048,7 +3056,41 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { u64::from(!skip_wal) ); } - reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let engine = if close { + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + assert_eq!( + region.version().flushed_sequence, + if flush || skip_wal { 4 } else { 0 } + ); + engine + } else { + // Stopping the engine directly exercises recovery without close-time flush. + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: store_api::region_request::PathType::Bare, + options: HashMap::new(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + engine + }; + assert!( + !engine + .get_region(region_id) + .unwrap() + .version() + .memtables + .has_skip_wal_writes() + ); let stream = engine .scan_to_stream(region_id, ScanRequest::default()) .await @@ -3060,7 +3102,7 @@ async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { .iter() .map(|b| b.num_rows()) .sum::(), - if skip_wal && !flush { 0 } else { 4 } + if skip_wal && !flush && !close { 0 } else { 4 } ); } diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index 363318d1cf5..915099e7228 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -34,7 +35,7 @@ use crate::engine::listener::AlterFlushListener; use crate::test_util::{ CreateRequestBuilder, LogStoreFactory, TestEnv, build_rows, flush_region, kafka_log_store_factory, multiple_log_store_factories, prepare_test_for_kafka_log_store, - put_rows, raft_engine_log_store_factory, rows_schema, + put_rows, raft_engine_log_store_factory, reopen_region, rows_schema, }; fn set_skip_wal_request(skip_wal: bool) -> RegionRequest { @@ -920,3 +921,192 @@ async fn test_close_region_after_truncate_skip_wal() { let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(3, total_rows); } + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal( + #[case] skip_wal: bool, + #[values(false, true)] flat_format: bool, +) { + let mut env = TestEnv::with_prefix("close-request-skip-wal").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + assert!(!region.skip_wal()); + assert_eq!(region.version().memtables.has_skip_wal_writes(), skip_wal); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + // Logged writes are recovered without a close-time flush. + assert_eq!( + region.version().flushed_sequence, + if skip_wal { 3 } else { 0 } + ); + let reopened = engine.get_region(region_id).unwrap(); + assert!(!reopened.version().memtables.has_skip_wal_writes()); + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[case(false)] +#[case(true)] +#[tokio::test] +async fn test_close_request_skip_wal_during_flush(#[case] new_skip_wal: bool) { + let mut env = TestEnv::with_prefix("close-request-skip-wal-during-flush").await; + let listener = Arc::new(AlterFlushListener::default()); + let engine = env + .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + engine + .handle_request(region_id, RegionRequest::Create(request.clone())) + .await + .unwrap(); + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: true, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let engine_cloned = engine.clone(); + let flush_job = tokio::spawn(async move { + flush_region(&engine_cloned, region_id, None).await; + }); + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + + // These rows belong to the next mutable, while the skipped rows are frozen. + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: new_skip_wal, + rows: Rows { + schema: rows_schema(&request), + rows: build_rows(3, 6), + }, + hint: None, + partition_expr_version: None, + }), + ) + .await + .unwrap(); + let expected = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + let request_count = listener.request_count(); + let engine_cloned = engine.clone(); + let close_job = tokio::spawn(async move { + engine_cloned + .handle_request( + region_id, + RegionRequest::Close(RegionCloseRequest::default()), + ) + .await + .unwrap(); + }); + tokio::time::timeout( + Duration::from_secs(5), + listener.wait_request_count(request_count + 1), + ) + .await + .unwrap(); + listener.wake_flush(); + tokio::time::timeout(Duration::from_secs(5), flush_job) + .await + .unwrap() + .unwrap(); + if new_skip_wal { + tokio::time::timeout(Duration::from_secs(5), listener.wait_flush_begin()) + .await + .unwrap(); + listener.wake_flush(); + } + tokio::time::timeout(Duration::from_secs(5), close_job) + .await + .unwrap() + .unwrap(); + assert_eq!( + region.version().flushed_sequence, + if new_skip_wal { 6 } else { 3 } + ); + assert!(!region.version().memtables.has_skip_wal_writes()); + + reopen_region(&engine, region_id, request.table_dir, false, HashMap::new()).await; + let actual = common_recordbatch::RecordBatches::try_collect( + engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(), + ) + .await + .unwrap() + .pretty_print() + .unwrap(); + assert_eq!(actual, expected); +} diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index b5d8c4e56ea..e910c81f896 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -232,7 +232,7 @@ pub enum FlushReason { Repartition, /// Flush triggered by remote WAL pruning. RemoteWalPrune, - /// Flush before closing a Noop WAL region. + /// Flush before closing a region. Closing, } diff --git a/src/mito2/src/memtable/version.rs b/src/mito2/src/memtable/version.rs index d707128c93a..cd5dae10315 100644 --- a/src/mito2/src/memtable/version.rs +++ b/src/mito2/src/memtable/version.rs @@ -15,6 +15,7 @@ //! Memtable version. use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use common_time::Timestamp; @@ -23,8 +24,9 @@ use store_api::metadata::RegionMetadataRef; use store_api::storage::SequenceNumber; use crate::error::Result; -use crate::memtable::time_partition::TimePartitionsRef; -use crate::memtable::{MemtableId, MemtableRef}; +use crate::memtable::bulk::part::BulkPart; +use crate::memtable::time_partition::TimePartitions; +use crate::memtable::{KeyValues, MemtableBuilderRef, MemtableId, MemtableRef}; pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; @@ -32,26 +34,35 @@ pub(crate) type SmallMemtableVec = SmallVec<[MemtableRef; 2]>; #[derive(Debug, Clone)] pub(crate) struct MemtableVersion { /// Mutable memtable. - pub(crate) mutable: TimePartitionsRef, + pub(crate) mutable: MutableMemtablesRef, /// Immutable memtables. /// /// We only allow one flush job per region but if a flush job failed, then we /// might need to store more than one immutable memtable on the next time we /// flush the region. immutables: SmallMemtableVec, + /// Whether immutable memtables may contain skip-WAL writes. + /// Flushes are serialized per region and cover all immutable memtables. + immutable_has_skip_wal_writes: bool, } pub(crate) type MemtableVersionRef = Arc; impl MemtableVersion { /// Returns a new [MemtableVersion] with specific mutable memtable. - pub(crate) fn new(mutable: TimePartitionsRef) -> MemtableVersion { + pub(crate) fn new(mutable: TimePartitions) -> MemtableVersion { MemtableVersion { - mutable, + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: SmallVec::new(), + immutable_has_skip_wal_writes: false, } } + /// Returns whether current memtables may contain skip-WAL writes. + pub(crate) fn has_skip_wal_writes(&self) -> bool { + self.mutable.has_skip_wal_writes() || self.immutable_has_skip_wal_writes + } + /// Immutable memtables. pub(crate) fn immutables(&self) -> &[MemtableRef] { &self.immutables @@ -94,21 +105,25 @@ impl MemtableVersion { } // Update the time window. - let mutable = self.mutable.new_with_part_duration(time_window, None); + let mutable = self + .mutable + .partitions + .new_with_part_duration(time_window, None); common_telemetry::debug!( "Freeze empty memtable, update partition duration from {:?} to {:?}", self.mutable.part_duration(), time_window ); return Ok(Some(MemtableVersion { - mutable: Arc::new(mutable), + mutable: Arc::new(MutableMemtables::new(mutable)), immutables: self.immutables.clone(), + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes, })); } // Marks the mutable memtable as immutable so it can free the memory usage from our // soft limit. - self.mutable.freeze()?; + self.mutable.partitions.freeze()?; // Fork the memtable. if Some(self.mutable.part_duration()) != time_window { common_telemetry::debug!( @@ -117,17 +132,23 @@ impl MemtableVersion { time_window ); } - let mutable = Arc::new(self.mutable.fork(metadata, time_window)); + let mutable = Arc::new(MutableMemtables::new( + self.mutable.partitions.fork(metadata, time_window), + )); let mut immutables = SmallVec::with_capacity(self.immutables.len() + self.mutable.num_partitions()); immutables.extend(self.immutables.iter().cloned()); // Pushes the mutable memtable to immutable list. - self.mutable.list_memtables_to_small_vec(&mut immutables); + self.mutable + .partitions + .list_memtables_to_small_vec(&mut immutables); Ok(Some(MemtableVersion { mutable, immutables, + immutable_has_skip_wal_writes: self.immutable_has_skip_wal_writes + || self.mutable.has_skip_wal_writes(), })) } @@ -139,6 +160,9 @@ impl MemtableVersion { .filter(|mem| !ids.contains(&mem.id())) .cloned() .collect(); + if self.immutables.is_empty() { + self.immutable_has_skip_wal_writes = false; + } } /// Returns the memory usage of the mutable memtable. @@ -182,3 +206,213 @@ impl MemtableVersion { self.mutable.is_empty() && self.immutables.is_empty() } } + +/// Mutable time partitions and their shared write state. +/// +/// Version clones share this object. Freezing creates a new object so writes to +/// the new mutable memtables cannot change the old version's state. +#[derive(Debug)] +pub(crate) struct MutableMemtables { + partitions: TimePartitions, + /// Set before installation because failed writes may install some rows. + has_skip_wal_writes: AtomicBool, +} + +pub(crate) type MutableMemtablesRef = Arc; + +impl MutableMemtables { + fn new(partitions: TimePartitions) -> Self { + Self { + partitions, + has_skip_wal_writes: AtomicBool::new(false), + } + } + + /// Marks that these memtables may contain skip-WAL writes. + pub(crate) fn mark_skip_wal_writes(&self) { + self.has_skip_wal_writes.store(true, Ordering::Relaxed); + } + + fn has_skip_wal_writes(&self) -> bool { + self.has_skip_wal_writes.load(Ordering::Relaxed) + } + + /// Writes key values to the mutable time partitions. + pub(crate) fn write(&self, kvs: &KeyValues) -> Result<()> { + self.partitions.write(kvs) + } + + /// Writes a bulk part to the mutable time partitions. + pub(crate) fn write_bulk(&self, part: BulkPart) -> Result<()> { + self.partitions.write_bulk(part) + } + + /// Returns whether all mutable time partitions are empty. + pub(crate) fn is_empty(&self) -> bool { + self.partitions.is_empty() + } + + /// Returns the memtable builder. + pub(crate) fn memtable_builder(&self) -> &MemtableBuilderRef { + self.partitions.memtable_builder() + } + + /// Returns the next memtable ID. + pub(crate) fn next_memtable_id(&self) -> MemtableId { + self.partitions.next_memtable_id() + } + + /// Returns the time partition duration. + pub(crate) fn part_duration(&self) -> Duration { + self.partitions.part_duration() + } + + /// Creates empty time partitions with an optional new duration and builder. + /// The returned partitions do not share data or write state with this object. + pub(crate) fn new_with_part_duration( + &self, + part_duration: Option, + memtable_builder: Option, + ) -> TimePartitions { + self.partitions + .new_with_part_duration(part_duration, memtable_builder) + } + + /// Returns the mutable memory usage. + pub(crate) fn memory_usage(&self) -> usize { + self.partitions.memory_usage() + } + + /// Returns the number of time series in mutable time partitions. + pub(crate) fn series_count(&self) -> usize { + self.partitions.series_count() + } + + fn num_rows(&self) -> u64 { + self.partitions.num_rows() + } + + fn num_partitions(&self) -> usize { + self.partitions.num_partitions() + } + + fn list_memtables(&self, memtables: &mut Vec) { + self.partitions.list_memtables(memtables); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memtable::time_series::TimeSeriesMemtableBuilder; + use crate::test_util::memtable_util; + + #[test] + fn test_skip_wal_writes_follow_memtable_lifecycle() { + for new_skip_wal_writes in [false, true] { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let snapshot = version.clone(); + assert!(!version.has_skip_wal_writes()); + + let kvs = memtable_util::build_key_values( + &metadata, + "hello".to_string(), + 0, + &[1000, 7000], + 1, + ); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + assert!(snapshot.has_skip_wal_writes()); + + let mut frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + assert_eq!(2, frozen.immutables().len()); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + assert!(!Arc::ptr_eq(&version.mutable, &frozen.mutable)); + + let during_flush = frozen.clone(); + if new_skip_wal_writes { + frozen.mutable.mark_skip_wal_writes(); + } + frozen.mutable.write(&kvs).unwrap(); + let ids: Vec<_> = frozen.immutables().iter().map(|mem| mem.id()).collect(); + frozen.remove_memtables(&ids); + assert_eq!(new_skip_wal_writes, frozen.has_skip_wal_writes()); + assert_eq!( + new_skip_wal_writes, + during_flush.mutable.has_skip_wal_writes() + ); + // Removing flushed memtables does not mutate the old flush snapshot. + assert_eq!(2, during_flush.immutables().len()); + assert!(during_flush.has_skip_wal_writes()); + + let mut next_flush = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + let ids: Vec<_> = next_flush.immutables().iter().map(|mem| mem.id()).collect(); + next_flush.remove_memtables(&ids); + assert!(!next_flush.has_skip_wal_writes()); + } + } + + #[test] + fn test_skip_wal_writes_survive_flush_retry() { + let metadata = memtable_util::metadata_for_test(); + let partitions = TimePartitions::new( + metadata.clone(), + Arc::new(TimeSeriesMemtableBuilder::default()), + 0, + Some(Duration::from_secs(5)), + ); + let version = MemtableVersion::new(partitions); + let kvs = memtable_util::build_key_values(&metadata, "hello".to_string(), 0, &[1000], 1); + version.mutable.mark_skip_wal_writes(); + version.mutable.write(&kvs).unwrap(); + let frozen = version + .freeze_mutable(&metadata, Some(Duration::from_secs(5))) + .unwrap() + .unwrap(); + + // A failed flush leaves immutable memtables installed. Changing the + // empty mutable's partition duration must retain their write state. + let frozen = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert!(frozen.has_skip_wal_writes()); + assert!(!frozen.mutable.has_skip_wal_writes()); + + // A retry includes the old immutable memtables and new logged writes. + frozen.mutable.write(&kvs).unwrap(); + let mut retried = frozen + .freeze_mutable(&metadata, Some(Duration::from_secs(10))) + .unwrap() + .unwrap(); + assert_eq!(2, retried.immutables().len()); + assert!(retried.has_skip_wal_writes()); + let ids: Vec<_> = retried.immutables().iter().map(|mem| mem.id()).collect(); + retried.remove_memtables(&ids); + assert!(!retried.has_skip_wal_writes()); + + let replacement = MemtableVersion::new(TimePartitions::new( + metadata, + frozen.mutable.memtable_builder().clone(), + frozen.mutable.next_memtable_id(), + None, + )); + assert!(!replacement.has_skip_wal_writes()); + assert!(frozen.has_skip_wal_writes()); + } +} diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 5cfd6a29844..0f9cea941b3 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -2888,12 +2888,12 @@ mod tests { }) .primary_key(vec![0]); let metadata = Arc::new(builder.build().unwrap()); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), ); @@ -3167,12 +3167,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let version = Arc::new( crate::region::version::VersionBuilder::new(metadata, mutable) .options(crate::region::options::RegionOptions { @@ -3207,12 +3207,12 @@ mod tests { ..Default::default() }; let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); - let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + let mutable = crate::memtable::time_partition::TimePartitions::new( metadata.clone(), Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), 0, None, - )); + ); let target_region_id = metadata.region_id; // Files at or below the cursor are excluded. Among the remaining files, @@ -3224,23 +3224,25 @@ mod tests { (None, false, true, false), (None, true, true, true), ] { - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: target_region_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: target_region_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let (files, range) = exact_sequence_range(&request, &version).unwrap(); assert_eq!(expected_files, !files.is_empty()); assert_eq!(expected_range, range.is_some()); @@ -3255,24 +3257,26 @@ mod tests { (None, true, 1, true), ] { let file_id = store_api::storage::FileId::random(); - let version = - crate::region::version::VersionBuilder::new(metadata.clone(), mutable.clone()) - .options(crate::region::options::RegionOptions { - preserve_row_sequence: true, - ..Default::default() - }) - .add_files( - new_noop_file_purger(), - [FileMeta { - region_id: RegionId::new(1, 2), - file_id, - sequence, - preserve_row_sequence: marked, - ..Default::default() - }] - .into_iter(), - ) - .build(); + let version = crate::region::version::VersionBuilder::new( + metadata.clone(), + mutable.new_with_part_duration(None, None), + ) + .options(crate::region::options::RegionOptions { + preserve_row_sequence: true, + ..Default::default() + }) + .add_files( + new_noop_file_purger(), + [FileMeta { + region_id: RegionId::new(1, 2), + file_id, + sequence, + preserve_row_sequence: marked, + ..Default::default() + }] + .into_iter(), + ) + .build(); let result = exact_sequence_range(&request, &version); if expected_error { assert!(result.is_err()); diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 84a0e13833f..71c1192cbf2 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -58,7 +58,7 @@ use crate::error::{ use crate::manifest::action::RegionManifest; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::memtable::bulk::part::BulkPart; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::{MemtableBuilderProvider, ensure_json2_not_use_time_series_memtable}; use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES}; use crate::read::series_candidate::is_sparse_metric_metadata; @@ -68,7 +68,7 @@ use crate::region::{ ManifestContext, ManifestStats, MitoRegion, MitoRegionRef, RegionLeaderState, RegionRoleState, RegionStats, }; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::OptionOutputTx; use crate::schedule::scheduler::SchedulerRef; use crate::series_index::{IndexFilePurger, load_version_control}; @@ -410,12 +410,8 @@ impl RegionOpener { let memtable_builder = self.memtable_builder_provider.builder_for_options(&options); let part_duration = options.compaction.time_window(); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); debug!( "Create region {} with options: {:?}, default_flat_format: {}", @@ -587,12 +583,8 @@ impl RegionOpener { .time_window() .or(manifest.compaction_time_window); // Initial memtable id is 0. - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - memtable_builder.clone(), - 0, - part_duration, - )); + let mutable = + TimePartitions::new(metadata.clone(), memtable_builder.clone(), 0, part_duration); // Updates region options by manifest before creating version. let version_builder = version_builder_from_manifest( @@ -785,7 +777,7 @@ pub(crate) fn version_builder_from_manifest( manifest: &RegionManifest, metadata: RegionMetadataRef, file_purger: FilePurgerRef, - mutable: TimePartitionsRef, + mutable: TimePartitions, region_options: RegionOptions, ) -> VersionBuilder { VersionBuilder::new(metadata, mutable) @@ -961,6 +953,7 @@ where provider.clone(), // For WAL replay, we don't need to track the write bytes rate. None, + WriteSource::WalReplay, ); for mutation in entry.mutations { rows_replayed += mutation @@ -1009,11 +1002,9 @@ where // set next_entry_id and write to memtable. region_write_ctx.set_next_entry_id(last_entry_id + 1); - region_write_ctx.write_memtable().await; - region_write_ctx.write_bulk().await; // Publish the replayed sequences only after all rows (including bulk // parts) are installed, matching the write path ordering. - region_write_ctx.publish_sequence_and_entry_id(); + region_write_ctx.write_memtables().await; } // TODO(weny): We need to update `flushed_entry_id` in the region manifest diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index a56cef1f544..1c595fb0b9f 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -32,7 +32,7 @@ use store_api::storage::{RegionId, SequenceNumber}; use crate::error::Result; use crate::manifest::action::{RegionEdit, TruncateKind}; -use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef}; +use crate::memtable::time_partition::TimePartitions; use crate::memtable::version::{MemtableVersion, MemtableVersionRef}; use crate::memtable::{MemtableBuilderRef, MemtableId}; use crate::region::options::RegionOptions; @@ -319,13 +319,13 @@ impl VersionControl { version: &Version, metadata: RegionMetadataRef, memtable_builder: MemtableBuilderRef, - ) -> TimePartitionsRef { - Arc::new(TimePartitions::new( + ) -> TimePartitions { + TimePartitions::new( metadata, memtable_builder, version.memtables.mutable.next_memtable_id(), Some(version.memtables.mutable.part_duration()), - )) + ) } } @@ -402,7 +402,7 @@ pub(crate) struct VersionBuilder { impl VersionBuilder { /// Returns a new builder. - pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitionsRef) -> Self { + pub(crate) fn new(metadata: RegionMetadataRef, mutable: TimePartitions) -> Self { VersionBuilder { metadata: metadata.clone(), memtables: Arc::new(MemtableVersion::new(mutable)), diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index a64c37dac7f..f8d7eb9c840 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -70,8 +70,19 @@ impl Drop for WriteNotify { } } +/// Source of writes installed into memtables. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WriteSource { + Request, + WalReplay, +} + /// Context to keep region metadata and buffer write requests. pub(crate) struct RegionWriteCtx { + /// Source of the buffered writes. + source: WriteSource, + /// Whether accepted non-empty writes bypass WAL. WAL replay never sets this flag. + has_skip_wal_writes: bool, /// Id of region to write. region_id: RegionId, /// Version of the region while creating the context. @@ -120,6 +131,7 @@ impl RegionWriteCtx { version_control: &VersionControlRef, provider: Provider, written_bytes: Option>, + source: WriteSource, ) -> RegionWriteCtx { let VersionControlData { version, @@ -129,6 +141,8 @@ impl RegionWriteCtx { } = version_control.current(); RegionWriteCtx { + source, + has_skip_wal_writes: false, region_id, version, version_control: version_control.clone(), @@ -162,6 +176,7 @@ impl RegionWriteCtx { self.next_sequence = sequence; } let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); + self.has_skip_wal_writes |= num_rows > 0 && self.is_skip_wal_write(skip_wal); let mutation = Mutation { op_type, sequence: self.next_sequence, @@ -255,8 +270,29 @@ impl RegionWriteCtx { self.next_entry_id } - /// Consumes mutations and writes them into mutable memtable. - pub(crate) async fn write_memtable(&mut self) { + /// Installs ordinary and bulk writes, then publishes their sequence and entry ID. + pub(crate) async fn write_memtables(&mut self) { + if self.failed { + return; + } + // A failed installation may leave some rows in memory. + if self.has_skip_wal_writes { + self.version.memtables.mutable.mark_skip_wal_writes(); + } + self.write_mutations().await; + self.write_bulk().await; + self.publish_sequence_and_entry_id(); + } + + fn is_skip_wal_write(&self, request_skip_wal: bool) -> bool { + self.source == WriteSource::Request + && (request_skip_wal + || self.provider == Provider::Noop + || self.version.options.skip_wal) + } + + /// Consumes mutations and writes them into mutable memtables. + async fn write_mutations(&mut self) { debug_assert_eq!(self.wal_notifiers.len(), self.wal_entry.mutations.len()); if self.failed { @@ -336,6 +372,7 @@ impl RegionWriteCtx { self.wal_entry.bulk_entries.push(entry); } + self.has_skip_wal_writes |= bulk.num_rows() > 0 && self.is_skip_wal_write(skip_wal); self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); @@ -344,7 +381,7 @@ impl RegionWriteCtx { true } - pub(crate) async fn write_bulk(&mut self) { + async fn write_bulk(&mut self) { if self.failed || self.bulk_parts.is_empty() { return; } @@ -398,10 +435,10 @@ impl RegionWriteCtx { } /// Publishes the assigned sequences and entry id to the region's committed - /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// watermark. Only call after both [`write_mutations`](Self::write_mutations) /// and [`write_bulk`](Self::write_bulk) have completed; a failed context /// must not publish. - pub(crate) fn publish_sequence_and_entry_id(&self) { + fn publish_sequence_and_entry_id(&self) { if self.failed { return; } @@ -576,6 +613,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); for (op_type, skip, num_rows) in [ (OpType::Put, skip_wal, 2), @@ -657,6 +695,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); ctx.push_mutation( OpType::Delete as i32, @@ -693,6 +732,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -736,6 +776,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); // Alternate policies in one context, retaining all parts for the memtable. for skip in [skip_wal, false, skip_wal] { @@ -772,8 +813,13 @@ mod tests { let builder = VersionControlBuilder::new(); let region_id = builder.region_id(); let version_control = Arc::new(builder.build()); - let mut ctx = - RegionWriteCtx::new(region_id, &version_control, Provider::noop_provider(), None); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::noop_provider(), + None, + WriteSource::Request, + ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None, skip_wal)); @@ -791,6 +837,142 @@ mod tests { assert!(result.is_err(), "bulk notifier should report WAL error"); } + #[tokio::test] + async fn test_skip_wal_writes_source_and_policy() { + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Value}; + + for source in [WriteSource::Request, WriteSource::WalReplay] { + for skip_wal in [false, true] { + for bulk in [false, true] { + for fail_before_install in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + source, + ); + let (tx, rx) = oneshot::channel(); + if bulk { + let mut part = new_bulk_part(); + part.fill_missing_columns(&ctx.version.metadata).unwrap(); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), part, None, skip_wal)); + } else { + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![ + time_index_column_schema( + "ts", + ColumnDataType::TimestampMillisecond, + ), + tag_column_schema("tag_0", ColumnDataType::String), + ], + rows: vec![Row { + values: vec![ + crate::test_util::ts_ms_value(0), + Value { + value_data: Some(ValueData::StringValue( + "a".to_string(), + )), + }, + ], + }], + }), + None, + OptionOutputTx::from(tx), + None, + skip_wal, + ); + } + if fail_before_install { + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "WAL append failed".to_string(), + } + .build(), + )); + } + ctx.write_memtables().await; + assert_eq!( + version_control + .current() + .version + .memtables + .has_skip_wal_writes(), + source == WriteSource::Request && skip_wal && !fail_before_install, + "source={source:?}, skip_wal={skip_wal}, bulk={bulk}, failed={fail_before_install}", + ); + assert_eq!( + version_control.committed_sequence(), + if fail_before_install { + 0 + } else if bulk { + 2 + } else { + 1 + } + ); + drop(ctx); + assert_eq!(rx.await.unwrap().is_err(), fail_before_install); + } + } + } + } + } + + #[test] + fn test_skip_wal_writes_region_policy_and_empty_requests() { + for source in [WriteSource::Request, WriteSource::WalReplay] { + for noop in [false, true] { + for region_skip_wal in [false, true] { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let mut version = builder.build_version(); + version.options.skip_wal = region_skip_wal; + let version_control = + Arc::new(crate::region::version::VersionControl::new(version)); + let provider = if noop { + Provider::Noop + } else { + Provider::raft_engine_provider(region_id.as_u64()) + }; + let mut ctx = + RegionWriteCtx::new(region_id, &version_control, provider, None, source); + ctx.push_mutation( + OpType::Put as i32, + None, + None, + OptionOutputTx::none(), + None, + true, + ); + assert!(!ctx.has_skip_wal_writes); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default()], + }), + None, + OptionOutputTx::none(), + None, + false, + ); + assert_eq!( + ctx.has_skip_wal_writes, + source == WriteSource::Request && (noop || region_skip_wal) + ); + } + } + } + } + fn new_bulk_part() -> BulkPart { let schema = Arc::new(Schema::new(vec![Field::new( "ts", diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 502f6bcdc05..be52170f713 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -1611,12 +1611,12 @@ mod tests { file_purger: FilePurgerRef, files: HashMap, ) -> VersionControlRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(EmptyMemtableBuilder::default()), 0, None, - )); + ); let version_builder = VersionBuilder::new(metadata, mutable) .add_files(file_purger, files.values().cloned()) .build(); diff --git a/src/mito2/src/sst/primary_key.rs b/src/mito2/src/sst/primary_key.rs index 5fb258052c7..7752a652f80 100644 --- a/src/mito2/src/sst/primary_key.rs +++ b/src/mito2/src/sst/primary_key.rs @@ -270,12 +270,12 @@ mod tests { } fn version_with_files(metadata: RegionMetadataRef, files: Vec) -> VersionRef { - let mutable = Arc::new(TimePartitions::new( + let mutable = TimePartitions::new( metadata.clone(), Arc::new(TimeSeriesMemtableBuilder::default()), 0, None, - )); + ); Arc::new( VersionBuilder::new(metadata, mutable) .add_files(new_noop_file_purger(), files.into_iter()) diff --git a/src/mito2/src/test_util/version_util.rs b/src/mito2/src/test_util/version_util.rs index 1c0d7bfebc8..44827b35cc1 100644 --- a/src/mito2/src/test_util/version_util.rs +++ b/src/mito2/src/test_util/version_util.rs @@ -132,12 +132,7 @@ impl VersionControlBuilder { pub(crate) fn build_version(&self) -> Version { let metadata = Arc::new(self.metadata.clone()); - let mutable = Arc::new(TimePartitions::new( - metadata.clone(), - self.memtable_builder.clone(), - 0, - None, - )); + let mutable = TimePartitions::new(metadata.clone(), self.memtable_builder.clone(), 0, None); VersionBuilder::new(metadata, mutable) .add_files(self.file_purger.clone(), self.files.values().cloned()) .build() diff --git a/src/mito2/src/worker/handle_close.rs b/src/mito2/src/worker/handle_close.rs index 5a0406c655f..ede17511ad7 100644 --- a/src/mito2/src/worker/handle_close.rs +++ b/src/mito2/src/worker/handle_close.rs @@ -36,16 +36,12 @@ impl RegionWorkerLoop { info!("Try to close region {}, worker: {}", region_id, self.id); - // If the close request asks for a flush, or the region skips WAL, - // and has data in memtable and region is flushable (like, not in follower state), - // we should flush it before closing to ensure durability. - if (request.flush_on_close || region.skip_wal()) - && !region - .version_control - .current() - .version - .memtables - .is_empty() + // Memtables containing skip-WAL writes must be flushed before closing. + // Followers cannot flush; explicit flush-on-close and region WAL options + // retain their existing behavior. + let version = region.version(); + if (request.flush_on_close || region.skip_wal() || version.memtables.has_skip_wal_writes()) + && !version.memtables.is_empty() && region.is_flushable() { info!("Region {} has pending data, waiting for flush", region_id); @@ -69,7 +65,7 @@ impl RegionWorkerLoop { return; } - // WAL configured or memtable is empty, flush is not necessary. + // No close-time flush is required, or the region cannot flush. self.remove_region(region_id).await; info!("Region {} closed, worker: {}", region_id, self.id); sender.send(Ok(0)) diff --git a/src/mito2/src/worker/handle_manifest.rs b/src/mito2/src/worker/handle_manifest.rs index b574e490ac7..8aac67dc96f 100644 --- a/src/mito2/src/worker/handle_manifest.rs +++ b/src/mito2/src/worker/handle_manifest.rs @@ -257,13 +257,10 @@ impl RegionWorkerLoop { } else { None }; - let new_mutable = Arc::new( - region - .version() - .memtables - .mutable - .new_with_part_duration(version.compaction_time_window, memtable_builder), - ); + let new_mutable = version + .memtables + .mutable + .new_with_part_duration(version.compaction_time_window, memtable_builder); // Here it assumes the leader has backfilled the partition_expr of the metadata. let metadata = manifest.metadata.clone(); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 323d73c60ec..41af569c649 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -33,7 +33,7 @@ use crate::metrics::{ WRITE_REJECT_TOTAL, WRITE_ROWS_TOTAL, WRITE_STAGE_ELAPSED, WRITE_STALL_TOTAL, }; use crate::region::{RegionLeaderState, RegionRoleState}; -use crate::region_write_ctx::RegionWriteCtx; +use crate::region_write_ctx::{RegionWriteCtx, WriteSource}; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; use crate::wal::Wal; use crate::worker::RegionWorkerLoop; @@ -126,9 +126,7 @@ impl RegionWorkerLoop { if region_ctxs.len() == 1 { // fast path for single region. let mut region_ctx = region_ctxs.into_values().next().unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; } else { @@ -137,9 +135,7 @@ impl RegionWorkerLoop { .map(|mut region_ctx| { // use tokio runtime to schedule tasks. common_runtime::spawn_global(async move { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; (region_ctx.put_num, region_ctx.delete_num) }) }) @@ -348,6 +344,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -492,6 +489,7 @@ impl RegionWorkerLoop { ®ion.version_control, region.provider.clone(), Some(region.region_stats.written_bytes.clone()), + WriteSource::Request, ); e.insert(region_ctx); @@ -897,6 +895,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); ctx.push_mutation( @@ -946,6 +945,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(1), None, skip_wal)); @@ -959,9 +959,7 @@ mod tests { if skip_wal { let ctx = contexts.get_mut(®ion_id).unwrap(); assert_eq!(ctx.next_entry_id(), 1); - ctx.write_memtable().await; - ctx.write_bulk().await; - ctx.publish_sequence_and_entry_id(); + ctx.write_memtables().await; assert_eq!(version_control.committed_sequence(), 1); assert_eq!(version_control.current().last_entry_id, 0); } @@ -995,9 +993,7 @@ mod tests { assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); for region_ctx in region_ctxs.values_mut() { - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; } assert_eq!( @@ -1034,6 +1030,7 @@ mod tests { &version_control, Provider::raft_engine_provider(region_id.as_u64()), None, + WriteSource::Request, ); let (tx, rx) = oneshot::channel(); assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None, skip_wal)); @@ -1050,9 +1047,7 @@ mod tests { let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); - region_ctx.write_memtable().await; - region_ctx.write_bulk().await; - region_ctx.publish_sequence_and_entry_id(); + region_ctx.write_memtables().await; }); tokio::time::timeout( From 6e1a1374111887fa31df1b9573e27d37a2b06329 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 9 Oct 2026 12:39:51 +0000 Subject: [PATCH 127/133] ci: enhance benchmark workflows with version comparisons and diagnostics (#9383) * ci: update Agent and Traces benchmark runtime Signed-off-by: WenyXu * ci: use unrestricted benchmark runtimes by default Signed-off-by: WenyXu * ci: retry benchmark image pulls before generating datasets Signed-off-by: WenyXu * ci: increase tracesbench defaults to 8 CPU and 32 GiB Signed-off-by: WenyXu * ci: configure tracesbench query timeout and update runtime Signed-off-by: WenyXu * ci: update tracesbench runtime and configure import scan workers Signed-off-by: WenyXu * ci: use shared parallel tracesbench import preparation Signed-off-by: WenyXu * ci: update tracesbench runtime to preserve Tempo trace requests Signed-off-by: WenyXu * ci: summarize tracesbench query results and lifecycle Signed-off-by: WenyXu * ci: show benchmark resource usage and update parallel loader runtime Signed-off-by: WenyXu * fix(ci): resolve shared summary module in existing test harness Signed-off-by: WenyXu * ci: select tracesbench engines and GreptimeDB image version Signed-off-by: WenyXu * ci: update benchmark runtime with parallel Agent preparation Signed-off-by: WenyXu * ci: select benchmark runtime registry by cloud provider Signed-off-by: WenyXu * ci: retry long-range benchmark image pulls Signed-off-by: WenyXu * ci: compare baseline and candidate GreptimeDB benchmarks Signed-off-by: WenyXu * ci: pin published dual-version benchmark runtime Signed-off-by: WenyXu * ci: prepare agent benchmark indexes with updated runtime Signed-off-by: WenyXu * ci: validate agent admission and render traces version comparison Signed-off-by: WenyXu * ci: use runtime with traces index preparation Signed-off-by: WenyXu * ci: use consolidated runtime for benchmark version reports Signed-off-by: WenyXu * test: add license header to benchmark comparison tests Signed-off-by: WenyXu * fix(ci): read legacy Tempo ingestion metrics in traces summaries Signed-off-by: WenyXu * ci: require tracesbench query plan artifacts Signed-off-by: WenyXu * ci: use optimized tracesbench report runtime Signed-off-by: WenyXu * ci: update tracesbench runtime to merged report validation Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../scripts/agent-observability-summary.py | 3 + .../scripts/benchmark-version-comparison.py | 191 +++++++++++++++ .github/scripts/benchmark_resource_summary.py | 99 ++++++++ .github/scripts/pull-benchmark-image.sh | 25 ++ .github/scripts/tracesbench-summary.py | 169 +++++++++++++ .github/workflows/agent-observability.yml | 69 +++++- .github/workflows/checks.yml | 1 + .github/workflows/tracesbench.yml | 226 +++++++++++++++--- .github/workflows/vmbench-long-range.yml | 41 +++- .../perf/test_agent_observability_summary.py | 2 + .../perf/test_benchmark_version_comparison.py | 138 +++++++++++ tests/perf/test_ci_runner_scripts.py | 210 ++++++++++++++++ tests/perf/test_tracesbench_summary.py | 54 +++++ 13 files changed, 1187 insertions(+), 41 deletions(-) create mode 100644 .github/scripts/benchmark-version-comparison.py create mode 100644 .github/scripts/benchmark_resource_summary.py create mode 100644 .github/scripts/pull-benchmark-image.sh create mode 100644 .github/scripts/tracesbench-summary.py create mode 100644 tests/perf/test_benchmark_version_comparison.py create mode 100644 tests/perf/test_tracesbench_summary.py diff --git a/.github/scripts/agent-observability-summary.py b/.github/scripts/agent-observability-summary.py index 38ba4e9dbdd..33ba3ad7e6e 100644 --- a/.github/scripts/agent-observability-summary.py +++ b/.github/scripts/agent-observability-summary.py @@ -18,6 +18,8 @@ import json import math from pathlib import Path +from benchmark_resource_summary import render_resources + TARGETS = ('greptimedb', 'clickhouse', 'victorialogs') @@ -201,6 +203,7 @@ def lifecycle(root, targets): rows.append([context, entry.get('phase'), duration(entry.get('elapsed_ms')) if valid else 'invalid timestamp units (legacy artifact)', entry.get('exit_code')]) lines += ['## Phase timings', '', 'Phase duration includes orchestration overhead. Container cleanup does not prove ECS teardown.', '', table(['DB / scope', 'Phase', 'Duration', 'Exit code'], rows), ''] + lines.append(render_resources([('dataset', root)] + [(t, root / f'{t}-1') for t in targets], table)) return '\n'.join(lines) diff --git a/.github/scripts/benchmark-version-comparison.py b/.github/scripts/benchmark-version-comparison.py new file mode 100644 index 00000000000..1aa744df9ad --- /dev/null +++ b/.github/scripts/benchmark-version-comparison.py @@ -0,0 +1,191 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare validated GreptimeDB query results; never gate CI on noisy latency deltas.""" +import argparse +import html +import json +import math +import os +from pathlib import Path +import re +import statistics + + +def read(path): + return json.loads(path.read_text()) if path.is_file() else {} + + +def validate(env): + enabled = env.get('COMPARE_GREPTIMEDB', 'false') + if enabled not in ('true', 'false'): + raise ValueError('COMPARE_GREPTIMEDB must be true or false') + if enabled == 'false': + return + if env.get('RUN_GREPTIMEDB') != 'true': + raise ValueError('Version comparison requires Run GreptimeDB') + for key in ('BASELINE_GREPTIMEDB_TAG', 'GREPTIMEDB_TAG'): + if not re.fullmatch(r'[a-zA-Z0-9_][a-zA-Z0-9_.-]{0,127}(@sha256:[0-9a-f]{64})?', env.get(key, '')): + raise ValueError(f'Invalid {key}') + if env['BASELINE_GREPTIMEDB_TAG'] == env['GREPTIMEDB_TAG']: + raise ValueError('Baseline and candidate must have different image tags') + + +def same(a, b, fields): + return all(a.get(k) is not None and a.get(k) == b.get(k) for k in fields) + + +def finite(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + + +def improvement(old, new, throughput=False): + if not finite(old) or not finite(new) or old <= 0: + return None + return (new / old - 1 if throughput else 1 - new / old) * 100 + + +def row(query, metric, old, new, compatible, reason): + delta = improvement(old, new, metric == 'QPS') if compatible else None + return dict(query=query, metric=metric, baseline=old, candidate=new, + improvement_percent=delta, status='comparable' if delta is not None else reason) + + +def agent(root): + def load(r): + return read(r / 'greptimedb-1/greptimedb-1/measured-summary.json') + a, b = load(root / 'baseline'), load(root) + finished = [r / 'greptimedb-1/exit-code.txt' for r in (root / 'baseline', root)] + compatible = (all(p.is_file() and p.read_text().strip() == '0' for p in finished) + and a.get('passed') is True and b.get('passed') is True + and same(a, b, ('contract_sha256', 'data_profile', 'query_window', 'warmup_runs', 'measured_runs'))) + qa, qb = ({q['query_id']: q for q in x.get('per_query', [])} for x in (a, b)) + rows = [] + for q in sorted(qa.keys() | qb.keys()): + old, new = qa.get(q, {}), qb.get(q, {}) + ok = compatible and same(old, new, ('result_fingerprint', 'success')) and old.get('success', 0) > 0 + ok = ok and all(v.get('errors') == 0 and v.get('timeouts') == 0 for v in (old, new)) + for key in ('p50_ms', 'p95_ms', 'p99_ms'): + rows.append(row(q, key, old.get(key), new.get(key), ok, 'failed, missing or incompatible result')) + return rows + + +def traces(root): + a, b = (read(r / 'runs/greptimedb/run.json') for r in (root / 'baseline', root)) + compatible = (a.get('success') is True and b.get('success') is True + and same(a, b, ('expected', 'protocol', 'audit_expected', 'corpus')) + and same(a.get('execution', {}), b.get('execution', {}), + ('db_cpus', 'db_memory', 'load_workers_requested', 'max_in_flight'))) + rows = [] + for q in (f'T{i:02}' for i in range(1, 9)): + values, identities, counts = [], [], [] + ok = compatible + for run in (a, b): + correctness = [s for s in run.get('correctness', []) if s.get('query_id') == q] + all_samples = [s for s in run.get('samples', []) if s.get('query_id') == q] + measured = [s for s in all_samples if s.get('phase') == 'measured'] + passed = bool(correctness and measured) and all(s.get('outcome') == 'passed' for s in correctness + all_samples) + expected = run.get('expected', {}).get('contracts', {}).get(q, {}) + passed = passed and all(s.get('fingerprint') == expected.get('fingerprint') and s.get('fingerprint') is not None + and s.get('row_count') == expected.get('row_count') for s in correctness + all_samples) + elapsed = [s.get('elapsed_ms') for s in measured] + passed = passed and all(finite(v) for v in elapsed) + ok = ok and passed + values.append(statistics.median(elapsed) if passed else None) + identities.append(expected) + counts.append((len(correctness), len(all_samples), len(measured))) + ok = ok and identities[0] == identities[1] and counts[0] == counts[1] + rows.append(row(q, 'median_ms', *values, ok, 'failed, skipped, missing or incompatible result')) + return rows + + +def long_range(root): + rows = [] + manifest = read(root / 'run-manifest.json') + healthy = all( + all(field in (manifest.get(phase, {}).get(target, {}).get('docker', {}) if phase == 'post_load_results' + else manifest.get(phase, {}).get(target, {})).get('state', '').split() + for field in ('Running=true', 'ExitCode=0', 'OOMKilled=false')) + for phase in ('post_load_results', 'post_query_results') + for target in ('greptimedb-baseline', 'greptimedb')) + for mode in ('serial', 'concurrent-duration-c1', 'concurrent-duration-c4'): + a, b = (read(root / name / 'results-vmbench-long-range' / (mode + '-summary.json')) + for name in ('greptimedb-baseline', 'greptimedb')) + fields = ('workload', 'query_file_sha256', 'start', 'end', 'step', 'profile', 'data_duration', + 'scrape_interval', 'cache_policy_effective', 'mode', 'connection_mode', 'request_timeout_seconds') + fields += ('tries', 'warmup_runs') if mode == 'serial' else ('concurrency', 'duration_seconds', 'warmup_seconds') + compatible = healthy and same(a, b, fields) and all(x.get('errors') == 0 and x.get('timeouts') == 0 for x in (a, b)) + qa, qb = ({q['query_index']: q for q in x.get('per_query', [])} for x in (a, b)) + for q in sorted(qa.keys() | qb.keys()): + old, new = qa.get(q, {}), qb.get(q, {}) + ok = compatible and same(old, new, ('query', 'series', 'points')) + ok = ok and all(v.get('success', 0) > 0 and v.get('series', 0) > 0 and v.get('points', 0) > 0 + and v.get('errors') == 0 and v.get('timeouts') == 0 for v in (old, new)) + for key in ('p50_ms', 'qps') if mode != 'serial' else ('p50_ms',): + rows.append(row(f'{mode}/Q{q}', 'QPS' if key == 'qps' else key, + old.get(key), new.get(key), ok, 'failed, missing or incompatible result')) + return rows + + +def render(document): + def cell(value): + return html.escape(str(value)).replace('|', '|').replace('\n', ' ') + def value(v, metric): + if not finite(v): + return '—' + if metric == 'QPS': + return f'{v:.3f}' + return f'{v / 1000:.2f} s' if v >= 1000 else f'{v:.2f} ms' + lines = ['# GreptimeDB version comparison', '', + f'Baseline: `{cell(document["baseline_tag"])}` → candidate: `{cell(document["candidate_tag"])}`.', '', + 'Same-host sequential runs with fresh databases and shared input data. Positive = improvement; ' + 'latency: (baseline − candidate) / baseline; QPS: (candidate − baseline) / baseline. ' + 'Single runs are diagnostic, not a statistically established regression. No performance threshold gates CI.', '', + '| Query | Metric | Baseline | Candidate | Improvement | Status |', + '| --- | --- | --- | --- | --- | --- |'] + for r in document['rows']: + delta = r['improvement_percent'] + lines.append('| ' + ' | '.join(map(cell, [r['query'], r['metric'], value(r['baseline'], r['metric']), + value(r['candidate'], r['metric']), f'{delta:+.1f}%' if delta is not None else '—', r['status']])) + ' |') + if document.get('kind') == 'long-range': + lines += ['', 'Long-range retains the existing non-empty series/point-count checks; these do not establish equality of every returned value.'] + if not document['rows']: + lines += ['', 'Results unavailable; comparison not performed.'] + return '\n'.join(lines) + '\n' + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('command', choices=('validate', 'report')) + parser.add_argument('--kind', choices=('agent', 'traces', 'long-range')) + parser.add_argument('--root', type=Path) + args = parser.parse_args() + if args.command == 'validate': + validate(os.environ) + return + if args.root is None or args.kind is None: + parser.error('report requires --root and --kind') + document = {'kind': args.kind, 'baseline_tag': os.environ.get('BASELINE_GREPTIMEDB_TAG', ''), + 'candidate_tag': os.environ.get('GREPTIMEDB_TAG', ''), + 'rows': {'agent': agent, 'traces': traces, 'long-range': long_range}[args.kind](args.root)} + text = render(document) + # Preserve partial comparisons even if a target failed before creating output. + args.root.mkdir(parents=True, exist_ok=True) + (args.root / 'version-comparison.json').write_text(json.dumps(document, indent=2, allow_nan=False) + '\n') + (args.root / 'version-comparison.md').write_text(text) + print(text, end='') + + +if __name__ == '__main__': + main() diff --git a/.github/scripts/benchmark_resource_summary.py b/.github/scripts/benchmark_resource_summary.py new file mode 100644 index 00000000000..3a0a896c46e --- /dev/null +++ b/.github/scripts/benchmark_resource_summary.py @@ -0,0 +1,99 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Render shared O11yBench sampled resource artifacts for lifecycle summaries.""" +import json +import math + + +def number(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + + +def metric(values, key, statistic): + entry = values.get(key, {}) + value = entry.get(statistic) + return value if number(entry.get('samples')) and entry['samples'] > 0 and number(value) else None + + +def decimal(value, suffix=''): + return f'{value:.2f}{suffix}' if number(value) else '—' + + +def memory(value): + if not number(value): + return '—' + for unit in ('B', 'KiB', 'MiB', 'GiB', 'TiB'): + if value < 1024 or unit == 'TiB': + return decimal(value, f' {unit}') + value /= 1024 + + +def capacities(path): + # Quotas are metadata, not recomputed CPU statistics. Bound reads on long runs. + result = {} + try: + with path.open() as source: + for index, line in enumerate(source): + if index == 256: + break + record = json.loads(line) + cores = record.get('host', {}).get('cpu_count') + if number(cores) and cores > 0: + result.setdefault('host', cores) + for role, values in record.get('containers', {}).items(): + quota = values.get('configured_cpu_limit_cores') + if number(quota): + result.setdefault(role, quota) + if all(role in result for role in ('host', 'runtime', 'db')): + break + except (OSError, ValueError, TypeError, AttributeError): + pass + return result + + +def render_resources(scopes, table): + rows = [] + for scope, directory in scopes: + try: + summary = json.loads((directory / 'resources.summary.json').read_text()) + phases = summary.get('phases', {}) if summary.get('schema') == 'o11ybench-resources-v1' else {} + except (OSError, ValueError, TypeError, AttributeError): + phases = {} + if not isinstance(phases, dict) or not phases: + rows.append([scope, 'unavailable', *['—'] * 9]) + continue + limits = capacities(directory / 'resources.jsonl') + for phase, data in phases.items(): + for role in ('host', 'runtime', 'db'): + values = data.get('host', {}) if role == 'host' else data.get('containers', {}).get(role, {}) + if role == 'db' and scope == 'dataset': + continue + quota = limits.get(role) + cpu_samples = values.get('cpu_used_cores', {}).get('samples') + rows.append([scope, phase, role, cpu_samples, + decimal(metric(values, 'cpu_used_cores', 'mean')), + decimal(metric(values, 'cpu_used_cores', 'max')), + 'unlimited' if quota == 0 else decimal(quota), + memory(metric(values, 'memory_used_bytes' if role == 'host' else 'memory_usage_bytes', 'max')), + decimal(metric(values, 'io_wait_percent', 'mean'), '%'), + decimal(metric(values, 'io_wait_percent', 'max'), '%'), + memory(metric(values, 'swap_used_bytes', 'max'))]) + return '\n'.join(['## Sampled resources', '', + 'CPU is measured in cores. Mean/max summarize available samples, not continuous peaks; ' + 'short phases may have no CPU sample. CPU capacity is host cores or the configured container quota. ' + 'Host memory is total minus available; Docker memory excludes cache where supported. ' + 'Swap is host-wide usage, not swap activity. Missing evidence stays —.', '', + table(['Scope', 'Phase', 'Component', 'CPU samples', 'Mean CPU cores', 'Max CPU cores', + 'CPU capacity', 'Max memory', 'Mean I/O wait', 'Max I/O wait', 'Max swap'], rows), '']) diff --git a/.github/scripts/pull-benchmark-image.sh b/.github/scripts/pull-benchmark-image.sh new file mode 100644 index 00000000000..1a26a4aeaee --- /dev/null +++ b/.github/scripts/pull-benchmark-image.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Pull only: never retry database writes, benchmarks, or cloud provisioning. +set -euo pipefail + +if [[ $# != 1 || -z "$1" || "$1" == -* ]]; then + echo "Usage: $0 IMAGE" >&2 + exit 2 +fi +image=$1 +attempts=3 +for ((attempt=1; attempt<=attempts; attempt++)); do + echo "Pulling $image (attempt $attempt/$attempts)" >&2 + if timeout --kill-after=10s 180s docker pull "$image"; then + exit 0 + else + status=$? + fi + if ((attempt == attempts)); then + echo "Failed to pull $image after $attempts attempts (exit $status)" >&2 + exit "$status" + fi + delay=$((attempt * 10)) + echo "Pull failed (exit $status); retrying in ${delay}s" >&2 + sleep "$delay" +done diff --git a/.github/scripts/tracesbench-summary.py b/.github/scripts/tracesbench-summary.py new file mode 100644 index 00000000000..087a67498f4 --- /dev/null +++ b/.github/scripts/tracesbench-summary.py @@ -0,0 +1,169 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import html +import math +from pathlib import Path + +from benchmark_resource_summary import render_resources + +TARGETS = ('greptimedb', 'victoriatraces', 'tempo') + + +def duration(ms): + if not isinstance(ms, (int, float)) or not math.isfinite(ms) or ms < 0: + return '—' + if ms < 1: + return f'{ms * 1000:.0f} µs' + if ms < 1000: + return f'{ms:.1f} ms' + if ms < 60_000: + return f'{ms / 1000:.2f} s' + minutes, seconds = divmod(round(ms / 1000), 60) + hours, minutes = divmod(minutes, 60) + return f'{hours}h {minutes}m {seconds}s' if hours else f'{minutes}m {seconds}s' + + +def cell(value): + return html.escape(str(value if value is not None else '—')).replace('|', '|').replace('\n', ' ') + + +def read(path): + if not path.exists(): + return {} + return json.loads(path.read_text()) + + +def table(headers, rows): + return '\n'.join(['| ' + ' | '.join(map(cell, headers)) + ' |', + '| ' + ' | '.join('---' for _ in headers) + ' |'] + + ['| ' + ' | '.join(map(cell, row)) + ' |' for row in rows]) + + +def delta(value, baseline): + if not all(isinstance(v, (int, float)) and math.isfinite(v) for v in (value, baseline)): + return '—' + return f'{(value / baseline - 1) * 100:+.1f}%' if baseline > 0 and value >= 0 else '—' + + + +def load(root): + report = read(root / 'report.json') + runs = [read(p) for p in sorted((root / 'runs').glob('*/run.json'))] + return report, runs + + +def benchmark(report, runs, targets=TARGETS): + lines = ['# Tracesbench benchmark summary', '', + 'Median Δ = (target / GreptimeDB − 1) × 100%; negative is faster. ' + 'Min/median/max come from the validated report, not recalculated samples. ' + 'Counts exclude warmup. Missing or excluded results have no latency comparison.', + '', f'Classification: {cell(report.get("classification"))}. ' + 'See the uploaded HTML for query semantics and client-side compensation.', ''] + statistics = {} + for corpus in report.get('per_corpus', []): + for stat in corpus.get('statistics', []): + if stat['target'] not in targets: + continue + statistics[(corpus['corpus'], stat['round'], stat['target'], stat['query_id'])] = stat + rows, details = [], [] + for target in targets: + selected = [run for run in runs if run.get('target') == target] + if not selected: + lines.append(f'- **{target}**: unavailable / not run.') + for run in selected: + lines.append(f'- **{target}**: lifecycle ' + f'{"completed" if run.get("success") is True else "failed / incomplete"}; ' + f'query failures: {len(run.get("query_failures", []))}.') + for qid in (f'T{i:02}' for i in range(1, 9)): + correct = [q for q in run.get('correctness', []) if q.get('query_id') == qid] + measured = [q for q in run.get('samples', []) + if q.get('query_id') == qid and q.get('phase') == 'measured'] + errors = [q for q in correct + run.get('samples', []) + if q.get('query_id') == qid and q.get('outcome') not in ('passed', 'not_executed')] + reason = (run.get('exclusions', {}).get(qid) or + run.get('execution', {}).get('query_skips', {}).get(qid)) + stat = statistics.get((run.get('corpus'), run.get('round'), target, qid), {}) + base = statistics.get((run.get('corpus'), run.get('round'), 'greptimedb', qid), {}) + status = ('failed' if errors else 'excluded' if reason else + 'passed' if stat else + 'correctness passed; report unavailable' if correct and + all(q.get('outcome') == 'passed' for q in correct) else 'unavailable / incomplete') + rows.append([run.get('corpus'), run.get('round'), qid, target, status, + duration(stat.get('min_ms')), duration(stat.get('median_ms')), + duration(stat.get('max_ms')), + delta(stat.get('median_ms'), base.get('median_ms')) if target != 'greptimedb' else '—', + sum(q.get('outcome') == 'passed' for q in measured), + sum(q.get('outcome') != 'passed' for q in measured)]) + if reason or errors: + messages = [str(q.get('error') or q.get('outcome')) for q in errors] + details.append(f'- **{target} / {qid}**: {cell(str(reason or "") + " " + "; ".join(messages))[:4000]}') + lines += ['', table(['Dataset', 'Round', 'Query', 'DB', 'Status', 'Min', 'Median', 'Max', + 'Median Δ', 'Measured success', 'Measured errors'], rows), '', + '
', 'Query failures and exclusions', '', *details, '', '
'] + if not report: + lines += ['', 'Validated report unavailable; raw outcomes are shown without latency statistics.'] + return '\n'.join(lines) + + +def lifecycle(report, runs, root, targets=TARGETS): + rows = [] + for target in targets: + selected = [r for r in runs if r.get('target') == target] or [{}] + for run in selected: + load = run.get('load', {}) + ingest = load.get('ingest', {}) + if target == 'tempo' and 'ingest' not in load: + ingest = load.get('transport', {}).get('ingest', {}) + health = run.get('post_load_health', {}) + execution = run.get('execution', {}) + start, finish = run.get('started_unix_ns'), run.get('finished_unix_ns') + elapsed = (finish - start) / 1_000_000 if isinstance(start, int) and isinstance(finish, int) else None + seconds = ingest.get('wall_seconds') + rows.append([target, run.get('corpus'), run.get('load_summary', {}).get('after'), + duration(elapsed), duration(seconds * 1000 if isinstance(seconds, (int, float)) else None), + ingest.get('successful_requests'), ingest.get('failed_requests'), + execution.get('load_workers_requested'), execution.get('max_in_flight'), + execution.get('db_cpus'), execution.get('db_memory'), + health.get('running'), health.get('oom_killed'), + run.get('cleanup', {}).get('verified_absent')]) + return '\n'.join(['# Tracesbench load / lifecycle summary', '', + 'Total is the recorded target lifecycle; ingest is the sender wall time, ' + 'not the entire load/readiness phase. Health is the post-load snapshot. ' + 'Missing evidence remains blank.', '', + table(['DB', 'Dataset', 'Loaded spans', 'Total', 'Ingest', + 'Write success', 'Write errors', 'Load workers', 'Max in flight', + 'DB CPU', 'DB memory', 'Running after load', 'OOM after load', + 'Container removed'], rows), '', + render_resources([('dataset', root / 'runs' / 'generate')] + + [(t, root / 'runs' / t) for t in targets], table)]) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--root', type=Path, required=True) + parser.add_argument('--section', choices=('benchmark', 'lifecycle'), required=True) + parser.add_argument('--targets', default=','.join(TARGETS), help='Comma-separated selected engines') + args = parser.parse_args() + targets = args.targets.split(',') + if not targets or any(t not in TARGETS for t in targets) or len(set(targets)) != len(targets): + parser.error('targets must be a non-empty comma-separated subset of greptimedb,victoriatraces,tempo') + report, runs = load(args.root) + print(benchmark(report, runs, targets) if args.section == 'benchmark' else lifecycle(report, runs, args.root, targets)) + + +if __name__ == '__main__': + main() diff --git a/.github/workflows/agent-observability.yml b/.github/workflows/agent-observability.yml index 86713c911ed..d265801d4a0 100644 --- a/.github/workflows/agent-observability.yml +++ b/.github/workflows/agent-observability.yml @@ -26,8 +26,18 @@ on: - M default: S required: true + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false greptimedb_tag: - description: Docker Hub greptime/greptimedb tag + description: Candidate Docker Hub greptime/greptimedb tag type: string default: latest required: true @@ -46,9 +56,9 @@ on: type: string required: false runtime_image: - description: Aliyun runtime image (tag or digest reference) + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260917-8111034e@sha256:ce9dee2abb1e70f907a09ea01854904a91cc97748b0c4ef6f33d9ec34044bec0 + default: auto required: true provider: description: Cloud provider for the benchmark runner @@ -89,12 +99,14 @@ on: permissions: contents: read env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} PROFILE: ${{ inputs.dataset }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} CLICKHOUSE_TAG: ${{ inputs.clickhouse_tag }} VICTORIALOGS_TAG: ${{ inputs.victorialogs_tag }} O11YBENCH_REF: ${{ inputs.o11ybench_ref }} - RUNTIME_IMAGE: ${{ inputs.runtime_image }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:2f40fb3fdbdc@sha256:8940782216bf5d066308fd71f6e346d950073c41cfb5756c5f3dc8aacf069cbc' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:2f40fb3fdbdc@sha256:8940782216bf5d066308fd71f6e346d950073c41cfb5756c5f3dc8aacf069cbc' }} PROVIDER: ${{ inputs.provider }} REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} @@ -118,6 +130,12 @@ jobs: with: path: greptimedb persist-credentials: false + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate - name: Resolve CI runner id: runner run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" @@ -126,6 +144,16 @@ jobs: python3 greptimedb/tests/perf/test_aliyun_ecs_runner_scripts.py python3 greptimedb/tests/perf/test_ci_runner_scripts.py python3 greptimedb/tests/perf/test_agent_observability_summary.py + - name: Check runtime supports shared version-comparison data + if: inputs.compare_greptimedb + run: | + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" + docker run --rm --entrypoint grep "$RUNTIME_IMAGE" -q -- '--data-root)' /opt/o11ybench/scripts/run-agent-observability.sh \ + || { echo 'Comparison requires a runtime with --data-root support'; exit 1; } + - name: Validate dataset admission in runtime + run: | + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" + docker run --rm "$RUNTIME_IMAGE" scripts/logbench_agent_observability_measured.py --preflight-profile "$PROFILE" - name: Validate requested benchmark id: inputs shell: bash @@ -262,11 +290,11 @@ jobs: run: | set -euo pipefail actual= - [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] images='{}' resolve_image() { local target=$1 reference=$2 variable=$3 metadata resolved revision - docker pull "$reference" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$reference" metadata=$(docker image inspect "$reference") resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") if [[ "$target" == runtime ]]; then @@ -289,6 +317,9 @@ jobs: victorialogs) resolve_image "$target" "victoriametrics/victoria-logs:$VICTORIALOGS_TAG" VICTORIALOGS_IMAGE ;; esac done + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + resolve_image greptimedb-baseline "greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG" BASELINE_GREPTIMEDB_IMAGE + fi jq -n --arg sha "$actual" --arg targets "$TARGETS" --arg profile "$PROFILE" \ --arg cpus "$DB_CPUS" --arg memory "$DB_MEMORY" --argjson images "$images" \ --arg provider "$PROVIDER" --arg instance_type "$INSTANCE_TYPE" --argjson disk_gib "$SYSTEM_DISK_GIB" \ @@ -312,6 +343,15 @@ jobs: --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" \ --root "$ARTIFACT_ROOT" --profile "$PROFILE" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" cp "$MANIFEST_PATH" "$ARTIFACT_ROOT/run-manifest.json" + - name: Load and benchmark baseline GreptimeDB + id: baseline_greptimedb + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' && inputs.compare_greptimedb }} + run: |- + GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage target --execute \ + --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" --root "$ARTIFACT_ROOT/baseline" --data-root "$ARTIFACT_ROOT" \ + --profile "$PROFILE" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --target greptimedb --image "$BASELINE_GREPTIMEDB_IMAGE" \ + --greptimedb-http-port "$GREPTIMEDB_HTTP_PORT" --greptimedb-grpc-port "$GREPTIMEDB_GRPC_PORT" - name: Load and benchmark greptimedb id: greptimedb if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' }} @@ -357,6 +397,13 @@ jobs: GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage cleanup --execute \ --owner "$OWNER" --root "$ARTIFACT_ROOT" fi + - name: Clean up baseline containers + if: always() && inputs.compare_greptimedb + run: |- + if [[ -f o11ybench/scripts/run-agent-observability.sh ]]; then + GITHUB_STEP_SUMMARY= bash o11ybench/scripts/run-agent-observability.sh --stage cleanup --execute \ + --owner "$OWNER-base" --root "$ARTIFACT_ROOT/baseline" + fi - name: Benchmark summary if: always() run: | @@ -367,6 +414,16 @@ jobs: run: | python3 greptimedb/.github/scripts/agent-observability-summary.py \ --root "$ARTIFACT_ROOT" --targets "$TARGETS" --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: Baseline load / lifecycle summary + if: always() && inputs.compare_greptimedb + run: | + echo '## Baseline GreptimeDB' >> "$GITHUB_STEP_SUMMARY" + python3 greptimedb/.github/scripts/agent-observability-summary.py \ + --root "$ARTIFACT_ROOT/baseline" --targets greptimedb --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb + run: | + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind agent --root "$ARTIFACT_ROOT" >> "$GITHUB_STEP_SUMMARY" - name: Upload results (no dataset) if: always() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index d3bad06f41d..deba5010c80 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -99,6 +99,7 @@ jobs: python3 tests/perf/test_aliyun_ecs_runner_scripts.py python3 tests/perf/test_ci_runner_scripts.py python3 tests/perf/test_agent_observability_summary.py + python3 tests/perf/test_tracesbench_summary.py check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' && (github.event_name != 'pull_request' || github.event.pull_request.draft == false || github.run_attempt > 1) }} diff --git a/.github/workflows/tracesbench.yml b/.github/workflows/tracesbench.yml index 60081d69ad9..049faaef1a7 100644 --- a/.github/workflows/tracesbench.yml +++ b/.github/workflows/tracesbench.yml @@ -2,6 +2,36 @@ name: Tracesbench on: workflow_dispatch: inputs: + run_greptimedb: + description: Run GreptimeDB + type: boolean + default: true + required: false + run_victoriatraces: + description: Run VictoriaTraces + type: boolean + default: false + required: false + run_tempo: + description: Run Tempo + type: boolean + default: false + required: false + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false + greptimedb_tag: + description: Candidate Docker Hub greptime/greptimedb tag + type: string + default: latest + required: true dataset: description: Dataset profile (main is experimental; native queries may exceed memory limits) type: choice @@ -9,9 +39,9 @@ on: default: sanity required: true runtime_image: - description: Aliyun runtime image (tag or digest reference) + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260928-4e52d0b0@sha256:39d2d87593788e0c47a54768f6323679b8e03d7dc72a8158a9b22cf29f883baa + default: auto required: true provider: description: Cloud provider for the benchmark runner @@ -20,7 +50,7 @@ on: default: Aliyun required: true instance_type: - description: 'Instance type; auto = Aliyun: ecs.c9i.2xlarge, AWS: c7i.2xlarge (both 8 vCPU / 16 GiB). Enter a specific type to override.' + description: 'Instance type; auto = Aliyun: ecs.g9i.2xlarge, AWS: m7i.2xlarge (both 8 vCPU / 32 GiB). Enter a specific type to override.' type: string default: auto required: true @@ -45,33 +75,46 @@ on: default: 8 required: true db_cpus: - description: CPU limit per DB container + description: CPU limit per DB container (default 8 vCPU) type: string default: '8' required: true db_memory: - description: Memory limit per DB container + description: Memory limit per DB container (default 32 GiB; shared host also runs the runtime and runner) type: string - default: 16g + default: 32g + required: true + load_workers: + description: CPU workers for parallel import reading/parsing across all engines (separate from write concurrency) + type: number + default: 4 + required: true + query_timeout_seconds: + description: Timeout per query request in seconds (correctness, warmup and measurement; 1..3600) + type: number + default: 300 required: true runtime_cpus: - description: CPU limit per runtime container + description: CPU limit per runtime container (0 = unrestricted) type: string - default: '2' + default: '0' required: true runtime_memory: - description: Memory limit per runtime container (separate from DB memory) + description: Memory limit per runtime container (0 = unrestricted; separate from DB memory) type: string - default: 4g + default: '0' required: true permissions: contents: read env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} DATASET: ${{ inputs.dataset }} + GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} RUNTIME_CPUS: ${{ inputs.runtime_cpus }} RUNTIME_MEMORY: ${{ inputs.runtime_memory }} BENCHMARK_ROOT: ${{ inputs.benchmark_root }} - RUNTIME_IMAGE: ${{ inputs.runtime_image }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:77a84e1c8136@sha256:c482993ec177b026235d6cd0c61d4fb6d6f1088bc78bdc0e5964a5398de6dba0' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:77a84e1c8136@sha256:c482993ec177b026235d6cd0c61d4fb6d6f1088bc78bdc0e5964a5398de6dba0' }} PROVIDER: ${{ inputs.provider }} REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} @@ -79,12 +122,18 @@ env: JANITOR_TTL_HOURS: ${{ inputs.janitor_ttl_hours }} DB_CPUS: ${{ inputs.db_cpus }} DB_MEMORY: ${{ inputs.db_memory }} + QUERY_TIMEOUT_SECONDS: ${{ inputs.query_timeout_seconds }} + LOAD_WORKERS: ${{ inputs.load_workers }} jobs: validate: runs-on: ubuntu-latest timeout-minutes: 15 outputs: instance_type: ${{ steps.runner.outputs.instance_type }} + targets: ${{ steps.inputs.outputs.targets }} + greptimedb: ${{ steps.inputs.outputs.greptimedb }} + victoriatraces: ${{ steps.inputs.outputs.victoriatraces }} + tempo: ${{ steps.inputs.outputs.tempo }} system_disk_gib: ${{ steps.runtime.outputs.system_disk_gib }} runtime_image: ${{ steps.runtime.outputs.runtime_image }} o11ybench_ref: ${{ steps.runtime.outputs.o11ybench_ref }} @@ -93,9 +142,44 @@ jobs: with: path: greptimedb persist-credentials: false + - name: Validate selected targets + id: inputs + env: + greptimedb: ${{ inputs.run_greptimedb }} + victoriatraces: ${{ inputs.run_victoriatraces }} + tempo: ${{ inputs.run_tempo }} + run: | + set -euo pipefail + selected= + for target in greptimedb victoriatraces tempo; do + enabled=${!target} + [[ "$enabled" == true || "$enabled" == false ]] + printf '%s=%s\n' "$target" "$enabled" >> "$GITHUB_OUTPUT" + if [[ "$enabled" == true ]]; then selected="${selected:+$selected,}$target"; fi + done + [[ -n "$selected" ]] || { echo 'Select at least one database' >&2; exit 1; } + if [[ "$greptimedb" == true ]]; then + [[ "$GREPTIMEDB_TAG" =~ ^[a-zA-Z0-9_][a-zA-Z0-9_.-]{0,127}$ ]] || { echo 'Invalid GreptimeDB image tag' >&2; exit 1; } + fi + printf 'targets=%s\n' "$selected" >> "$GITHUB_OUTPUT" + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate - name: Resolve CI runner id: runner - run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" + run: | + set -euo pipefail + instance_type="$REQUESTED_INSTANCE_TYPE" + if [[ "$instance_type" == auto ]]; then + case "$PROVIDER" in + Aliyun) instance_type=ecs.g9i.2xlarge ;; + AWS) instance_type=m7i.2xlarge ;; + esac + fi + python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$instance_type" - name: Validate budgets and shared runtime id: runtime run: | @@ -115,19 +199,38 @@ jobs: ((BENCHMARK_TIMEOUT_MINUTES <= 360 && JANITOR_TTL_HOURS <= 168)) ((JANITOR_TTL_HOURS * 60 >= BENCHMARK_TIMEOUT_MINUTES + 120)) [[ "$DB_CPUS" =~ ^[1-9][0-9]*$ && "$DB_MEMORY" =~ ^[1-9][0-9]*[mMgG]$ ]] - [[ "$RUNTIME_CPUS" =~ ^[1-9][0-9]*$ && "$RUNTIME_MEMORY" =~ ^[1-9][0-9]*[mMgG]$ ]] - [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] - docker pull "$RUNTIME_IMAGE" + [[ "$RUNTIME_CPUS" =~ ^(0|[1-9][0-9]*)$ && "$RUNTIME_MEMORY" =~ ^(0|[1-9][0-9]*[mMgG])$ ]] + [[ "$QUERY_TIMEOUT_SECONDS" =~ ^[1-9][0-9]{0,3}$ ]] + ((QUERY_TIMEOUT_SECONDS <= 3600)) + [[ "$LOAD_WORKERS" =~ ^[1-9][0-9]{0,2}$ ]] + ((LOAD_WORKERS <= 256)) + runtime_limits=() + [[ "$RUNTIME_CPUS" == 0 ]] || runtime_limits+=(--cpus "$RUNTIME_CPUS") + [[ "$RUNTIME_MEMORY" == 0 ]] || runtime_limits+=(--memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY") + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" metadata=$(docker image inspect "$RUNTIME_IMAGE") resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + docker run --rm "$resolved" scripts/run_tracesbench_docker.py --help | grep -q -- '--data-root' \ + || { echo 'Comparison requires a runtime with --data-root support'; exit 1; } + fi revision=$(jq -er '.[0].Config.Labels["org.opencontainers.image.revision"]' <<< "$metadata") [[ "$revision" =~ ^[0-9a-f]{40}$ ]] # Refuse runtimes without the trace runner and its validated dataset/report contract. - help=$(docker run --rm --cpus "$RUNTIME_CPUS" --memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY" "$resolved" scripts/run_tracesbench_docker.py --help) + help=$(docker run --rm "${runtime_limits[@]}" "$resolved" scripts/run_tracesbench_docker.py --help) + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + grep -F -- '--baseline-run' <<< "$help" \ + || { echo 'Comparison requires a runtime with --baseline-run report support'; exit 1; } + fi grep -F -- '--dataset {sanity,preflight,main}' <<< "$help" grep -F -- '--runtime-cpus' <<< "$help" grep -F -- '--runtime-memory' <<< "$help" - docker run --rm --cpus "$RUNTIME_CPUS" --memory "$RUNTIME_MEMORY" --memory-swap "$RUNTIME_MEMORY" "$resolved" scripts/check_run_tracesbench_docker.py + grep -F -- '--query-timeout-seconds' <<< "$help" + grep -F -- '--load-workers' <<< "$help" + grep -F -- '--targets' <<< "$help" + grep -F -- '--greptimedb-image' <<< "$help" + docker run --rm "${runtime_limits[@]}" "$resolved" scripts/check_run_tracesbench_docker.py printf 'runtime_image=%s\no11ybench_ref=%s\nsystem_disk_gib=%s\n' "$resolved" "$revision" "$SYSTEM_DISK_GIB" >> "$GITHUB_OUTPUT" provision: needs: validate @@ -190,6 +293,7 @@ jobs: runs-on: ${{ needs.provision.outputs.label }} timeout-minutes: ${{ fromJSON(format('{0}', inputs.benchmark_timeout_minutes)) }} env: + TARGETS: ${{ needs.validate.outputs.targets }} RESOLVED_RUNTIME_IMAGE: ${{ needs.validate.outputs.runtime_image }} RESOLVED_O11YBENCH_REF: ${{ needs.validate.outputs.o11ybench_ref }} RUN_ROOT: ${{ inputs.benchmark_root }}/tracesbench-${{ github.run_id }}-${{ github.run_attempt }} @@ -224,36 +328,83 @@ jobs: available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') ((available >= minimum_gib * 1024 * 1024)) || { echo "Need at least $minimum_gib GiB free on $path"; exit 1; } done - docker pull "$RESOLVED_RUNTIME_IMAGE" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE" docker image inspect "$RESOLVED_RUNTIME_IMAGE" > "$METADATA_ROOT/runtime-image.json" test "$(jq -r '.[0].Config.Labels["org.opencontainers.image.revision"]' "$METADATA_ROOT/runtime-image.json")" = "$RESOLVED_O11YBENCH_REF" - # Preserve the pinned catalog images and require all three targets in the report. + - name: Pull database images before generating data + run: | + set -euo pipefail + IFS=, read -ra targets <<< "$TARGETS" + for target in "${targets[@]}"; do + if [[ "$target" == greptimedb ]]; then + image="greptime/greptimedb:$GREPTIMEDB_TAG" + else + image=$(jq -er '.target_identity.image | select(test("@sha256:[a-f0-9]{64}$"))' \ + "o11ybench/workloads/tracesbench/catalogs/$target.json") + fi + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$image" + docker image inspect "$image" > "$METADATA_ROOT/$target-image.json" + if [[ "$target" == greptimedb ]]; then + resolved=$(jq -er '.[0].RepoDigests[0] | select(test("@sha256:[a-f0-9]{64}$"))' "$METADATA_ROOT/$target-image.json") + printf 'GREPTIMEDB_IMAGE=%s\n' "$resolved" >> "$GITHUB_ENV" + fi + done + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + image="greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$image" + docker image inspect "$image" > "$METADATA_ROOT/greptimedb-baseline-image.json" + resolved=$(jq -er '.[0].RepoDigests[0] | select(test("@sha256:[a-f0-9]{64}$"))' "$METADATA_ROOT/greptimedb-baseline-image.json") + printf 'BASELINE_GREPTIMEDB_IMAGE=%s\n' "$resolved" >> "$GITHUB_ENV" + fi - name: Generate selected dataset once + id: generate run: | python3 o11ybench/scripts/run_tracesbench_docker.py --stage generate \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" - - name: GreptimeDB load, correctness and benchmark + - name: Baseline GreptimeDB load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' && inputs.compare_greptimedb }} run: | - python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 \ - --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 --greptimedb-image "$BASELINE_GREPTIMEDB_IMAGE" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ + --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ + --root "$RUN_ROOT/baseline" --data-root "$RUN_ROOT" --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + python3 o11ybench/scripts/check_tracesbench_query_plans.py "$RUN_ROOT/baseline/runs/greptimedb" + - name: GreptimeDB load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.greptimedb == 'true' }} + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target greptimedb --port 14000 --greptimedb-image "$GREPTIMEDB_IMAGE" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + python3 o11ybench/scripts/check_tracesbench_query_plans.py "$RUN_ROOT/runs/greptimedb" - name: VictoriaTraces load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.victoriatraces == 'true' }} run: | python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target victoriatraces --port 19428 \ - --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" - name: Tempo load, correctness and benchmark + if: ${{ !cancelled() && steps.generate.outcome == 'success' && needs.validate.outputs.tempo == 'true' }} run: | python3 o11ybench/scripts/run_tracesbench_docker.py --stage target --target tempo --port 13200 \ - --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" \ + --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --query-timeout-seconds "$QUERY_TIMEOUT_SECONDS" --load-workers "$LOAD_WORKERS" \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" - - name: Validate and render report + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb run: | - python3 o11ybench/scripts/run_tracesbench_docker.py --stage report \ + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind traces --root "$RUN_ROOT" >> "$GITHUB_STEP_SUMMARY" + - name: Validate and render report + env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + run: | + comparison_args=() + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + comparison_args+=(--baseline-run "$RUN_ROOT/baseline/runs/greptimedb/run.json") + fi + python3 o11ybench/scripts/run_tracesbench_docker.py --stage report --targets "$TARGETS" "${comparison_args[@]}" \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" - name: Clean up owned containers @@ -262,6 +413,23 @@ jobs: python3 o11ybench/scripts/run_tracesbench_docker.py --stage cleanup \ --dataset "$DATASET" --runtime-cpus "$RUNTIME_CPUS" --runtime-memory "$RUNTIME_MEMORY" \ --root "$RUN_ROOT" --owner "$OWNER" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Clean up baseline containers + if: always() && steps.runtime.outcome == 'success' && inputs.compare_greptimedb + run: | + python3 o11ybench/scripts/run_tracesbench_docker.py --stage cleanup \ + --root "$RUN_ROOT/baseline" --owner "$OWNER-base" --runtime-image "$RESOLVED_RUNTIME_IMAGE" + - name: Benchmark summary + if: always() + run: python3 greptimedb/.github/scripts/tracesbench-summary.py --root "$RUN_ROOT" --targets "$TARGETS" --section benchmark >> "$GITHUB_STEP_SUMMARY" + - name: Load / lifecycle summary + if: always() + run: python3 greptimedb/.github/scripts/tracesbench-summary.py --root "$RUN_ROOT" --targets "$TARGETS" --section lifecycle >> "$GITHUB_STEP_SUMMARY" + - name: Baseline load / lifecycle summary + if: always() && inputs.compare_greptimedb + run: | + echo '## Baseline GreptimeDB' >> "$GITHUB_STEP_SUMMARY" + python3 greptimedb/.github/scripts/tracesbench-summary.py \ + --root "$RUN_ROOT/baseline" --targets greptimedb --section lifecycle >> "$GITHUB_STEP_SUMMARY" - name: Stage benchmark artifacts if: always() run: | @@ -269,11 +437,15 @@ jobs: if [[ -d "$METADATA_ROOT" ]]; then rsync -a --no-owner --no-group "$METADATA_ROOT" "$RUNNER_TEMP/tracesbench-artifacts/" fi - for entry in runs report report.json report.md; do + for entry in runs report report.json report.md version-comparison.json version-comparison.md; do if [[ -e "$RUN_ROOT/$entry" ]]; then rsync -a --no-owner --no-group "$RUN_ROOT/$entry" "$RUNNER_TEMP/tracesbench-artifacts/" fi done + if [[ -d "$RUN_ROOT/baseline/runs" ]]; then + mkdir -p "$RUNNER_TEMP/tracesbench-artifacts/baseline" + rsync -a --no-owner --no-group "$RUN_ROOT/baseline/runs" "$RUNNER_TEMP/tracesbench-artifacts/baseline/" + fi - name: Upload HTML and results (no corpus) if: always() uses: actions/upload-artifact@v4 diff --git a/.github/workflows/vmbench-long-range.yml b/.github/workflows/vmbench-long-range.yml index 357d2b21c69..aaf2dd7f437 100644 --- a/.github/workflows/vmbench-long-range.yml +++ b/.github/workflows/vmbench-long-range.yml @@ -20,8 +20,18 @@ on: - 1mo default: 1w required: true + compare_greptimedb: + description: Compare baseline and candidate sequentially on the same host (adds one full GreptimeDB run) + type: boolean + default: false + required: false + baseline_greptimedb_tag: + description: Baseline Docker Hub greptime/greptimedb tag; required when comparison is enabled + type: string + default: v1.2.1 + required: false greptimedb_tag: - description: Docker Hub greptime/greptimedb tag + description: Candidate Docker Hub greptime/greptimedb tag type: string default: latest required: true @@ -31,9 +41,9 @@ on: default: v1.114.0@sha256:227ad53fd57a4b4430624898a1be0cac39431e1235fa0dc92f567cce38969344 required: true runtime_image: - description: Aliyun runtime image (tag or digest reference) + description: 'Runtime image: auto = AWS ECR Public or Aliyun ACR by provider; explicit tag/digest overrides' type: string - default: greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:20260918-545dc570@sha256:0ff03e7e690575e875bb85b3054a03486dd9359073546dc715da177153cf7a29 + default: auto required: true provider: description: Cloud provider for the benchmark runner @@ -79,11 +89,13 @@ on: permissions: contents: read env: + COMPARE_GREPTIMEDB: ${{ inputs.compare_greptimedb }} + BASELINE_GREPTIMEDB_TAG: ${{ inputs.baseline_greptimedb_tag }} BENCHMARK_ROOT: ${{ inputs.benchmark_root }} RANGE_PRESET: ${{ inputs.range_preset }} GREPTIMEDB_TAG: ${{ inputs.greptimedb_tag }} VICTORIAMETRICS_TAG: ${{ inputs.victoriametrics_tag }} - RUNTIME_IMAGE: ${{ inputs.runtime_image }} + RUNTIME_IMAGE: ${{ inputs.runtime_image != 'auto' && inputs.runtime_image || inputs.provider == 'AWS' && 'public.ecr.aws/i8k6a5e1/greptime/o11ybenchruntime:b43785b56f42@sha256:e38811eb68caf3682ba56f5591a3e8e27e78e7983a923abb29fa2d7cd8fdbd46' || 'greptime-registry.cn-hangzhou.cr.aliyuncs.com/tools/o11ybench-runtime:b43785b56f42@sha256:e38811eb68caf3682ba56f5591a3e8e27e78e7983a923abb29fa2d7cd8fdbd46' }} PROVIDER: ${{ inputs.provider }} REQUESTED_INSTANCE_TYPE: ${{ inputs.instance_type }} SYSTEM_DISK_GIB: ${{ inputs.system_disk_gib }} @@ -107,6 +119,12 @@ jobs: with: path: greptimedb persist-credentials: false + - name: Validate version comparison + env: + RUN_GREPTIMEDB: ${{ inputs.run_greptimedb }} + run: | + python3 greptimedb/tests/perf/test_benchmark_version_comparison.py + python3 greptimedb/.github/scripts/benchmark-version-comparison.py validate - name: Resolve CI runner id: runner run: python3 greptimedb/.github/scripts/ci-runner-config.py --provider "$PROVIDER" --instance-type "$REQUESTED_INSTANCE_TYPE" @@ -135,8 +153,8 @@ jobs: ((BENCHMARK_TIMEOUT_MINUTES <= 360 && JANITOR_TTL_HOURS <= 168)) # Match the existing ECS workflow's provision, teardown and queue headroom. ((JANITOR_TTL_HOURS * 60 >= BENCHMARK_TIMEOUT_MINUTES + 120)) - [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* ]] - docker pull "$RUNTIME_IMAGE" + [[ "$RUNTIME_IMAGE" == *.cr.aliyuncs.com/* || "$RUNTIME_IMAGE" == public.ecr.aws/* ]] + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RUNTIME_IMAGE" metadata=$(docker image inspect "$RUNTIME_IMAGE") resolved=$(jq -er '.[0].RepoDigests[0]' <<< "$metadata") revision=$(jq -er '.[0].Config.Labels["org.opencontainers.image.revision"]' <<< "$metadata") @@ -146,6 +164,9 @@ jobs: args=(--workload vmbench-long-range --range-preset "$RANGE_PRESET" --db-cpus "$DB_CPUS" --db-memory "$DB_MEMORY" --run-classification controlled_server_benchmark --no-open-browser) + if [[ "$COMPARE_GREPTIMEDB" == true ]]; then + args+=(--target "greptimedb:name=greptimedb-baseline,image=greptime/greptimedb:$BASELINE_GREPTIMEDB_TAG,url=http://127.0.0.1:14000/v1/prometheus/api/v1/query_range") + fi # Tags become target-spec fields; reject separators before passing them to the planner. for target in greptimedb victoriametrics; do case "$target" in @@ -273,13 +294,13 @@ jobs: available=$(df -Pk "$path" | awk 'NR == 2 {print $4}') ((available >= 10 * 1024 * 1024)) || { echo "Need at least 10 GiB free on $path"; exit 1; } done - docker pull "$RESOLVED_RUNTIME_IMAGE" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE" docker image inspect "$RESOLVED_RUNTIME_IMAGE" > "$METADATA_ROOT/runtime-image.json" test "$(jq -r '.[0].Config.Labels["org.opencontainers.image.revision"]' "$METADATA_ROOT/runtime-image.json")" = "$RESOLVED_O11YBENCH_REF" plan="$METADATA_ROOT/warm.json" while IFS=$'\t' read -r target reference; do echo "Resolving $target: $reference" - docker pull "$reference" + bash greptimedb/.github/scripts/pull-benchmark-image.sh "$reference" docker image inspect "$reference" > "$METADATA_ROOT/$target-image.json" digest=$(jq -er '.[0].RepoDigests[0]' "$METADATA_ROOT/$target-image.json") jq --arg target "$target" --arg image "$digest" '.targets |= map(if .name == $target then .image = $image else . end)' "$plan" > "$plan.tmp" @@ -303,6 +324,10 @@ jobs: if: always() run: | docker ps -aq --filter "label=o11ybench.run-id=$RUN_ID-warm" | xargs -r docker rm -fv + - name: GreptimeDB version comparison + if: always() && inputs.compare_greptimedb + run: | + python3 greptimedb/.github/scripts/benchmark-version-comparison.py report --kind long-range --root "$ARTIFACT_ROOT-warm" >> "$GITHUB_STEP_SUMMARY" - name: Stage benchmark artifacts if: always() run: | diff --git a/tests/perf/test_agent_observability_summary.py b/tests/perf/test_agent_observability_summary.py index 8223231c908..a6c1774ba0f 100644 --- a/tests/perf/test_agent_observability_summary.py +++ b/tests/perf/test_agent_observability_summary.py @@ -14,11 +14,13 @@ import importlib.util import json +import sys import tempfile import unittest from pathlib import Path SCRIPT = Path(__file__).parents[2] / '.github/scripts/agent-observability-summary.py' +sys.path.insert(0, str(SCRIPT.parent)) spec = importlib.util.spec_from_file_location('observability_summary', SCRIPT) summary = importlib.util.module_from_spec(spec) spec.loader.exec_module(summary) diff --git a/tests/perf/test_benchmark_version_comparison.py b/tests/perf/test_benchmark_version_comparison.py new file mode 100644 index 00000000000..56c2619e81c --- /dev/null +++ b/tests/perf/test_benchmark_version_comparison.py @@ -0,0 +1,138 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest + +SCRIPTS = Path(__file__).resolve().parents[2] / '.github/scripts' +spec = importlib.util.spec_from_file_location('comparison', SCRIPTS / 'benchmark-version-comparison.py') +comparison = importlib.util.module_from_spec(spec) +spec.loader.exec_module(comparison) + + +def write(path, value): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value)) + + +class VersionComparison(unittest.TestCase): + def test_opt_in_validation(self): + comparison.validate({}) + comparison.validate({'COMPARE_GREPTIMEDB': 'false', 'GREPTIMEDB_TAG': 'invalid\n'}) + good = dict(COMPARE_GREPTIMEDB='true', RUN_GREPTIMEDB='true', + BASELINE_GREPTIMEDB_TAG='v1.2.1', GREPTIMEDB_TAG='v1.3.0-beta.1') + comparison.validate(good) + for key, value in [('RUN_GREPTIMEDB', 'false'), ('BASELINE_GREPTIMEDB_TAG', ''), + ('BASELINE_GREPTIMEDB_TAG', 'v1.3.0-beta.1'), + ('BASELINE_GREPTIMEDB_TAG', 'a,image=other'), ('GREPTIMEDB_TAG', 'x\ny'), + ('COMPARE_GREPTIMEDB', 'yes')]: + with self.subTest(key=key, value=value), self.assertRaises(ValueError): + comparison.validate(good | {key: value}) + + def test_improvement_direction(self): + self.assertAlmostEqual(comparison.improvement(100, 80), 20) + self.assertAlmostEqual(comparison.improvement(100, 120), -20) + self.assertAlmostEqual(comparison.improvement(100, 120, True), 20) + for old, new in [(0, 3), (None, 4), (1, -1), (1, float('nan')), (True, 2)]: + self.assertIsNone(comparison.improvement(old, new)) + + def test_agent_result_gates(self): + summary = dict(passed=True, contract_sha256='contract', data_profile='S', query_window=[1, 2], + warmup_runs=1, measured_runs=5, per_query=[dict(query_id='Q01', result_fingerprint='answer', + success=5, errors=0, timeouts=0, p50_ms=100, p95_ms=100, p99_ms=100)]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + base = root / 'baseline/greptimedb-1/greptimedb-1/measured-summary.json' + candidate = root / 'greptimedb-1/greptimedb-1/measured-summary.json' + for r in (root / 'baseline', root): + (r / 'greptimedb-1').mkdir(parents=True, exist_ok=True) + (r / 'greptimedb-1/exit-code.txt').write_text('0\n') + write(base, summary) + changed = copy.deepcopy(summary); changed['per_query'][0]['p50_ms'] = 80 + write(candidate, changed) + self.assertAlmostEqual(comparison.agent(root)[0]['improvement_percent'], 20) + for key, value in [('errors', 1), ('timeouts', 1), ('result_fingerprint', 'wrong'), ('success', 0)]: + broken = copy.deepcopy(changed); broken['per_query'][0][key] = value + write(candidate, broken) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + (root / 'greptimedb-1/exit-code.txt').write_text('1\n') + write(candidate, changed) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + (root / 'greptimedb-1/exit-code.txt').write_text('0\n') + changed['query_window'] = [2, 3]; write(candidate, changed) + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + candidate.unlink() + self.assertTrue(all(r['improvement_percent'] is None for r in comparison.agent(root))) + + def test_traces_uses_measured_only_and_rejects_failed_query(self): + answer = dict(query_id='T01', outcome='passed', fingerprint='answer', row_count=1, elapsed_ms=100) + sample = dict(answer, phase='measured') + record = dict(success=True, expected={'contracts': {'T01': {'fingerprint': 'answer', 'row_count': 1}}}, + protocol={'runs': 5}, audit_expected={}, corpus='sanity', + execution=dict(db_cpus='8', db_memory='16g', load_workers_requested=4, max_in_flight=8), + correctness=[answer], samples=[dict(sample, phase='warmup', elapsed_ms=99999), sample, sample]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp);base = root / 'baseline/runs/greptimedb/run.json'; candidate = root / 'runs/greptimedb/run.json' + write(base, record); changed = copy.deepcopy(record) + for s in changed['samples']: + if s['phase'] == 'measured': s['elapsed_ms'] = 80 + write(candidate, changed) + self.assertAlmostEqual(comparison.traces(root)[0]['improvement_percent'], 20) + self.assertIsNone(comparison.traces(root)[1]['improvement_percent']) + for key, value in [('outcome', 'failed'), ('fingerprint', 'wrong'), ('row_count', 2)]: + broken = copy.deepcopy(changed); broken['samples'][1][key] = value; write(candidate, broken) + self.assertIsNone(comparison.traces(root)[0]['improvement_percent']) + + def test_long_range_gates(self): + summary = dict(workload='vmbench-long-range', query_file_sha256='hash', start='start', end='end', step='30m', + profile='small', data_duration='32d', scrape_interval='30s', cache_policy_effective='warm', + mode='serial', connection_mode='force-close', request_timeout_seconds=120, tries=5, warmup_runs=1, + errors=0, timeouts=0, per_query=[dict(query_index=1, query='up', series=1, points=2, + success=5, errors=0, timeouts=0, p50_ms=100)]) + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp); base = root/'greptimedb-baseline/results-vmbench-long-range/serial-summary.json' + candidate = root/'greptimedb/results-vmbench-long-range/serial-summary.json' + health = {phase: {target: {'state': 'Running=true ExitCode=0 OOMKilled=false'} + for target in ('greptimedb-baseline', 'greptimedb')} + for phase in ('post_load_results', 'post_query_results')} + health['post_load_results'] = {k: {'docker': v} for k, v in health['post_load_results'].items()} + write(root / 'run-manifest.json', health) + write(base, summary); write(candidate, summary) + self.assertEqual(comparison.long_range(root)[0]['improvement_percent'], 0) + for key, value in [('start','different'),('tries',1),('query_file_sha256','other'),('errors',1)]: + write(candidate, summary | {key:value}) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + write(candidate, summary) + broken=copy.deepcopy(health);broken['post_query_results']['greptimedb']['state']='Running=false ExitCode=137 OOMKilled=true' + write(root / 'run-manifest.json', broken) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + write(root / 'run-manifest.json', health) + empty=copy.deepcopy(summary);empty['per_query'][0]['points']=0;write(candidate,empty) + self.assertIsNone(comparison.long_range(root)[0]['improvement_percent']) + + def test_missing_results_and_rendering(self): + with tempfile.TemporaryDirectory() as tmp: + for load in (comparison.agent, comparison.traces, comparison.long_range): + self.assertTrue(all(r['improvement_percent'] is None for r in load(Path(tmp)))) + rendered = comparison.render(dict(baseline_tag='old', candidate_tag='', rows=[])) + self.assertIn('<new|tag>', rendered) + self.assertIn('Results unavailable', rendered) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/perf/test_ci_runner_scripts.py b/tests/perf/test_ci_runner_scripts.py index f9742079b49..c703e41c62b 100644 --- a/tests/perf/test_ci_runner_scripts.py +++ b/tests/perf/test_ci_runner_scripts.py @@ -16,6 +16,7 @@ """Offline regression coverage for CI runner selection and lifecycle.""" import importlib.util +import json import os import re import subprocess @@ -77,6 +78,67 @@ class ProviderConfigTest(unittest.TestCase): with self.subTest(provider=provider, requested=requested): self.assertEqual(config.resolve_instance_type(provider, requested), expected) + def test_traces_workflow_runner_defaults_and_overrides(self): + workflow = (SCRIPTS.parent / "workflows/tracesbench.yml").read_text() + step = workflow.split(" - name: Resolve CI runner\n", 1)[1].split(" - name:", 1)[0] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "greptimedb").symlink_to(SCRIPTS.parents[1], target_is_directory=True) + output = root / "output" + for provider, requested, expected in ( + ("Aliyun", "auto", "ecs.g9i.2xlarge"), + ("AWS", "auto", "m7i.2xlarge"), + ("Aliyun", "ecs.c9i.4xlarge", "ecs.c9i.4xlarge"), + ("AWS", "c7i.4xlarge", "c7i.4xlarge"), + ("invalid", "auto", None), + ("AWS", "ecs.g9i.2xlarge", None), + ): + with self.subTest(provider=provider, requested=requested): + output.write_text("") + result = subprocess.run( + ["bash", "-c", command], cwd=root, capture_output=True, text=True, + env=dict(os.environ, PROVIDER=provider, REQUESTED_INSTANCE_TYPE=requested, + GITHUB_OUTPUT=str(output)), timeout=10, + ) + if expected is None: + self.assertNotEqual(result.returncode, 0) + self.assertEqual(output.read_text(), "") + else: + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(output.read_text(), f"instance_type={expected}\n") + self.assertRegex(workflow, r"(?s)db_cpus:.*?default: '8'") + self.assertRegex(workflow, r"(?s)db_memory:.*?default: 32g") + + def test_traces_selected_targets_and_image_tag(self): + workflow = (SCRIPTS.parent / "workflows/tracesbench.yml").read_text() + step = workflow.split(" - name: Validate selected targets\n", 1)[1].split(" - name:", 1)[0] + command = textwrap.dedent(step.split(" run: |\n", 1)[1]) + with tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "output" + for enabled, tag, expected in ( + (("true", "false", "false"), "latest", "greptimedb"), + (("false", "true", "false"), "", "victoriatraces"), + (("false", "false", "true"), "", "tempo"), + (("true", "true", "true"), "v1.3.0-beta.1", "greptimedb,victoriatraces,tempo"), + (("false", "false", "false"), "latest", None), + (("true", "false", "false"), "latest; touch /tmp/unexpected", None), + (("true", "false", "false"), "", None), + (("invalid", "false", "false"), "latest", None), + ): + with self.subTest(enabled=enabled, tag=tag): + output.write_text("") + env = dict(os.environ, GREPTIMEDB_TAG=tag, GITHUB_OUTPUT=str(output)) + env.update(zip(("greptimedb", "victoriatraces", "tempo"), enabled)) + result = subprocess.run(["bash", "-c", command], env=env, + capture_output=True, text=True, timeout=5) + if expected is None: + self.assertNotEqual(result.returncode, 0) + self.assertNotIn("targets=", output.read_text()) + else: + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn(f"targets={expected}\n", output.read_text()) + def test_wrong_provider_and_shell_fragments_rejected(self): for provider, requested in (("ec2", "auto"), ("AWS", "ecs.c9i.2xlarge"), ("Aliyun", "c7i.2xlarge"), ("AWS", ""), @@ -401,6 +463,154 @@ class AwsTeardownTest(unittest.TestCase): self.assertEqual(teardown.teardown(Mock(), "owner/repo", "token", run_id="123-2"), 1) self.assertEqual([c.args[1] for c in delete.call_args_list], ["i-a", "i-b"]) +class BenchmarkImagePullTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.root = Path(self.tmp.name) + self.env = dict(os.environ, PATH=f"{self.root}:{os.environ['PATH']}", + PULL_TEST_ROOT=str(self.root)) + self.tool("docker", ''' + echo "$*" >> "$PULL_TEST_ROOT/docker.log" + count=$(wc -l < "$PULL_TEST_ROOT/docker.log") + codes=${PULL_TEST_CODES:-0} + code=$(printf '%s' "$codes" | cut -d, -f"$count") + exit "${code:-0}" + ''') + self.tool("timeout", ''' + echo "$*" >> "$PULL_TEST_ROOT/timeout.log" + shift 2 + if [[ ${PULL_TEST_TIMEOUT:-0} == 1 ]]; then exit 124; fi + exec "$@" + ''') + self.tool("sleep", 'echo "$*" >> "$PULL_TEST_ROOT/sleep.log"') + + def tool(self, name, body): + p = self.root / name + p.write_text("#!/bin/bash\nset -eu\n" + textwrap.dedent(body)) + p.chmod(0o755) + + def lines(self, name): + p = self.root / (name + ".log") + return p.read_text().splitlines() if p.exists() else [] + + def pull(self, codes, *args): + self.env["PULL_TEST_CODES"] = codes + return subprocess.run(["bash", str(SCRIPTS / "pull-benchmark-image.sh"), *args], + env=self.env, capture_output=True, text=True, timeout=5) + + def test_success_does_not_retry(self): + result = self.pull("0", "registry/test@sha256:" + "a" * 64) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(len(self.lines("docker")), 1) + self.assertEqual(self.lines("sleep"), []) + self.assertTrue(self.lines("timeout")[0].startswith("--kill-after=10s 180s docker pull ")) + + def test_transient_failure_then_success(self): + result = self.pull("1,1,0", "registry/test:tag") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(self.lines("docker"), ["pull registry/test:tag"] * 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + self.assertIn("attempt 3/3", result.stderr) + + def test_exhaustion_preserves_failure(self): + result = self.pull("1,1,7", "registry/test:tag") + self.assertEqual(result.returncode, 7) + self.assertEqual(len(self.lines("docker")), 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + + def test_timeout_is_bounded_and_fatal(self): + self.env["PULL_TEST_TIMEOUT"] = "1" + result = self.pull("0", "registry/test:tag") + self.assertEqual(result.returncode, 124) + self.assertEqual(len(self.lines("timeout")), 3) + self.assertEqual(self.lines("sleep"), ["10", "20"]) + + def test_bad_arguments_do_not_call_docker(self): + for args in ([], [""], ["--help"], ["a", "b"]): + with self.subTest(args=args): + self.assertEqual(self.pull("0", *args).returncode, 2) + self.assertEqual(self.lines("docker"), []) + + def test_traces_prefetch_executes_catalog_refs_and_stops_on_failure(self): + traces = (SCRIPTS.parent / 'workflows/tracesbench.yml').read_text() + start = traces.index(' - name: Pull database images before generating data') + end = traces.index(' - name: Generate selected dataset once') + command = textwrap.dedent(traces[start:end].split(' run: |\n', 1)[1]) + helper = self.root / 'greptimedb/.github/scripts/pull-benchmark-image.sh' + helper.parent.mkdir(parents=True) + helper.write_text((SCRIPTS / helper.name).read_text()) + catalogs = self.root / 'o11ybench/workloads/tracesbench/catalogs' + catalogs.mkdir(parents=True) + images = [f'registry/{name}@sha256:' + 'a' * 64 + for name in ('greptimedb', 'victoriatraces', 'tempo')] + for name, image in zip(('greptimedb', 'victoriatraces', 'tempo'), images): + (catalogs / f'{name}.json').write_text(json.dumps({'target_identity': {'image': image}})) + self.tool('docker', ''' + echo "$*" >> "$PULL_TEST_ROOT/docker.log" + if [[ $1 == pull && $2 == ${PULL_FAIL_IMAGE:-} ]]; then exit 1; fi + if [[ $1 == image ]]; then printf '[{"RepoDigests":["greptime/greptimedb@sha256:%s"]}]\\n' "$(printf 'a%.0s' {1..64})"; fi + ''') + images[0] = 'greptime/greptimedb:v1.3.0-beta.1' + self.env.update(METADATA_ROOT=str(self.root), GREPTIMEDB_TAG='v1.3.0-beta.1', + GITHUB_ENV=str(self.root / 'env'), BASELINE_GREPTIMEDB_TAG='v1.2.1') + for selected, compare in ((('greptimedb',), False), (('tempo',), False), + (('greptimedb', 'victoriatraces', 'tempo'), False), + (('greptimedb',), True), (('greptimedb', 'victoriatraces', 'tempo'), True)): + with self.subTest(selected=selected): + self.env['TARGETS'] = ','.join(selected) + self.env['COMPARE_GREPTIMEDB'] = str(compare).lower() + (self.root / 'env').write_text('') + (self.root / 'docker.log').unlink(missing_ok=True) + result = subprocess.run(['bash', '-c', command], cwd=self.root, env=self.env, + capture_output=True, text=True, timeout=5) + self.assertEqual(result.returncode, 0, result.stderr) + selected_images = [image for target, image in zip(('greptimedb', 'victoriatraces', 'tempo'), images) + if target in selected] + if compare: + selected_images.append('greptime/greptimedb:v1.2.1') + self.assertEqual(self.lines('docker'), + [call for image in selected_images for call in (f'pull {image}', f'image inspect {image}')]) + expected_env = 'GREPTIMEDB_IMAGE=greptime/greptimedb@sha256:' + 'a' * 64 + '\n' + expected_env = expected_env if 'greptimedb' in selected else '' + if compare: + expected_env += 'BASELINE_GREPTIMEDB_IMAGE=greptime/greptimedb@sha256:' + 'a' * 64 + '\n' + self.assertEqual((self.root / 'env').read_text(), expected_env) + (self.root / 'docker.log').unlink() + self.env['PULL_FAIL_IMAGE'] = images[1] + result = subprocess.run(['bash', '-c', command], cwd=self.root, env=self.env, + capture_output=True, text=True, timeout=5) + self.assertNotEqual(result.returncode, 0) + self.assertEqual(self.lines('docker'), [f'pull {images[0]}', f'image inspect {images[0]}'] + + [f'pull {images[1]}'] * 3) + self.assertNotIn(f'pull {images[2]}', self.lines('docker')) + + def test_workflows_pull_before_generation(self): + workflows = SCRIPTS.parent / "workflows" + agent = (workflows / "agent-observability.yml").read_text() + traces = (workflows / "tracesbench.yml").read_text() + long_range = (workflows / "vmbench-long-range.yml").read_text() + for content in (agent, traces, long_range): + self.assertNotIn("docker pull ", content) + for reference in ("RUNTIME_IMAGE", "RESOLVED_RUNTIME_IMAGE", "reference"): + self.assertIn(f'pull-benchmark-image.sh "${reference}"', long_range) + self.assertLess(long_range.index('pull-benchmark-image.sh "$reference"'), + long_range.index('- name: Warm serial and concurrent benchmarks')) + self.assertIn('pull-benchmark-image.sh "$reference"', agent) + self.assertLess(agent.index('resolve_image runtime'), agent.index('- name: Generate dataset once')) + self.assertIn('pull-benchmark-image.sh "$RUNTIME_IMAGE"', traces) + self.assertIn('pull-benchmark-image.sh "$RESOLVED_RUNTIME_IMAGE"', traces) + start = traces.index(' - name: Pull database images before generating data') + end = traces.index(' - name: Generate selected dataset once') + self.assertLess(start, end) + step = traces[start:end] + self.assertIn('IFS=, read -ra targets <<< "$TARGETS"', step) + self.assertIn('for target in "${targets[@]}"', step) + self.assertIn('o11ybench/workloads/tracesbench/catalogs/$target.json', step) + self.assertIn('pull-benchmark-image.sh "$image"', step) + self.assertIn('set -euo pipefail', step) + self.assertNotIn('continue-on-error', step) + if __name__ == "__main__": unittest.main() diff --git a/tests/perf/test_tracesbench_summary.py b/tests/perf/test_tracesbench_summary.py new file mode 100644 index 00000000000..cb7c4d8e06e --- /dev/null +++ b/tests/perf/test_tracesbench_summary.py @@ -0,0 +1,54 @@ +# Copyright 2023 Greptime Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPT = Path(__file__).parents[2] / '.github/scripts/tracesbench-summary.py' +sys.path.insert(0, str(SCRIPT.parent)) +spec = importlib.util.spec_from_file_location('tracesbench_summary', SCRIPT) +summary = importlib.util.module_from_spec(spec) +spec.loader.exec_module(summary) + + +class LifecycleTest(unittest.TestCase): + def test_ingestion_layouts(self): + metrics = {'wall_seconds': 0.121723417, 'successful_requests': 17, 'failed_requests': 0} + cases = [ + (target, {'ingest': metrics}, ['121.7 ms', '17', '0']) + for target in summary.TARGETS + ] + [ + ('tempo', {'transport': {'ingest': metrics}}, ['121.7 ms', '17', '0']), + ('tempo', {'ingest': {'wall_seconds': 0, 'successful_requests': 0, 'failed_requests': 0}, + 'transport': {'ingest': metrics}}, ['0 µs', '0', '0']), + ('tempo', {'ingest': {}, 'transport': {'ingest': metrics}}, ['—', '—', '—']), + ('tempo', {'transport': {}}, ['—', '—', '—']), + ] + [(target, {}, ['—', '—', '—']) for target in summary.TARGETS] + [ + (target, {'transport': {'ingest': metrics}}, ['—', '—', '—']) + for target in ('greptimedb', 'victoriatraces') + ] + with tempfile.TemporaryDirectory() as tmp: + for target, load, expected in cases: + with self.subTest(target=target, load=load): + text = summary.lifecycle({}, [{'target': target, 'load': load}], Path(tmp), [target]) + row = next(line for line in text.splitlines() if line.startswith('| ' + target + ' |')) + cells = [cell.strip() for cell in row.split('|')[1:-1]] + self.assertEqual(cells[4:7], expected) + + +if __name__ == '__main__': + unittest.main() From a907cbfbe4197512d568c811c3daed99b19e436d Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:45:43 +0000 Subject: [PATCH 128/133] perf: cache OTLP trace service and operation writes (#9432) * perf: cache OTLP trace service and operation writes Signed-off-by: shuiyisong * chore: fix CR issues for tests Signed-off-by: shuiyisong * fix: bound asynchronous trace auxiliary writes Signed-off-by: shuiyisong * perf: bound trace auxiliary cache by estimated memory Signed-off-by: shuiyisong * chore: size trace auxiliary cache from host or pod memory Signed-off-by: shuiyisong * chore: add metrics for cache hit and miss Signed-off-by: shuiyisong * fix: test Signed-off-by: shuiyisong * chore: unify frontend cache metrics Signed-off-by: shuiyisong * test: merge and reduce Signed-off-by: shuiyisong * chore: add init for the metrics Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- Cargo.lock | 6 + config/config.md | 2 + config/frontend.example.toml | 7 + config/standalone.example.toml | 7 + src/cmd/tests/load_config_test.rs | 23 + src/common/frontend/Cargo.toml | 2 + src/common/frontend/src/lib.rs | 1 + src/common/frontend/src/metrics.rs | 117 ++++ src/common/query/src/lib.rs | 39 +- src/frontend/AGENTS.md | 14 + src/frontend/Cargo.toml | 1 + src/frontend/src/instance.rs | 34 +- src/frontend/src/instance/builder.rs | 17 +- src/frontend/src/instance/otlp.rs | 1 + src/frontend/src/instance/otlp/README.md | 52 ++ .../src/instance/otlp/trace_ingest.rs | 346 ++++++++++- .../src/instance/otlp/trace_ingest/tests.rs | 250 +++++++- src/frontend/src/metrics.rs | 7 + src/frontend/src/service_config/otlp.rs | 44 ++ src/operator/AGENTS.md | 2 + src/operator/src/batcher.rs | 4 +- src/operator/src/insert.rs | 15 +- src/pipeline/Cargo.toml | 1 + src/pipeline/src/manager/pipeline_cache.rs | 218 +++++-- src/pipeline/src/manager/pipeline_operator.rs | 5 +- src/servers/AGENTS.md | 2 + src/servers/Cargo.toml | 1 + src/servers/src/batcher/table.rs | 74 ++- src/servers/src/mysql/handler.rs | 23 +- tests-integration/Cargo.toml | 1 + tests-integration/src/otlp.rs | 577 +++++++++++++++++- tests-integration/tests/http.rs | 1 + 32 files changed, 1784 insertions(+), 110 deletions(-) create mode 100644 src/common/frontend/src/metrics.rs diff --git a/Cargo.lock b/Cargo.lock index c02c7fdcc60..9b447169fbc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2459,7 +2459,9 @@ dependencies = [ "common-macro", "common-meta", "greptime-proto", + "lazy_static", "meta-client", + "prometheus 0.14.0", "serde", "serde_json", "snafu 0.8.6", @@ -5604,6 +5606,7 @@ dependencies = [ "lazy_static", "log-query", "meta-client", + "moka", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", @@ -10860,6 +10863,7 @@ dependencies = [ "chrono-tz", "common-catalog", "common-error", + "common-frontend", "common-macro", "common-meta", "common-query", @@ -13522,6 +13526,7 @@ dependencies = [ "common-catalog", "common-decimal", "common-error", + "common-frontend", "common-grpc", "common-macro", "common-mem-prof", @@ -15122,6 +15127,7 @@ dependencies = [ "store-api", "substrait 1.4.0-alpha.0", "table", + "temp-env", "tempfile", "time", "tokio", diff --git a/config/config.md b/config/config.md index c538559e7cf..612a5cd43c0 100644 --- a/config/config.md +++ b/config/config.md @@ -95,6 +95,7 @@ | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | @@ -368,6 +369,7 @@ | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.trace_aux_cache_size` | String | Auto | Estimated memory budget for cached trace service/operation keys per frontend,
shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching.
Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB.
Uses 32MiB if memory detection is unavailable.
Excludes cache and allocator overhead; shared table names are charged per entry. | | `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 21e7bbb40f3..e4ef6e36935 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -310,6 +310,13 @@ enable = true enable = true ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 7e1f8dcf6b0..c656049c5d2 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -282,6 +282,13 @@ enable = true enable = true ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Estimated memory budget for cached trace service/operation keys per frontend, +## shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. +## Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. +## Uses 32MiB if memory detection is unavailable. +## Excludes cache and allocator overhead; shared table names are charged per entry. +## @toml2docs:none-default="Auto" +#+ trace_aux_cache_size = "32MiB" ## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric ## resource attributes, so metrics-only services reach the semantic graph. experimental_enable_resource_info = false diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 1f4c69c8938..00cdafd9b8a 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -731,6 +731,29 @@ fn test_example_configs_document_actual_defaults() { assert!(errors.is_empty(), "{}", errors.join("\n")); } +#[test] +fn test_load_trace_aux_cache_size() { + let default = FrontendOptions::default().otlp.trace_aux_cache_size; + for (setting, size) in [ + ("", default), + ("trace_aux_cache_size = 0", ReadableSize(0)), + ("trace_aux_cache_size = \"1MiB\"", ReadableSize::mb(1)), + ] { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(config.path(), format!("[otlp]\n{setting}\n")).unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + for options in [frontend.component, standalone.component.frontend_options()] { + assert_eq!(options.otlp.trace_aux_cache_size, size); + } + } +} + #[test] fn test_load_removed_histogram_options() { for enabled in [false, true] { diff --git a/src/common/frontend/Cargo.toml b/src/common/frontend/Cargo.toml index 26092ae13f7..84a8373dcbb 100644 --- a/src/common/frontend/Cargo.toml +++ b/src/common/frontend/Cargo.toml @@ -13,7 +13,9 @@ common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true greptime-proto.workspace = true +lazy_static.workspace = true meta-client.workspace = true +prometheus.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true diff --git a/src/common/frontend/src/lib.rs b/src/common/frontend/src/lib.rs index dacffd19968..b6ec72da1a3 100644 --- a/src/common/frontend/src/lib.rs +++ b/src/common/frontend/src/lib.rs @@ -18,6 +18,7 @@ use std::str::FromStr; use snafu::OptionExt; pub mod error; +pub mod metrics; pub mod selector; pub mod slow_query_event; diff --git a/src/common/frontend/src/metrics.rs b/src/common/frontend/src/metrics.rs new file mode 100644 index 00000000000..7a2c152fa1f --- /dev/null +++ b/src/common/frontend/src/metrics.rs @@ -0,0 +1,117 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Cache lookup counters shared by frontend, pipeline, and protocol handlers. +//! +//! The `type` label identifies `otlp_trace_aux`, `otlp_metrics_legacy`, +//! `pipeline_table`, `pipeline`, `pipeline_original`, `pipeline_failover`, or +//! `mysql_prepared_stmt`. Counters aggregate across catalogs and connections. +//! Trace lookups count distinct service/operation keys once per request, excluding +//! rechecks; legacy metric lookups count each requested metric name. Pipeline +//! loaders count successful coalesced waiters as hits and failed loads as misses. +//! Failover counts one outcome after resolving all candidate schemas; an ambiguous +//! result is a miss. Table and prepared-statement caches count each lookup. +//! Both outcomes for all known caches are exported from the first metric update. + +use lazy_static::lazy_static; +use prometheus::{IntCounterVec, register_int_counter_vec}; + +const CACHE_TYPES: [&str; 7] = [ + "otlp_trace_aux", + "otlp_metrics_legacy", + "pipeline_table", + "pipeline", + "pipeline_original", + "pipeline_failover", + "mysql_prepared_stmt", +]; + +lazy_static! { + static ref CACHE_COUNTERS: (IntCounterVec, IntCounterVec) = { + let hits = register_int_counter_vec!( + "greptime_frontend_cache_hit", + "frontend cache hit", + &["type"] + ) + .unwrap(); + let misses = register_int_counter_vec!( + "greptime_frontend_cache_miss", + "frontend cache miss", + &["type"] + ) + .unwrap(); + // Register both outcomes together so an all-hit or all-miss cache has both series. + for cache_type in CACHE_TYPES { + hits.with_label_values(&[cache_type]); + misses.with_label_values(&[cache_type]); + } + (hits, misses) + }; + + /// Frontend cache hits by cache type. + pub static ref CACHE_HIT: IntCounterVec = CACHE_COUNTERS.0.clone(); + /// Frontend cache misses by cache type. + pub static ref CACHE_MISS: IntCounterVec = CACHE_COUNTERS.1.clone(); +} + +/// Records one completed cache lookup. +pub fn record_cache_lookup(cache_type: &str, hit: bool) { + if hit { + CACHE_HIT.with_label_values(&[cache_type]).inc(); + } else { + CACHE_MISS.with_label_values(&[cache_type]).inc(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cache_lookup_counters() { + record_cache_lookup("pipeline", false); + // Scrape before accessing CACHE_HIT, which would initialize it and hide a missing series. + let families = prometheus::gather(); + for (name, expected) in [ + ("greptime_frontend_cache_hit", 0.0), + ("greptime_frontend_cache_miss", 1.0), + ] { + let family = families + .iter() + .find(|family| family.name() == name) + .unwrap(); + assert_eq!(family.get_metric().len(), CACHE_TYPES.len()); + for cache_type in CACHE_TYPES { + let metric = family + .get_metric() + .iter() + .find(|metric| metric.get_label()[0].value() == cache_type) + .unwrap(); + assert_eq!(metric.get_label()[0].name(), "type"); + assert_eq!( + metric.get_counter().value(), + if cache_type == "pipeline" { + expected + } else { + 0.0 + } + ); + } + } + record_cache_lookup("pipeline", true); + record_cache_lookup("pipeline", false); + assert_eq!(CACHE_HIT.with_label_values(&["pipeline"]).get(), 1); + assert_eq!(CACHE_MISS.with_label_values(&["pipeline"]).get(), 2); + } +} diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs index 8ce028ef9fe..8151dfc5731 100644 --- a/src/common/query/src/lib.rs +++ b/src/common/query/src/lib.rs @@ -25,10 +25,13 @@ pub mod stream; pub mod test_util; use std::fmt::{Debug, Display, Formatter}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use api::greptime_proto::v1::AddColumnLocation as Location; use api::greptime_proto::v1::add_column_location::LocationType; +use common_error::ext::BoxedError; use common_recordbatch::{ RecordBatches, SendableRecordBatchMapper, SendableRecordBatchStream, map_dictionary_to_values, map_dictionary_to_values_schema, @@ -74,6 +77,30 @@ pub struct OutputMeta { /// May exist for query output. One can retrieve execution metrics from this plan. pub plan: Option>, pub cost: OutputCost, + /// Process-local storage results attached by asynchronous ordinary-table batching. + /// Callers needing confirmed writes must await these. Dropping them does not cancel writes. + pub write_completions: Vec, +} + +/// Actual storage completion of an already accepted asynchronous write. +pub struct WriteCompletion(Pin> + Send>>); + +impl WriteCompletion { + /// Wraps the storage result without changing how the write is scheduled. + pub fn new(completion: impl Future> + Send + 'static) -> Self { + Self(Box::pin(completion)) + } + + /// Waits for storage success or failure, independently of queue admission. + pub async fn wait(self) -> Result<(), BoxedError> { + self.0.await + } +} + +impl Debug for WriteCompletion { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WriteCompletion").finish_non_exhaustive() + } } impl Output { @@ -163,18 +190,26 @@ impl Debug for OutputData { impl OutputMeta { pub fn new(plan: Option>, cost: usize) -> Self { - Self { plan, cost } + Self { + plan, + cost, + ..Default::default() + } } pub fn new_with_plan(plan: Arc) -> Self { Self { plan: Some(plan), cost: 0, + ..Default::default() } } pub fn new_with_cost(cost: usize) -> Self { - Self { plan: None, cost } + Self { + cost, + ..Default::default() + } } } diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index bab27676655..fdaa5dbcfe4 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -34,6 +34,9 @@ remote datanodes via `operator`/`client`. | `heartbeat` | `src/frontend/src/heartbeat.rs` | Heartbeat to metasrv; handles suspend / cache invalidation | | `service_config` | `src/frontend/src/service_config/` | Per-protocol option structs | +Shared frontend cache counters live in `src/common/frontend/src/metrics.rs`, +including caches owned by pipeline and protocol handlers. + ## Request lifecycles - **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, @@ -58,6 +61,17 @@ remote datanodes via `operator`/`client`. the lazy schema handshake; checks permissions and reconciles missing columns through `Inserter` once per stream, then reuses the refreshed table. +- **OTLP trace lookup tables** (`instance/otlp/trace_ingest.rs`): `Instance` owns + a shared cache of confirmed service/operation writes, which bypass batching. + `otlp.trace_aux_cache_size` sets its estimated byte budget at startup; zero disables it. + Async main writes retain completion results; a frontend task writes auxiliary + rows only for confirmed chunks before populating the cache. + Cold async trace requests reserve a shared auxiliary slot before main writes, + bounded separately by the table batcher's `max_inflight_requests`; confirmed + cache hits bypass this admission. The slot lasts through auxiliary completion. + These tables are ingestion-managed; manual mutation requires restarting + serving frontends. + - **Logical-table batching** (`instance/logical_batcher.rs`): `Services` initializes one shared batcher for opted-in HTTP Prom and nonlegacy OTLP metric-engine writes. Downstream routers can initialize it when enabling replacement endpoints. diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 1aa057d0fcd..92803933a8b 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -58,6 +58,7 @@ humantime-serde.workspace = true lazy_static.workspace = true log-query.workspace = true meta-client.workspace = true +moka = { workspace = true, features = ["sync"] } opentelemetry-proto.workspace = true opentelemetry-semantic-conventions = { version = "0.31", features = ["semconv_experimental"] } operator.workspace = true diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index d3e6f2addfc..b91290d2a55 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -48,8 +48,10 @@ use catalog::process_manager::{ use client::OutputData; use common_base::Plugins; use common_base::cancellation::CancellableFuture; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_event_recorder::EventRecorderRef; +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::key::TableMetadataManagerRef; @@ -141,6 +143,8 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + trace_aux_cache: otlp::TraceAuxCache, + trace_aux_limiter: Option, otlp_resource_info: bool, suspend: Arc, @@ -148,7 +152,7 @@ pub struct Instance { // first layer key: db-string // key: direct input metric name // value: if runs in legacy mode - otlp_metrics_table_legacy_cache: DashMap>, + otlp_metrics_table_legacy_cache: Arc>>, } impl Instance { @@ -570,6 +574,12 @@ fn fast_legacy_check( .iter() .filter_map(|name| cache.get(name)) .collect::>(); + CACHE_HIT + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by(hit_cache.len() as u64); + CACHE_MISS + .with_label_values(&["otlp_metrics_legacy"]) + .inc_by((names.len() - hit_cache.len()) as u64); if !hit_cache.is_empty() { let hit_legacy = hit_cache.iter().any(|en| *en.value()); let hit_prom = hit_cache.iter().any(|en| !*en.value()); @@ -3418,12 +3428,16 @@ mod tests { } #[tokio::test] - async fn test_event_recorder_is_exposed() -> TestResult<()> { + async fn test_instance_shared_state() -> TestResult<()> { let instance = test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) .await?; let _event_recorder = instance.event_recorder(); + assert!(Arc::ptr_eq( + &instance.otlp_metrics_table_legacy_cache, + &instance.clone().otlp_metrics_table_legacy_cache, + )); Ok(()) } @@ -3554,11 +3568,20 @@ mod tests { })?; let ctx = test_query_ctx(1); let handler: PipelineHandlerRef = Arc::new(instance.clone()); + let hits = CACHE_HIT.with_label_values(&["pipeline_table"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_table"]); + let before = (hits.get(), misses.get()); handler .get_pipeline("pipeline", None, ctx.clone()) .await .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); + handler + .get_pipeline("pipeline", None, ctx.clone()) + .await + .unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); assert_permission_denied( PipelineHandler::get_pipeline_str(&instance, "pipeline", None, ctx.clone()).await, ); @@ -3877,11 +3900,15 @@ mod tests { #[test] fn test_fast_legacy_check_is_read_only() { + let hits = CACHE_HIT.with_label_values(&["otlp_metrics_legacy"]); + let misses = CACHE_MISS.with_label_values(&["otlp_metrics_legacy"]); + let before = (hits.get(), misses.get()); let cache = DashMap::new(); cache.insert("metric1".to_string(), true); let names = vec!["metric1".to_string(), "metric2".to_string()]; assert_eq!(Some(true), fast_legacy_check(&cache, &names).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 1)); assert!(!cache.contains_key("metric2")); cache_legacy_mode(&cache, &names, true).unwrap(); @@ -3893,6 +3920,9 @@ mod tests { cache_incompatible.insert("metric1".to_string(), true); cache_incompatible.insert("metric2".to_string(), false); assert!(fast_legacy_check(&cache_incompatible, &names).is_err()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); + assert_eq!(None, fast_legacy_check(&cache, &[]).unwrap()); + assert_eq!((hits.get(), misses.get()), (before.0 + 3, before.1 + 1)); } #[test] diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index d261daed48e..032722499f9 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -22,6 +22,7 @@ use catalog::kvbackend::KvBackendCatalogManager; use catalog::process_manager::ProcessManagerRef; use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; +use common_batcher::request_limiter::RequestLimiter; use common_datasource::object_store::LocalFileAccess; use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableFlownodeSetCacheRef, TableRouteCacheRef}; @@ -50,8 +51,8 @@ use partition::manager::PartitionRuleManager; use pipeline::pipeline_operator::PipelineOperator; use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; -use servers::batcher::BatchingProtocol; use servers::batcher::table::TablePendingRowsBatcher; +use servers::batcher::{BatchingProtocol, pending_rows_batch_sync_enabled}; use snafu::{OptionExt, ResultExt}; use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; @@ -60,6 +61,7 @@ use crate::frontend::FrontendOptions; use crate::heartbeat::frontend_peer_addr; use crate::instance::Instance; use crate::instance::entity_graph::EntityGraphProviderImpl; +use crate::instance::otlp::TraceAuxCache; use crate::instance::region_query::FrontendRegionQueryHandler; use crate::service_config::BatcherOptions; @@ -385,6 +387,15 @@ impl FrontendBuilder { )); admin_event_recorder.install(&event_recorder); + let batcher_options = self.options.table_batcher_options(); + let trace_aux_limiter = if batcher_options.pending_rows_batching_enabled() + && !pending_rows_batch_sync_enabled() + { + RequestLimiter::try_new(batcher_options.max_inflight_requests) + } else { + None + }; + Ok(Instance { logical_batcher: Default::default(), frontend_peer_addr, @@ -400,10 +411,12 @@ impl FrontendBuilder { event_recorder, slow_query_recorder, process_manager, - otlp_metrics_table_legacy_cache: DashMap::new(), + otlp_metrics_table_legacy_cache: Arc::new(DashMap::new()), slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + trace_aux_cache: TraceAuxCache::new(self.options.otlp.trace_aux_cache_size.as_bytes()), + trace_aux_limiter, otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 5d96138319a..e4f395db457 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -52,6 +52,7 @@ use table::requests::{ SOURCE_OPENTELEMETRY, }; +pub(crate) use self::trace_ingest::TraceAuxCache; use self::trace_ingest::trace_conventions; use crate::instance::Instance; use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index c123f1caae1..e308ebec184 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -175,6 +175,58 @@ Accounting follows the main-table write: accepted or rejected span counts; and - failure details are bounded before they are folded into `TraceIngestOutcome`. +A cache shared by each frontend's `Instance` clones skips previously confirmed +auxiliary writes for v0, v1, and v2. `otlp.trace_aux_cache_size` sets an estimated +memory budget across all catalogs, schemas, and trace tables on each frontend. +By default, the budget is 1/128 of the pod's cgroup memory limit or total host +memory, with a minimum of `32MiB`. If memory detection is unavailable, it uses +`32MiB`. An explicit size overrides this default; `0` disables caching. +The estimate includes key structs and allocated string capacity, charging shared +table names per entry. It excludes cache and allocator overhead, so it is not a +strict process-memory limit. Changes take effect on frontend or standalone +restart. There is no time-based expiry. +Services and operations are cached independently. Only a successful +auxiliary write populates the cache. Auxiliary writes always bypass batching, +including when main spans use asynchronous batching. Failures, eviction, +frontend restarts, and concurrent misses can cause repeat writes. Write cost +reflects only the writes actually performed. + +`greptime_frontend_cache_hit{type="otlp_trace_aux"}` and +`greptime_frontend_cache_miss{type="otlp_trace_aux"}` count distinct service and +operation keys at the first cache check in each request. +Requests with auxiliary admission count before acquiring a slot; other requests +count when preparing auxiliary writes. Admission and write rechecks do not count. +When the cache is disabled, all checked keys count as misses. + +Successful auxiliary writes with `skip_wal` enabled are cached too. If a +datanode crashes before flushing those rows while the frontend survives, the +cache can suppress their recreation, even by later WAL-enabled requests, until +eviction or frontend restart. Missing service/operation discovery rows in this +case are an accepted limitation; the cache does not guarantee durability beyond +the write's WAL policy. + +With asynchronous main-table batching, the response still acknowledges queue +admission. A frontend task waits for the storage results and writes auxiliary +rows only for successful chunks. Failed chunks do not populate the auxiliary +tables or cache. Deferred auxiliary failures are logged; they cannot change an +already returned response. Deferred writes retain the original request's row +admission, and their write cost is not included in the early response. + +Async trace requests with cache misses acquire one shared auxiliary-work slot +before submitting main chunks. Each frontend limits these requests separately +from main-table batching, using `pending_rows_batcher.max_inflight_requests`. +The slot is held through auxiliary completion, including failures. Requests +whose service/operation entries are all cached bypass this admission and need +no deferred task. Cache lookups and updates hold no lock across storage writes. + +The auxiliary tables are ingestion-managed. When caching is enabled, manual +`DROP`, `TRUNCATE`, `DELETE`, or other mutations require clearing the affected +caches or restarting all serving frontends before relying on ingestion to +populate the lookup rows again. +There is no automatic invalidation or public cache-clear command. In v0, fewer +duplicate rows are appended; distinct service/operation lookup results remain +the same. + A rejection detail carries the failing cause, not just a status code, so an unusable attribute value names its column, source value, and target type. Repeats of the same `(site, cause)` collapse into one entry with an occurrence count, and diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index e5b16c7d191..8b90a19ba49 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, HashMap, HashSet}; +use std::mem; use std::sync::Arc; use api::helper::ColumnDataTypeWrapper; @@ -22,13 +23,17 @@ use api::v1::{ RowInsertRequest, RowInsertRequests, Rows, Value, }; use client::Output; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_meta::rpc::ddl::TriggerReason; +use common_query::WriteCompletion; use common_telemetry::{debug, warn}; use datatypes::prelude::ConcreteDataType; +use moka::sync::Cache; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::PipelineWay; +use prometheus::IntCounter; use servers::error::{self, Result as ServerResult}; use servers::otlp; use servers::otlp::coerce::{coerce_value_data, is_supported_trace_coercion, trace_value_datatype}; @@ -43,6 +48,8 @@ use table::requests::{ SEMANTIC_TRACE_CONVENTIONS, SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, SIGNAL_TYPE_TRACE, SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL_TRACE_V1, TABLE_DATA_MODEL_TRACE_V2, }; +use table::table_name::TableName; +use tokio::sync::OwnedSemaphorePermit; use crate::instance::Instance; use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; @@ -52,7 +59,131 @@ use crate::instance::otlp::trace_types::{ is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, truncate_for_diagnostics, }; -use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; +use crate::metrics::{ + OTLP_TRACE_AUX_CACHE_HIT, OTLP_TRACE_AUX_CACHE_MISS, OTLP_TRACES_FAILURE_COUNT, + OTLP_TRACES_ROWS, +}; + +/// Confirmed auxiliary writes, shared across requests and instance clones. +/// +/// ponytail: these tables are ingestion-managed; add invalidation if external +/// mutations must be supported without restarting every serving frontend. +#[derive(Clone)] +pub(crate) struct TraceAuxCache { + entries: Cache, +} + +#[derive(Clone, Hash, PartialEq, Eq)] +struct TraceAuxCacheKey { + table: Arc, + entry: TraceAuxEntry, +} + +impl TraceAuxCacheKey { + /// Estimated key memory, excluding cache and allocator overhead. + fn size(&self) -> usize { + let entry_size = match &self.entry { + TraceAuxEntry::Service(service) => service.capacity(), + TraceAuxEntry::Operation(service, span, kind) => { + service.capacity() + span.capacity() + kind.capacity() + } + }; + // Charge the shared table name per entry to keep the estimate conservative. + mem::size_of::() + + mem::size_of::() + + self.table.catalog_name.capacity() + + self.table.schema_name.capacity() + + self.table.table_name.capacity() + + entry_size + } +} + +#[derive(Clone, Hash, PartialEq, Eq)] +enum TraceAuxEntry { + Service(String), + Operation(String, String, String), +} + +impl TraceAuxCache { + /// Creates a cache with an estimated byte budget; zero disables caching. + pub(crate) fn new(max_size: u64) -> Self { + Self { + entries: Cache::builder() + .max_capacity(max_size) + .weigher(|key: &TraceAuxCacheKey, _: &()| { + u32::try_from(key.size()).unwrap_or(u32::MAX) + }) + .build(), + } + } + + /// Removes confirmed writes and returns keys to remember after insertion. + /// Only the first cache check in a request supplies metrics. + fn filter( + &self, + table: TableName, + aux_data: &mut TraceAuxData, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> Vec { + let lookups = aux_data.services.len() + aux_data.operations.len(); + let table = Arc::new(table); + let mut pending = Vec::new(); + let mut retain = |entry| { + let key = TraceAuxCacheKey { + table: table.clone(), + entry, + }; + if self.entries.get(&key).is_some() { + false + } else { + pending.push(key); + true + } + }; + aux_data + .services + .retain(|service| retain(TraceAuxEntry::Service(service.clone()))); + aux_data.operations.retain(|(service, span, kind)| { + retain(TraceAuxEntry::Operation( + service.clone(), + span.clone(), + kind.clone(), + )) + }); + if let Some((hits, misses)) = metrics { + hits.inc_by((lookups - pending.len()) as u64); + misses.inc_by(pending.len() as u64); + } + pending + } + + fn record(&self, keys: Vec) { + for key in keys { + self.entries.insert(key, ()); + } + } + + /// Reserves background capacity only for misses, without holding a cache lock. + async fn acquire_for_misses( + &self, + limiter: &RequestLimiter, + table: TableName, + mut aux_data: TraceAuxData, + metrics: (&IntCounter, &IntCounter), + ) -> ServerResult>> { + self.filter(table.clone(), &mut aux_data, Some(metrics)); + if aux_data.is_empty() { + return Ok(None); + } + let permit = limiter + .acquire() + .await + .map_err(|_| error::BatcherChannelClosedSnafu.build())?; + // Another request may have confirmed these entries while admission waited. + self.filter(table, &mut aux_data, None); + Ok((!aux_data.is_empty()).then_some(permit)) + } +} /// Merge converted V2 rows only when their fixed schemas agree. fn merge_trace_v2_rows(rows: &mut Rows, batch: Rows) -> ServerResult<()> { @@ -101,11 +232,52 @@ struct TraceChunkIngestContext<'a> { /// Accumulates trace outcomes, auxiliary rows, and bounded failure details. struct TraceIngestState { + needs_aux_write: bool, aux_data: TraceAuxData, + pending_aux_data: Vec<(TraceAuxData, Vec)>, outcome: TraceIngestOutcome, failure_messages: TraceFailureMessages, } +impl TraceIngestState { + fn add_aux_data(&mut self, aux_data: TraceAuxData, completions: Vec) { + if completions.is_empty() { + self.aux_data.services.extend(aux_data.services); + self.aux_data.operations.extend(aux_data.operations); + } else if !aux_data.is_empty() { + self.pending_aux_data.push((aux_data, completions)); + } + } +} + +/// Keeps confirmed chunks even if another chunk fails or loses its completion sender. +async fn confirmed_trace_aux_data( + pending: Vec<(TraceAuxData, Vec)>, +) -> TraceAuxData { + let results = futures::future::join_all(pending.into_iter().map( + |(aux_data, completions)| async move { + for completion in completions { + completion.wait().await?; + } + Ok::<_, BoxedError>(aux_data) + }, + )) + .await; + let mut confirmed = TraceAuxData::default(); + for result in results { + match result { + Ok(aux_data) => { + confirmed.services.extend(aux_data.services); + confirmed.operations.extend(aux_data.operations); + } + Err(err) => { + warn!(err; "Skipping auxiliary trace rows after asynchronous main write failure") + } + } + } + confirmed +} + /// Bounded, deduplicated failure details for one trace request. /// /// Occurrences past [`TRACE_FAILURE_MESSAGE_LIMIT`] distinct failures are @@ -876,6 +1048,32 @@ impl Instance { conventions: &str, ctx: QueryContextRef, ) -> ServerResult { + let aux_limiter = self + .trace_aux_limiter + .as_ref() + .filter(|_| ctx.batching_enabled()); + // Reserve once per request, before any main chunks are submitted. Main + // batcher permits remain independent so multi-chunk requests can progress. + let aux_permit = if let Some(limiter) = aux_limiter { + let mut aux_data = TraceAuxData::default(); + for span in groups.iter().flat_map(|group| &group.spans) { + aux_data.observe_span(span); + } + self.trace_aux_cache + .acquire_for_misses( + limiter, + TableName::new(ctx.current_catalog(), ctx.current_schema(), &table_name), + aux_data, + (&OTLP_TRACE_AUX_CACHE_HIT, &OTLP_TRACE_AUX_CACHE_MISS), + ) + .await? + } else { + None + }; + // Requests without auxiliary admission first check the cache when preparing writes. + let aux_cache_metrics = aux_limiter + .is_none() + .then_some((&*OTLP_TRACE_AUX_CACHE_HIT, &*OTLP_TRACE_AUX_CACHE_MISS)); let is_trace_v1_model = matches!(pipeline, PipelineWay::OtlpTraceDirectV1); let data_model = match pipeline { PipelineWay::OtlpTraceDirectV1 => Some(TABLE_DATA_MODEL_TRACE_V1), @@ -918,7 +1116,11 @@ impl Instance { table_name: &table_name, }; let mut ingest_state = TraceIngestState { + // Confirmed hits stay satisfied for this request even if evicted + // later: cache eviction does not delete the stored lookup rows. + needs_aux_write: aux_limiter.is_none() || aux_permit.is_some(), aux_data: TraceAuxData::default(), + pending_aux_data: Vec::new(), outcome: TraceIngestOutcome::default(), failure_messages: TraceFailureMessages::default(), }; @@ -988,20 +1190,64 @@ impl Instance { OTLP_TRACES_ROWS.inc_by(ingest_state.outcome.accepted_spans as u64); + if !ingest_state.pending_aux_data.is_empty() { + let pending = std::mem::take(&mut ingest_state.pending_aux_data); + let instance = self.clone(); + let pipeline = match pipeline { + PipelineWay::OtlpTraceDirectV1 => PipelineWay::OtlpTraceDirectV1, + PipelineWay::OtlpTraceDirectV2 => PipelineWay::OtlpTraceDirectV2, + _ => PipelineWay::OtlpTraceDirectV0, + }; + let table_name = table_name.clone(); + let ctx = ctx.clone(); + common_runtime::spawn_ingest(async move { + let _permit = aux_permit; + let aux_data = confirmed_trace_aux_data(pending).await; + let result: ServerResult<()> = async { + let (requests, keys) = instance.prepare_trace_aux_write( + aux_data, + &pipeline, + &table_name, + &ctx, + aux_cache_metrics, + )?; + if !requests.inserts.is_empty() { + instance + .insert_trace_aux_requests(requests, keys, &pipeline, ctx) + .await?; + } + Ok(()) + } + .await; + if let Err(err) = result { + OTLP_TRACES_FAILURE_COUNT + .with_label_values(&["aux_table_update_failed"]) + .inc(); + warn!(table_name, err = ?err, "Failed to write auxiliary trace rows after asynchronous main write"); + } + }); + } + if !ingest_state.aux_data.is_empty() { - // Auxiliary trace tables are derived from spans whose main-table - // writes are already confirmed, so they never create new accepted - // spans and they do not affect rejected span counts. - let aux_requests = otlp::trace::to_grpc_insert_requests_for_aux_tables( + // Auxiliary trace tables are derived from accepted spans, so they + // never create new accepted spans or affect rejected span counts. + let aux_requests = self.prepare_trace_aux_write( std::mem::take(&mut ingest_state.aux_data), ingest_ctx.pipeline, ingest_ctx.table_name, + &ctx, + aux_cache_metrics, ); match aux_requests { - Ok((aux_requests, _)) if !aux_requests.inserts.is_empty() => { + Ok((aux_requests, pending_aux_keys)) if !aux_requests.inserts.is_empty() => { match self - .insert_trace_requests(aux_requests, ingest_ctx.pipeline, ctx) + .insert_trace_aux_requests( + aux_requests, + pending_aux_keys, + ingest_ctx.pipeline, + ctx, + ) .await { Ok(output) => { @@ -1076,6 +1322,72 @@ impl Instance { Ok(ingest_state.outcome) } + fn prepare_trace_aux_write( + &self, + mut aux_data: TraceAuxData, + pipeline: &PipelineWay, + table_name: &str, + ctx: &QueryContextRef, + metrics: Option<(&IntCounter, &IntCounter)>, + ) -> ServerResult<(RowInsertRequests, Vec)> { + let keys = self.trace_aux_cache.filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + metrics, + ); + if aux_data.is_empty() { + return Ok((RowInsertRequests::default(), keys)); + } + let (mut requests, _) = + otlp::trace::to_grpc_insert_requests_for_aux_tables(aux_data, pipeline, table_name)?; + // Remove fully cached tables before v1 reconciliation performs catalog lookups. + requests.inserts.retain(|request| { + request + .rows + .as_ref() + .is_some_and(|rows| !rows.rows.is_empty()) + }); + Ok((requests, keys)) + } + + /// Checks deferred cache population without confusing visible rows with completion. + #[cfg(feature = "testing")] + pub fn trace_aux_data_cached_for_test( + &self, + table_name: &str, + mut aux_data: TraceAuxData, + ctx: &QueryContextRef, + ) -> bool { + self.trace_aux_cache + .filter( + TableName::new(ctx.current_catalog(), ctx.current_schema(), table_name), + &mut aux_data, + None, + ) + .is_empty() + } + + async fn insert_trace_aux_requests( + &self, + requests: RowInsertRequests, + keys: Vec, + pipeline: &PipelineWay, + ctx: QueryContextRef, + ) -> ServerResult { + // Cache entries require a completed write, even when main spans use + // batching with early acknowledgement. + let mut aux_ctx = ctx.fork(); + aux_ctx.set_batching_enabled(false); + let output = self + .insert_trace_requests(requests, pipeline, Arc::new(aux_ctx)) + .await?; + // ponytail: cache skip_wal writes too. After datanode data loss, lookup + // rows may stay missing until eviction/restart; require durable writes + // before caching if that limitation becomes unacceptable. + self.trace_aux_cache.record(keys); + Ok(output) + } + /// Ingest one owned trace chunk so successful spans can be moved into the /// accepted set without extra cloning. async fn ingest_trace_chunk( @@ -1167,8 +1479,12 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += chunk_rows; - for span in &chunk { - ingest_state.aux_data.observe_span(span); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + for span in &chunk { + aux_data.observe_span(span); + } + ingest_state.add_aux_data(aux_data, output.meta.write_completions); } Ok(()) @@ -1345,7 +1661,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += span_count; - retry.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + retry.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } Ok(()) } @@ -1428,7 +1748,11 @@ impl Instance { Self::add_trace_write_cost(&mut ingest_state.outcome, output.meta.cost); ingest_state.outcome.accepted_spans += 1; - span.add_to_aux_data(&mut ingest_state.aux_data); + if ingest_state.needs_aux_write { + let mut aux_data = TraceAuxData::default(); + span.add_to_aux_data(&mut aux_data); + ingest_state.add_aux_data(aux_data, output.meta.write_completions); + } } Ok(()) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 1459119efdd..d497db83e13 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use api::v1::column_data_type_extension::TypeExt; use api::v1::value::ValueData; @@ -20,26 +21,269 @@ use api::v1::{ ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row, RowInsertRequest, Rows, SemanticType, Value, }; +use common_batcher::request_limiter::RequestLimiter; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::WriteCompletion; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ ColumnSchema as DatatypesColumnSchema, SchemaBuilder as DatatypesSchemaBuilder, }; use opentelemetry_proto::tonic::common::v1::any_value::Value as OtlpValue; use opentelemetry_proto::tonic::common::v1::{AnyValue, ArrayValue, KeyValue}; -use servers::otlp::trace::SERVICE_NAME_COLUMN; +use prometheus::IntCounter; use servers::otlp::trace::attributes::Attributes; use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; +use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; +use table::table_name::TableName; use super::{ - ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, - TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + ChunkFailureReaction, Instance, TraceAuxCache, TraceAuxCacheKey, TraceAuxEntry, + TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, TraceRequestSchema, + TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; +fn trace_aux_data(service: &str, span: &str, kind: &str) -> TraceAuxData { + TraceAuxData { + services: HashSet::from([service.to_string()]), + operations: HashSet::from([(service.to_string(), span.to_string(), kind.to_string())]), + } +} + +fn trace_aux_cache_metrics() -> (IntCounter, IntCounter) { + ( + IntCounter::new("test_trace_aux_cache_hit", "Cache hits").unwrap(), + IntCounter::new("test_trace_aux_cache_miss", "Cache misses").unwrap(), + ) +} + +#[tokio::test] +async fn test_trace_aux_waits_for_successful_chunk_completions() { + let (success_tx, success_rx) = tokio::sync::oneshot::channel(); + let (failure_tx, failure_rx) = tokio::sync::oneshot::channel(); + let completion = |receiver: tokio::sync::oneshot::Receiver| { + WriteCompletion::new(async move { + if receiver.await.unwrap_or(false) { + Ok(()) + } else { + Err(common_error::ext::BoxedError::new( + servers::error::InternalSnafu { + err_msg: "main write failed", + } + .build(), + )) + } + }) + }; + let mut confirmed = Box::pin(super::confirmed_trace_aux_data(vec![ + ( + trace_aux_data("good", "span", "server"), + vec![completion(success_rx)], + ), + ( + trace_aux_data("bad", "span", "server"), + vec![completion(failure_rx)], + ), + ])); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + success_tx.send(true).unwrap(); + assert!(futures::poll!(confirmed.as_mut()).is_pending()); + // A lost completion sender must not count as successful storage. + drop(failure_tx); + let aux = confirmed.await; + assert_eq!(aux.services, HashSet::from(["good".to_string()])); + assert_eq!( + aux.operations, + trace_aux_data("good", "span", "server").operations + ); +} + +#[test] +fn test_trace_aux_cache_confirmed_writes() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let metrics = Some((&hits, &misses)); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + let pending = cache.filter(table.clone(), &mut data, metrics); + assert_eq!(pending.len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + + // Preparing a write must not suppress a retry after failure or cancellation. + let mut retry = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table.clone(), &mut retry, None).len(), 2); + assert_eq!((hits.get(), misses.get()), (0, 2)); + cache.record(pending); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert!( + cache + .clone() + .filter(table.clone(), &mut repeated, metrics) + .is_empty() + ); + assert!(repeated.is_empty()); + assert_eq!((hits.get(), misses.get()), (2, 2)); + cache.filter(table.clone(), &mut retry, None); + assert_eq!((hits.get(), misses.get()), (2, 2)); + + for (span, kind) in [("new_span", "server"), ("span", "client")] { + let mut data = trace_aux_data("svc", span, kind); + assert_eq!(cache.filter(table.clone(), &mut data, metrics).len(), 1); + assert!(data.services.is_empty()); + assert_eq!(data.operations.len(), 1); + } + assert_eq!((hits.get(), misses.get()), (4, 4)); + cache.filter(table.clone(), &mut TraceAuxData::default(), metrics); + assert_eq!((hits.get(), misses.get()), (4, 4)); + + let mut other_service = trace_aux_data("other", "span", "server"); + assert_eq!(cache.filter(table, &mut other_service, None).len(), 2); + + for other_table in [ + TableName::new("other", "public", "traces"), + TableName::new("greptime", "other", "traces"), + TableName::new("greptime", "public", "other"), + ] { + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(other_table, &mut data, None).len(), 2); + } + + // An operation hit must not suppress a missing service entry. + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + data.services.clear(); + cache.record(cache.filter(table.clone(), &mut data, None)); + + let mut data = trace_aux_data("svc", "span", "server"); + assert_eq!(cache.filter(table, &mut data, None).len(), 1); + assert_eq!(data.services.len(), 1); + assert!(data.operations.is_empty()); +} + +#[test] +fn test_trace_aux_cache_zero_size_disables_caching() { + let cache = TraceAuxCache::new(0); + let (hits, misses) = trace_aux_cache_metrics(); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, Some((&hits, &misses)))); + + let mut repeated = trace_aux_data("svc", "span", "server"); + assert_eq!( + cache + .filter(table, &mut repeated, Some((&hits, &misses))) + .len(), + 2 + ); + assert_eq!(repeated.services.len(), 1); + assert_eq!(repeated.operations.len(), 1); + assert_eq!(hits.get(), 0); + assert_eq!(misses.get(), 4); +} + +#[test] +fn test_trace_aux_cache_accounts_for_key_memory() { + let cache = TraceAuxCache::new(1024); + let table = TableName::new("greptime", "public", "traces"); + let mut data = trace_aux_data("svc", "span", "server"); + cache.record(cache.filter(table.clone(), &mut data, None)); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 2); + assert!(cache.entries.weighted_size() > 2); + assert!(cache.entries.weighted_size() <= 1024); + + // Spare string capacity counts even when the string content is empty. + let large = || String::with_capacity(2048); + let service = || TraceAuxEntry::Service("svc".to_string()); + for (table, entry) in [ + (TableName::new(large(), "public", "traces"), service()), + (TableName::new("greptime", large(), "traces"), service()), + (TableName::new("greptime", "public", large()), service()), + (table.clone(), TraceAuxEntry::Service(large())), + ( + table.clone(), + TraceAuxEntry::Operation(large(), String::new(), String::new()), + ), + ( + table.clone(), + TraceAuxEntry::Operation(String::new(), large(), String::new()), + ), + ( + table, + TraceAuxEntry::Operation(String::new(), String::new(), large()), + ), + ] { + let cache = TraceAuxCache::new(1024); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(table), + entry, + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); + } + + // Empty strings still retain the key and table structs. + let cache = TraceAuxCache::new(1); + cache.record(vec![TraceAuxCacheKey { + table: Arc::new(TableName::new("", "", "")), + entry: TraceAuxEntry::Service(String::new()), + }]); + cache.entries.run_pending_tasks(); + assert_eq!(cache.entries.entry_count(), 0); +} + +#[tokio::test] +async fn test_trace_aux_admission_rechecks_cache() { + let cache = TraceAuxCache::new(1024); + let (hits, misses) = trace_aux_cache_metrics(); + let limiter = RequestLimiter::try_new(1).unwrap(); + let table = TableName::new("greptime", "public", "traces"); + let data = || trace_aux_data("svc", "span", "server"); + let permit = cache + .acquire_for_misses(&limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .unwrap(); + let cloned_cache = cache.clone(); + let cloned_limiter = limiter.clone(); + let mut waiting = Box::pin(cloned_cache.acquire_for_misses( + &cloned_limiter, + table.clone(), + data(), + (&hits, &misses), + )); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + assert_eq!((hits.get(), misses.get()), (0, 4)); + + // Successful storage populates the shared cache before releasing admission. + cache.record(cache.filter(table.clone(), &mut data(), None)); + assert!( + cloned_cache + .acquire_for_misses(&cloned_limiter, table.clone(), data(), (&hits, &misses)) + .await + .unwrap() + .is_none() + ); + drop(permit); + assert!(waiting.await.unwrap().is_none()); + assert_eq!((hits.get(), misses.get()), (2, 4)); + + // Rechecking a now-warm request must release its briefly acquired slot. + let other = cache.acquire_for_misses( + &limiter, + table, + trace_aux_data("svc", "other", "server"), + (&hits, &misses), + ); + assert!(futures::poll!(Box::pin(other).as_mut()).is_ready()); + assert_eq!((hits.get(), misses.get()), (3, 5)); +} + #[test] fn test_chunk_owned() { let chunks = chunk_owned(vec![1, 2, 3], 2); diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index f649db0b46f..74574e56930 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; use lazy_static::lazy_static; use prometheus::*; @@ -68,6 +69,12 @@ lazy_static! { ) .unwrap(); + /// Initial trace auxiliary-cache lookups, excluding admission and write rechecks. + pub static ref OTLP_TRACE_AUX_CACHE_HIT: IntCounter = + CACHE_HIT.with_label_values(&["otlp_trace_aux"]); + pub static ref OTLP_TRACE_AUX_CACHE_MISS: IntCounter = + CACHE_MISS.with_label_values(&["otlp_trace_aux"]); + /// The number of OpenTelemetry logs send by frontend node. pub static ref OTLP_LOGS_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_logs_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index e67f4b5bcba..7a6d3f75b91 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -12,9 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_base::readable_size::ReadableSize; +use common_stat::get_total_memory_readable; use serde::{Deserialize, Serialize}; const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; +const MIN_TRACE_AUX_CACHE_SIZE: ReadableSize = ReadableSize::mb(32); #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] @@ -22,6 +25,12 @@ pub struct OtlpOptions { pub enable: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Estimated memory budget for cached trace service/operation keys per frontend, + /// shared across all catalogs, schemas, and trace tables. Set to 0 to disable caching. + /// Defaults to 1/128 of the host or pod memory limit, with a minimum of 32MiB. + /// Uses 32MiB if memory detection is unavailable. + /// Excludes cache and allocator overhead; shared table names are charged per entry. + pub trace_aux_cache_size: ReadableSize, /// Whether to synthesize the `greptime_otel_resource_info` descriptor /// table from the resource attributes of OTLP metrics, so metrics-only /// services reach the semantic graph. Off by default: it creates and @@ -34,24 +43,51 @@ impl Default for OtlpOptions { Self { enable: true, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + trace_aux_cache_size: default_trace_aux_cache_size(get_total_memory_readable()), experimental_enable_resource_info: false, } } } +fn default_trace_aux_cache_size(memory: Option) -> ReadableSize { + memory + .map(|size| size / 128) + .unwrap_or_default() + .max(MIN_TRACE_AUX_CACHE_SIZE) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_default_trace_aux_cache_size() { + for (memory, expected) in [ + (None, ReadableSize::mb(32)), + (Some(ReadableSize(0)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(1)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(4)), ReadableSize::mb(32)), + (Some(ReadableSize::gb(8)), ReadableSize::mb(64)), + (Some(ReadableSize::gb(64)), ReadableSize::mb(512)), + ] { + assert_eq!(default_trace_aux_cache_size(memory), expected); + } + } + #[test] fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); assert_eq!(default.trace_ingest_chunk_size, 512); + assert_eq!( + default.trace_aux_cache_size, + default_trace_aux_cache_size(get_total_memory_readable()) + ); assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); + assert_eq!(options.trace_aux_cache_size, default.trace_aux_cache_size); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -63,5 +99,13 @@ mod tests { let serialized = toml::to_string(&options).unwrap(); assert_eq!(toml::from_str::(&serialized).unwrap(), options); + + for (value, size) in [("0", ReadableSize(0)), ("\"1MiB\"", ReadableSize::mb(1))] { + let options: OtlpOptions = + toml::from_str(&format!("trace_aux_cache_size = {value}")).unwrap(); + assert_eq!(options.trace_aux_cache_size, size); + let serialized = toml::to_string(&options).unwrap(); + assert_eq!(toml::from_str::(&serialized).unwrap(), options); + } } } diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md index 62853a18f61..49cbd617567 100644 --- a/src/operator/AGENTS.md +++ b/src/operator/AGENTS.md @@ -46,3 +46,5 @@ the default `meter-macros/noop` build and active metering when changing this pat Finite requests split internally use `admit_write` / `admit_row_insert_batches` before dispatch. Their database-scoped `QueryContext` admission prevents a second row debit while `do_request` still records WCU for the actual writes. +Async table-batcher outputs carry storage completions in `OutputMeta`; keep +them when combining submissions so callers can distinguish admission from success. diff --git a/src/operator/src/batcher.rs b/src/operator/src/batcher.rs index 677524d53f3..bc9ef86b34e 100644 --- a/src/operator/src/batcher.rs +++ b/src/operator/src/batcher.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use arrow::record_batch::RecordBatch; use async_trait::async_trait; +use common_query::Output; use session::context::QueryContextRef; use table::metadata::TableInfoRef; use tokio::sync::OwnedSemaphorePermit; @@ -32,11 +33,12 @@ pub trait PendingRowsBatcher: Send + Sync { /// Submits rows according to the acknowledgement policy, retaining the slot until /// writing completes even when the response acknowledges queue admission only. /// Cancelling the response wait does not retract an already enqueued submission. + /// Early acknowledgements retain actual storage results in the output metadata. async fn submit( &self, table_info: TableInfoRef, batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result; + ) -> Result; } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 80866b45b66..c7f60c0883c 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -555,11 +555,14 @@ impl Inserter { // Observe every table completion even when another table fails. future::join_all(submissions).await }; - let affected_rows = results.into_iter().sum::>()?; - Ok(Output::new( - OutputData::AffectedRows(affected_rows), - OutputMeta::new_with_cost(write_cost as _), - )) + let mut affected_rows = 0; + let mut meta = OutputMeta::new_with_cost(write_cost as _); + for result in results { + let output = result?; + affected_rows += output.extract_rows_and_cost().0; + meta.write_completions.extend(output.meta.write_completions); + } + Ok(Output::new(OutputData::AffectedRows(affected_rows), meta)) } /// Handles row inserts request with metric engine. @@ -3302,7 +3305,7 @@ mod tests { _batch: arrow::record_batch::RecordBatch, _ctx: QueryContextRef, _permit: Arc, - ) -> Result { + ) -> Result { panic!("empty writes must not submit a batch") } } diff --git a/src/pipeline/Cargo.toml b/src/pipeline/Cargo.toml index 5a9cb3e9aa6..9ebe51ecf24 100644 --- a/src/pipeline/Cargo.toml +++ b/src/pipeline/Cargo.toml @@ -19,6 +19,7 @@ chrono.workspace = true chrono-tz.workspace = true common-catalog.workspace = true common-error.workspace = true +common-frontend.workspace = true common-macro.workspace = true common-meta.workspace = true common-query.workspace = true diff --git a/src/pipeline/src/manager/pipeline_cache.rs b/src/pipeline/src/manager/pipeline_cache.rs index e60963e7f84..f0e9ca52970 100644 --- a/src/pipeline/src/manager/pipeline_cache.rs +++ b/src/pipeline/src/manager/pipeline_cache.rs @@ -16,6 +16,7 @@ use std::future::Future; use std::sync::Arc; use std::time::Duration; +use common_frontend::metrics; use datatypes::timestamp::TimestampNanosecond; use moka::future::Cache; @@ -32,7 +33,7 @@ const PIPELINES_CACHE_SIZE: u64 = 10000; /// to encapsulate inner cache. Only public methods are exposed. /// /// `pipelines` and `original_pipelines` are keyed by the *requested* schema so -/// a lookup is a single key probe, as [`Cache::try_get_with`] requires; +/// a lookup is a single key probe through [`Cache::entry`]; /// resolving it to a stored schema is the loader's job. `failover_cache` has no /// loader and keeps the stored-schema key. pub(crate) struct PipelineCache { @@ -71,7 +72,7 @@ impl PipelineCache { } } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_with( &self, schema: &str, @@ -80,13 +81,17 @@ impl PipelineCache { init: impl Future>>, ) -> Result> { let key = generate_pipeline_cache_key(schema, name, version); - self.pipelines - .try_get_with(key, init) - .await + let entry = self.pipelines.entry(key).or_try_insert_with(init).await; + metrics::record_cache_lookup( + "pipeline", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } - /// Concurrent misses on the same key share one `init` call. + /// Concurrent misses share one `init` call; successful waiters count as hits. pub(crate) async fn get_pipeline_str_with( &self, schema: &str, @@ -95,9 +100,17 @@ impl PipelineCache { init: impl Future>, ) -> Result { let key = generate_pipeline_cache_key(schema, name, version); - self.original_pipelines - .try_get_with(key, init) - .await + let entry = self + .original_pipelines + .entry(key) + .or_try_insert_with(init) + .await; + metrics::record_cache_lookup( + "pipeline_original", + entry.as_ref().is_ok_and(|entry| !entry.is_fresh()), + ); + entry + .map(|entry| entry.into_value()) .map_err(|error| CacheLoadSnafu { error }.build()) } @@ -114,6 +127,7 @@ impl PipelineCache { generate_pipeline_cache_key(schema, name, version), ] { if let Some(content) = self.failover_cache.get(&key).await { + metrics::record_cache_lookup("pipeline_failover", true); return Ok(Some(content)); } } @@ -126,7 +140,7 @@ impl PipelineCache { .filter(|(k, _)| k.ends_with(&suffix)) .collect::>(); - match found.len() { + let result = match found.len() { 0 => Ok(None), 1 => Ok(Some(found.remove(0).1)), _ => MultiPipelineWithDiffSchemaSnafu { @@ -139,7 +153,12 @@ impl PipelineCache { .join(","), } .fail(), - } + }; + metrics::record_cache_lookup( + "pipeline_failover", + result.as_ref().is_ok_and(|content| content.is_some()), + ); + result } pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) { @@ -191,7 +210,8 @@ impl PipelineCache { mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; - use tokio::sync::Barrier; + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + use tokio::sync::Notify; use super::*; @@ -205,36 +225,149 @@ mod tests { } } - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[tokio::test] async fn test_concurrent_misses_run_one_loader() { const CONCURRENCY: usize = 8; - let cache = Arc::new(PipelineCache::new(Duration::from_secs(60))); - let loads = Arc::new(AtomicUsize::new(0)); - let barrier = Arc::new(Barrier::new(CONCURRENCY)); - - let handles = (0..CONCURRENCY) + let hits = CACHE_HIT.with_label_values(&["pipeline_original"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_original"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + let loads = AtomicUsize::new(0); + let release = Notify::new(); + let mut requests = (0..CONCURRENCY) .map(|_| { - let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone()); - tokio::spawn(async move { - barrier.wait().await; + Box::pin(cache.get_pipeline_str_with("db", "p", None, async { + loads.fetch_add(1, Ordering::SeqCst); + release.notified().await; + Ok(content_at(1)) + })) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request.as_mut()).is_pending()); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + release.notify_waiters(); + for request in requests { + assert_eq!(request.await.unwrap(), content_at(1)); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!( + (hits.get(), misses.get()), + (before.0 + CONCURRENCY as u64 - 1, before.1 + 1) + ); + } + + #[tokio::test] + async fn test_pipeline_lookup_metrics() { + let cache = PipelineCache::new(Duration::from_secs(60)); + for compiled in [false, true] { + let cache_type = if compiled { + "pipeline" + } else { + "pipeline_original" + }; + let hits = CACHE_HIT.with_label_values(&[cache_type]); + let misses = CACHE_MISS.with_label_values(&[cache_type]); + let before = (hits.get(), misses.get()); + let load = async |fail| { + if compiled { + cache + .get_pipeline_with("db", "p", None, async { + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } + Ok(Arc::new( + crate::manager::table::PipelineTable::compile_pipeline( + "transform:\n - fields: [message]\n type: string", + )?, + )) + }) + .await + .map(|_| ()) + } else { cache .get_pipeline_str_with("db", "p", None, async { - loads.fetch_add(1, Ordering::SeqCst); - // Hold the loader open so every caller is waiting on it. - tokio::time::sleep(Duration::from_millis(100)).await; + if fail { + return Err(crate::error::PipelineNotFoundSnafu { + name: "p", + version: None, + } + .build()); + } Ok(content_at(1)) }) .await - .unwrap() - }) - }) - .collect::>(); - - for handle in handles { - assert_eq!(handle.await.unwrap(), content_at(1)); + .map(|_| ()) + } + }; + assert!(load(true).await.is_err()); + load(false).await.unwrap(); + // A warm lookup must not evaluate the failing initializer. + load(true).await.unwrap(); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); } - assert_eq!(loads.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_failover_lookup_metrics() { + let hits = CACHE_HIT.with_label_values(&["pipeline_failover"]); + let misses = CACHE_MISS.with_label_values(&["pipeline_failover"]); + let before = (hits.get(), misses.get()); + let cache = PipelineCache::new(Duration::from_secs(60)); + assert!( + cache + .get_failover_cache("db", "p", None) + .await + .unwrap() + .is_none() + ); + let local = PipelineContent { + schema: "a".into(), + ..content_at(1) + }; + cache.insert_failover_cache(local.clone(), true).await; + // Direct schema lookup and cross-schema fallback each count one hit. + assert_eq!( + cache.get_failover_cache("a", "p", None).await.unwrap(), + Some(local.clone()) + ); + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(local) + ); + cache + .insert_failover_cache( + PipelineContent { + schema: "b".into(), + ..content_at(2) + }, + true, + ) + .await; + assert!(cache.get_failover_cache("db", "p", None).await.is_err()); + // Global pipelines win over local entries inserted before or after them. + for (schema, version) in [(EMPTY_SCHEMA_NAME, 3), ("c", 4)] { + cache + .insert_failover_cache( + PipelineContent { + schema: schema.into(), + ..content_at(version) + }, + true, + ) + .await; + assert_eq!( + cache.get_failover_cache("db", "p", None).await.unwrap(), + Some(content_at(3)) + ); + } + assert_eq!((hits.get(), misses.get()), (before.0 + 4, before.1 + 2)); } #[tokio::test] @@ -288,25 +421,4 @@ mod tests { let failover = cache.get_failover_cache("b", "p", None).await.unwrap(); assert_eq!(failover.map(|c| c.version), Some(v2.version)); } - - #[tokio::test] - async fn test_failover_serves_global_pipeline_to_unwarmed_schema() { - let cache = PipelineCache::new(Duration::from_secs(60)); - let content = content_at(1); - - cache.insert_failover_cache(content.clone(), true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content.clone())); - - // A same-named pipeline under another schema must not shadow the global one. - let schema_local = PipelineContent { - schema: "x".to_string(), - ..content_at(2) - }; - cache.insert_failover_cache(schema_local, true).await; - - let found = cache.get_failover_cache("b", "p", None).await.unwrap(); - assert_eq!(found, Some(content)); - } } diff --git a/src/pipeline/src/manager/pipeline_operator.rs b/src/pipeline/src/manager/pipeline_operator.rs index c9731b36f20..cace74171c2 100644 --- a/src/pipeline/src/manager/pipeline_operator.rs +++ b/src/pipeline/src/manager/pipeline_operator.rs @@ -19,6 +19,7 @@ use std::time::{Duration, Instant}; use api::v1::CreateTableExpr; use catalog::{CatalogManagerRef, RegisterSystemTableRequest}; use common_catalog::consts::{DEFAULT_PRIVATE_SCHEMA_NAME, default_engine}; +use common_frontend::metrics; use common_meta::rpc::ddl::TriggerReason; use common_telemetry::info; use common_time::FOREVER; @@ -164,7 +165,9 @@ impl PipelineOperator { /// Get a pipeline table from the cache. pub fn get_pipeline_table_from_cache(&self, catalog: &str) -> Option { - self.tables.read().unwrap().get(catalog).cloned() + let table = self.tables.read().unwrap().get(catalog).cloned(); + metrics::record_cache_lookup("pipeline_table", table.is_some()); + table } } diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md index 006ac59771c..d20704cf616 100644 --- a/src/servers/AGENTS.md +++ b/src/servers/AGENTS.md @@ -41,3 +41,5 @@ cargo nextest run -p servers ``` Keep protocol translation here and permissions/database behavior in frontend. +Async table batching returns storage-completion futures with its early +acknowledgements. Keep completion reporting independent of the caller's lifetime. diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 8a2b6ab2149..0be0e47eab7 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -40,6 +40,7 @@ common-batcher.workspace = true common-catalog.workspace = true common-decimal.workspace = true common-error.workspace = true +common-frontend.workspace = true common-grpc.workspace = true common-macro.workspace = true common-mem-prof = { workspace = true, optional = true } diff --git a/src/servers/src/batcher/table.rs b/src/servers/src/batcher/table.rs index 67b3431a810..daae05251f8 100644 --- a/src/servers/src/batcher/table.rs +++ b/src/servers/src/batcher/table.rs @@ -30,6 +30,8 @@ use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::timing::TimingFlushPolicy; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; +use common_error::ext::BoxedError; +use common_query::{Output, WriteCompletion}; use operator::batcher::PendingRowsBatcher; use operator::error::{BatchFlushSnafu, Result, UnexpectedSnafu}; use operator::insert::Inserter; @@ -171,10 +173,10 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { batch: RecordBatch, ctx: QueryContextRef, permit: Arc, - ) -> Result { + ) -> Result { let total_rows = batch.num_rows(); if total_rows == 0 { - return Ok(0); + return Ok(Output::new_with_affected_rows(0)); } if table_info.catalog_name != ctx.current_catalog() || table_info.schema_name != ctx.current_schema() @@ -213,19 +215,30 @@ impl PendingRowsBatcher for TablePendingRowsBatcher { } .fail(); } - if !self.pending_rows_batch_sync { - return Ok(total_rows); + let completion = async move { + response_rx + .await + .map_err(|_| { + UnexpectedSnafu { + violated: "batch worker stopped before reporting its write result" + .to_string(), + } + .build() + })? + .context(BatchFlushSnafu) + }; + let mut output = Output::new_with_affected_rows(total_rows); + if self.pending_rows_batch_sync { + completion.await?; + } else { + output + .meta + .write_completions + .push(WriteCompletion::new(async move { + completion.await.map_err(BoxedError::new) + })); } - response_rx - .await - .map_err(|_| { - UnexpectedSnafu { - violated: "batch worker stopped before reporting its write result".to_string(), - } - .build() - })? - .context(BatchFlushSnafu)?; - Ok(total_rows) + Ok(output) } } @@ -386,17 +399,20 @@ mod tests { .unwrap() .unwrap(); let mut submitted = Some(submitted); + let mut completion = None; if sync { assert!(!submitted.as_ref().unwrap().is_finished()); } else { - assert_eq!( - timeout(Duration::from_secs(5), submitted.take().unwrap()) - .await - .unwrap() - .unwrap() - .unwrap(), - 1 - ); + let mut output = timeout(Duration::from_secs(5), submitted.take().unwrap()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert_eq!(output.meta.write_completions.len(), 1); + let mut pending = Box::pin(output.meta.write_completions.pop().unwrap().wait()); + assert!(futures::poll!(pending.as_mut()).is_pending()); + completion = Some(pending); } // Early acknowledgement must not release capacity before completion. let acquire = batcher.acquire(); @@ -420,9 +436,15 @@ mod tests { .unwrap(); assert_eq!(result.is_err(), fail); if !fail { - assert_eq!(result.unwrap(), 1); + let output = result.unwrap(); + assert_eq!(output.extract_rows_and_cost().0, 1); + assert!(output.meta.write_completions.is_empty()); } } + if let Some(completion) = completion { + let result = timeout(Duration::from_secs(5), completion).await.unwrap(); + assert_eq!(result.is_err(), fail); + } timeout(Duration::from_secs(5), acquire) .await .unwrap() @@ -533,9 +555,13 @@ mod tests { let (first_result, second_result) = timeout(Duration::from_secs(5), async { if shared_request { let permit = batcher.acquire().await.unwrap(); - tokio::join!( + let (first, second) = tokio::join!( batcher.submit(table.clone(), first, influx_ctx, permit.clone()), batcher.submit(table.clone(), second, opentsdb_ctx, permit), + ); + ( + first.map(|output| output.extract_rows_and_cost().0), + second.map(|output| output.extract_rows_and_cost().0), ) } else { let inserter = Inserter::new( diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs index d0dd1f25279..bdfc7a72057 100644 --- a/src/servers/src/mysql/handler.rs +++ b/src/servers/src/mysql/handler.rs @@ -23,6 +23,7 @@ use async_trait::async_trait; use chrono::{NaiveDate, NaiveDateTime}; use common_catalog::parse_optional_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_frontend::metrics; use common_query::Output; use common_telemetry::{debug, error, tracing, warn}; use common_time::Timezone; @@ -182,8 +183,9 @@ impl MysqlInstanceShim { /// Retrieve the query and logical plan by a given statement key fn plan(&self, stmt_key: &str) -> Option { - let guard = self.prepared_stmts.read(); - guard.get(stmt_key).cloned() + let plan = self.prepared_stmts.read().get(stmt_key).cloned(); + metrics::record_cache_lookup("mysql_prepared_stmt", plan.is_some()); + plan } /// Save the prepared statement and return the parameters and result columns @@ -1053,6 +1055,23 @@ mod tests { ) } + #[test] + fn test_prepared_statement_cache_metrics() { + use common_frontend::metrics::{CACHE_HIT, CACHE_MISS}; + + let hits = CACHE_HIT.with_label_values(&["mysql_prepared_stmt"]); + let misses = CACHE_MISS.with_label_values(&["mysql_prepared_stmt"]); + let before = (hits.get(), misses.get()); + let mut shim = create_shim(); + assert!(shim.plan("stmt").is_none()); + shim.save_plan(SqlPlan::Shortcut("SELECT 1".into()), "stmt".into()) + .unwrap(); + assert!(shim.plan("stmt").is_some()); + shim.do_close("stmt".into()); + assert!(shim.plan("stmt").is_none()); + assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2)); + } + #[test] fn test_batching_context() { for enabled in [false, true] { diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 2c90f4ee2ae..79df85475d9 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -123,6 +123,7 @@ sqlx = { workspace = true, features = [ "chrono", ] } store-api.workspace = true +temp-env.workspace = true time = "0.3" tokio-postgres = { workspace = true } tokio-util.workspace = true diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 29c8891a695..3e447bb0194 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -14,17 +14,28 @@ #[cfg(test)] mod test { - use std::sync::Arc; + use std::sync::{Arc, Mutex}; + use std::time::Duration; use TimeUnit as ArrowTimeUnit; + use api::region::RegionResponse; + use api::v1::region::{RegionRequest, region_request}; use client::{DEFAULT_CATALOG_NAME, OutputData}; - use common_recordbatch::RecordBatches; + use common_meta::cache::LayeredCacheRegistryBuilder; + use common_meta::node_manager::{ + Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, NodeManagerRef, + }; + use common_meta::peer::Peer; + use common_query::request::QueryRequest; + use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::{ DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; use frontend::instance::Instance; + use frontend::instance::builder::FrontendBuilder; + use frontend::service_config::BatcherOptions; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::any_value::Value as Val; @@ -39,14 +50,574 @@ mod test { use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use pipeline::{GreptimePipelineParams, PipelineWay}; use serde_json::json; + use servers::batcher::BatchingProtocol; + use servers::otlp::trace::{TraceAuxData, span}; use servers::query_handler::OpenTelemetryProtocolHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; + use store_api::region_engine::RegionEngine; + use store_api::storage::RegionId; + use tokio::sync::oneshot; - use crate::standalone::GreptimeDbStandaloneBuilder; + use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; use crate::tests; + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache() { + check_trace_aux_cache(false).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_trace_aux_cache_with_batching_v0_v1() { + check_trace_aux_cache(true).await; + } + + #[test] + fn test_trace_aux_async_completion() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(check_trace_aux_async_completion()); + }); + } + + #[test] + fn test_trace_aux_admission() { + temp_env::with_var("PENDING_ROWS_BATCH_SYNC", Some("false"), || { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + tokio::time::timeout(Duration::from_secs(60), check_trace_aux_admission()) + .await + .expect("auxiliary admission must make progress"); + }); + }); + } + + struct TraceAuxWriteGate { + region_id: RegionId, + started: oneshot::Sender<()>, + release: oneshot::Receiver, + } + + struct GatedTraceNodeManager { + inner: NodeManagerRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl DatanodeManager for GatedTraceNodeManager { + async fn datanode(&self, peer: &Peer) -> DatanodeRef { + Arc::new(GatedTraceDatanode { + inner: self.inner.datanode(peer).await, + gate: self.gate.clone(), + }) + } + } + + #[async_trait::async_trait] + impl FlownodeManager for GatedTraceNodeManager { + async fn flownode(&self, peer: &Peer) -> FlownodeRef { + self.inner.flownode(peer).await + } + } + + struct GatedTraceDatanode { + inner: DatanodeRef, + gate: Arc>>, + } + + #[async_trait::async_trait] + impl Datanode for GatedTraceDatanode { + async fn handle( + &self, + request: RegionRequest, + ) -> common_meta::error::Result { + let gate = { + let mut gate = self.gate.lock().unwrap(); + if gate.as_ref().is_some_and(|gate| { + matches!(&request.body, Some(region_request::Body::Inserts(inserts)) + if inserts.requests.iter().any(|insert| insert.region_id == gate.region_id.as_u64())) + }) { + gate.take() + } else { + None + } + }; + if let Some(gate) = gate { + gate.started.send(()).unwrap(); + if !gate.release.await.unwrap() { + return common_meta::error::UnexpectedSnafu { + err_msg: "injected auxiliary write failure", + } + .fail(); + } + } + self.inner.handle(request).await + } + + async fn handle_query( + &self, + request: QueryRequest, + ) -> common_meta::error::Result { + self.inner.handle_query(request).await + } + } + + async fn check_trace_aux_admission() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_admission") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 1, + max_inflight_requests: 1, + ..Default::default() + }) + .build() + .await; + let source = standalone.fe_instance(); + let gate = Arc::new(Mutex::new(None)); + let cache_registry = Arc::new( + cache::with_default_composite_cache_registry( + LayeredCacheRegistryBuilder::default().add_cache_registry( + cache::build_fundamental_cache_registry(standalone.kv_backend.clone()), + ), + ) + .unwrap() + .build(), + ); + let mut options = standalone.opts.frontend_options(); + options.otlp.trace_ingest_chunk_size = 1; + let instance = Arc::new( + FrontendBuilder::new( + options, + standalone.kv_backend.clone(), + cache_registry, + source.catalog_manager().clone(), + Arc::new(GatedTraceNodeManager { + inner: source.node_manager().clone(), + gate: gate.clone(), + }), + source.procedure_executor().clone(), + Arc::new(catalog::process_manager::ProcessManager::new( + "trace_aux_admission".into(), + None, + )), + ) + .try_build() + .await + .unwrap(), + ); + + for version in 0..=1 { + let table_name = format!("trace_aux_admission_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let write = async |request| { + instance + .traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ) + .await + .unwrap() + }; + let cached = |request: ExportTraceServiceRequest| { + let mut data = TraceAuxData::default(); + for span in span::parse(request).iter().flat_map(|group| &group.spans) { + data.observe_span(span); + } + instance.trace_aux_data_cached_for_test(&table_name, data, &ctx) + }; + let warm = trace_aux_request(); + assert_eq!(1, write(warm.clone()).await.accepted_spans); + while !cached(warm.clone()) { + tokio::task::yield_now().await; + } + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}_operations"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + let region_id = route.region_routes[0].region.id; + + for success in [true, false] { + let mut cold = warm.clone(); + let spans = &mut cold.resource_spans[0].scope_spans[0].spans; + spans[0].name = format!("cold_{success}"); + spans.push(spans[0].clone()); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *gate.lock().unwrap() = Some(TraceAuxWriteGate { + region_id, + started: started_tx, + release: release_rx, + }); + + // Two chunks must progress with just one main and one auxiliary slot. + assert_eq!(2, write(cold.clone()).await.accepted_spans); + started_rx.await.unwrap(); + assert!( + !cached(cold.clone()), + "pending writes must not warm the cache" + ); + let before = trace_aux_sequences(&standalone, &table_name).await; + let mut waiting = Box::pin(write(cold.clone())); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + + // Main admission is free while auxiliary storage is blocked. Warm + // requests must bypass auxiliary admission and keep making progress. + assert_eq!(1, write(warm.clone()).await.accepted_spans); + assert!(futures::poll!(waiting.as_mut()).is_pending()); + release_tx.send(success).unwrap(); + assert_eq!(2, waiting.await.accepted_spans); + if !success { + // The failed write cannot satisfy the waiting request's miss. + while !cached(cold.clone()) { + tokio::task::yield_now().await; + } + } + let after = trace_aux_sequences(&standalone, &table_name).await; + assert_eq!(before[1], after[1], "cached service must not be rewritten"); + assert_eq!( + before[2] + 1, + after[2], + "exactly one successful operation write" + ); + } + } + } + + fn trace_aux_request() -> ExportTraceServiceRequest { + serde_json::from_value(json!({ + "resourceSpans": [{ + "resource": {"attributes": [ + {"key": "service.name", "value": {"stringValue": "svc"}} + ]}, + "scopeSpans": [{"spans": [{ + "traceId": "c05d7a4ec8e1f231f02ed6e8da8655b4", + "spanId": "9630f2916e2f7909", + "name": "operation", + "kind": 2, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000" + }]}] + }] + })) + .unwrap() + } + + async fn trace_aux_sequences(standalone: &GreptimeDbStandalone, table_name: &str) -> Vec { + let instance = standalone.fe_instance(); + let mut sequences = Vec::new(); + for suffix in ["", "_services", "_operations"] { + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None, + ) + .await + .unwrap() + .unwrap(); + let route = instance + .partition_manager() + .find_physical_table_route(table.table_info().table_id()) + .await + .unwrap(); + assert_eq!(route.region_routes.len(), 1); + sequences.push( + standalone + .mito_engine + .get_committed_sequence(route.region_routes[0].region.id) + .await + .unwrap(), + ); + } + sequences + } + + async fn check_trace_aux_async_completion() { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_async_completion") + .with_table_batcher(BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + // V2 bulk ingestion has a separate Variant/Struct conversion failure. + for version in 0..=1 { + let table_name = format!("trace_aux_async_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(true); + let ctx = Arc::new(context); + let request = trace_aux_request(); + let write = async |request, ctx| { + let outcome = tokio::time::timeout( + Duration::from_secs(30), + instance.traces( + instance.clone(), + request, + if version == 0 { + PipelineWay::OtlpTraceDirectV0 + } else { + PipelineWay::OtlpTraceDirectV1 + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx, + ), + ) + .await + .expect("async trace ingestion must return before the batch flush") + .unwrap(); + assert_eq!((outcome.accepted_spans, outcome.rejected_spans), (1, 0)); + assert!(outcome.error_message.is_none()); + }; + write(request.clone(), ctx.clone()).await; + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + for suffix in ["_services", "_operations"] { + assert!( + instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + "public", + &format!("{table_name}{suffix}"), + None + ) + .await + .unwrap() + .is_none() + ); + } + + // The second submission deterministically triggers the main flush. + // Omitting the service avoids a second concurrent auxiliary write. + let mut flush_request = request.clone(); + flush_request.resource_spans[0] + .resource + .as_mut() + .unwrap() + .attributes + .clear(); + write(flush_request, ctx.clone()).await; + tokio::time::timeout(Duration::from_secs(30), async { + loop { + let mut aux_data = TraceAuxData::default(); + for span in span::parse(request.clone()) + .into_iter() + .flat_map(|group| group.spans) + { + aux_data.observe_span(&span); + } + if instance.trace_aux_data_cached_for_test(&table_name, aux_data, &ctx) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("deferred auxiliary writes must populate the cache"); + for suffix in ["_services", "_operations"] { + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) = 1 FROM {table_name}{suffix}"), + ctx.clone(), + ) + .await + .unwrap(); + } + assert_trace_v2_query( + instance, + &format!("SELECT COUNT(*) > 0 FROM {table_name}"), + ctx.clone(), + ) + .await + .unwrap(); + + let before = trace_aux_sequences(&standalone, &table_name).await; + // A direct warm request finishes before we compare storage sequences. + let mut warm_ctx = ctx.fork(); + warm_ctx.set_batching_enabled(false); + write(request, Arc::new(warm_ctx)).await; + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + } + } + + async fn check_trace_aux_cache(batching: bool) { + let standalone = GreptimeDbStandaloneBuilder::new("trace_aux_cache") + .with_table_batcher(if batching { + BatcherOptions { + protocols: vec![BatchingProtocol::Otlp], + // Two main spans trigger a flush, while each auxiliary table + // has only one row and would wait an hour if batched. + pending_rows_flush_interval: Duration::from_secs(3600), + max_batch_rows: 2, + ..Default::default() + } + } else { + BatcherOptions::default() + }) + .build() + .await; + let instance = standalone.fe_instance(); + let cloned_instance = Arc::new(instance.as_ref().clone()); + + // V2 bulk writes currently fail on Variant/Struct schema conversion + // before reaching auxiliary ingestion; cover V2 with direct main writes. + let max_version = if batching { 1 } else { 2 }; + for version in 0..=max_version { + let table_name = format!("trace_aux_v{version}"); + let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + context.set_extension( + common_catalog::consts::TRACE_TABLE_NAME_SESSION_KEY, + &table_name, + ); + context.set_extension(table::requests::TRACE_TABLE_PARTITIONS_HINT_KEY, "1"); + context.set_batching_enabled(batching); + let ctx = Arc::new(context); + let mut request = trace_aux_request(); + if batching { + let spans = &mut request.resource_spans[0].scope_spans[0].spans; + spans.push(spans[0].clone()); + } + let accepted_spans = if batching { 2 } else { 1 }; + let write = async |frontend: &Arc, request| { + tokio::time::timeout( + Duration::from_secs(30), + frontend.traces( + frontend.clone(), + request, + match version { + 0 => PipelineWay::OtlpTraceDirectV0, + 1 => PipelineWay::OtlpTraceDirectV1, + _ => PipelineWay::OtlpTraceDirectV2, + }, + GreptimePipelineParams::default(), + table_name.clone(), + ctx.clone(), + ), + ) + .await + .expect("auxiliary writes must not wait for the batch flush interval") + .unwrap() + }; + + // An incompatible auxiliary column fails after the main span is + // accepted. No candidate key may be cached from the failed batch. + instance + .do_query( + &format!( + "CREATE TABLE {table_name}_operations (\ + \"timestamp\" TIMESTAMP(9) TIME INDEX, service_name STRING, \ + span_name STRING, span_kind BIGINT)" + ), + ctx.clone(), + ) + .await + .remove(0) + .unwrap(); + let failed_aux = write(instance, request.clone()).await; + assert_eq!( + (failed_aux.accepted_spans, failed_aux.rejected_spans), + (accepted_spans, 0) + ); + assert!( + failed_aux + .error_message + .unwrap() + .contains("Auxiliary trace tables") + ); + + instance + .do_query(&format!("DROP TABLE {table_name}_operations"), ctx.clone()) + .await + .remove(0) + .unwrap(); + let retry = write(instance, request.clone()).await; + assert_eq!( + (retry.accepted_spans, retry.rejected_spans), + (accepted_spans, 0) + ); + assert!(retry.error_message.is_none(), "{retry:?}"); + let before = trace_aux_sequences(&standalone, &table_name).await; + assert!(before.iter().all(|sequence| *sequence > 0)); + + // Storage sequences, unlike row counts in an upsert table, prove + // the writes were skipped. Cloned Instances must share the cache. + let warm = write(&cloned_instance, request.clone()).await; + assert_eq!( + (warm.accepted_spans, warm.rejected_spans), + (accepted_spans, 0) + ); + assert!(warm.error_message.is_none(), "{warm:?}"); + let after = trace_aux_sequences(&standalone, &table_name).await; + assert!(after[0] > before[0]); + assert_eq!(after[1..], before[1..], "v{version}"); + + request.resource_spans[0].scope_spans[0].spans[0].name = "new_operation".into(); + let new_operation = write(instance, request).await; + assert_eq!( + (new_operation.accepted_spans, new_operation.rejected_spans), + (accepted_spans, 0) + ); + assert!(new_operation.error_message.is_none(), "{new_operation:?}"); + let new_sequences = trace_aux_sequences(&standalone, &table_name).await; + assert!(new_sequences[0] > after[0]); + assert_eq!(new_sequences[1], after[1]); + assert!(new_sequences[2] > after[2]); + } + } + #[tokio::test(flavor = "multi_thread")] async fn test_trace_v2_on_standalone() -> Result<(), Box> { let standalone = GreptimeDbStandaloneBuilder::new("trace_v2_standalone") diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 54a4e2d67b8..b3293bf6421 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -3523,6 +3523,7 @@ fn drop_lines_with_inconsistent_results(input: String) -> String { "disable_ec2_metadata =", "cache_path =", "cache_capacity =", + "trace_aux_cache_size =", "memory_pool_size =", "scan_memory_limit =", "sas_token =", From d9bdded741fba4c513966a05ba2d80853bb6a74f Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:33:16 +0000 Subject: [PATCH 129/133] chore: open middleware (#9502) Signed-off-by: shuiyisong --- src/servers/src/http.rs | 1 + src/servers/src/http/memory_limit.rs | 10 ++++++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 4555b85ac69..2653ceddbe5 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -112,6 +112,7 @@ mod timeout; pub mod utils; mod workload_scheduler; +pub use memory_limit::decoded_body_accounting_middleware; use result::HttpOutputWriter; pub(crate) use timeout::DynamicTimeoutLayer; diff --git a/src/servers/src/http/memory_limit.rs b/src/servers/src/http/memory_limit.rs index 8d7b99c15fa..41905616f44 100644 --- a/src/servers/src/http/memory_limit.rs +++ b/src/servers/src/http/memory_limit.rs @@ -127,16 +127,18 @@ fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response response } -/// Route-local counterpart of [`memory_limit_middleware`] for routes that -/// install `RequestDecompressionLayer`: it must be layered *inside* the +/// Route-local accounting for routes that install `RequestDecompressionLayer`: +/// it must be layered *inside* the /// decompression layer, so the body it wraps is the **decompressed** stream. +/// It requires the global accounting installed by [`crate::http::HttpServer`] +/// and must use the same [`ServerMemoryLimiter`]. /// -/// Only requests marked [`ContentEncoded`] by the global middleware are +/// Only encoded requests marked by the global middleware are /// accounted; plain (uncompressed) requests are already covered by the global /// wire-byte accounting, and charging them here as well would double-count. /// For decompressed requests the `Content-Length` header has been stripped by /// the decompression layer, so the full decoded size is charged as it streams. -pub(crate) async fn decoded_body_accounting_middleware( +pub async fn decoded_body_accounting_middleware( State(limiter): State, req: Request, next: Next, From a8e293fd799438ea839df7e2af27fcd82e247b55 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sat, 10 Oct 2026 05:26:21 +0000 Subject: [PATCH 130/133] fix(promql): accept evaluation times before the Unix epoch (#9480) * fix(promql): accept evaluation times before the Unix epoch Prometheus accepts negative timestamps in query parameters; GreptimeDB rejected them while parsing and would panic converting them in the planner. Signed-off-by: Dennis Zhuang * test: update flow_tql error for a negative TQL end time A negative end now parses and fails the end >= start check instead of the timestamp conversion. Signed-off-by: Dennis Zhuang * fix(promql): reject evaluation times outside the millisecond range The signed millisecond conversion cast u128 to i64 unchecked, so a far-past time wrapped into the future. Reuse the checked conversion from label_values and return an error instead. Signed-off-by: Dennis Zhuang * chore: drop an unused import Signed-off-by: Dennis Zhuang * fix(promql): round query times to milliseconds like Prometheus Numeric times round the fraction of a second to milliseconds, and pre-epoch times floor sub-millisecond parts, matching Prometheus' parseTime and timestamp.FromTime. Signed-off-by: Dennis Zhuang * fix(promql): round the float fraction before converting to nanoseconds Signed-off-by: Dennis Zhuang * test: skip the overflow case where SystemTime cannot hold it Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/parser.rs | 55 +++++++++++++++++-- src/query/src/promql/label_values.rs | 15 +++-- src/query/src/promql/planner.rs | 15 ++--- src/query/src/promql/planner/test.rs | 34 +++++++++++- .../distributed/flow-tql/flow_tql.result | 2 +- .../cases/standalone/flow-tql/flow_tql.result | 2 +- 6 files changed, 102 insertions(+), 21 deletions(-) diff --git a/src/query/src/parser.rs b/src/query/src/parser.rs index 87b827e0cc1..c89c47392f6 100644 --- a/src/query/src/parser.rs +++ b/src/query/src/parser.rs @@ -224,11 +224,31 @@ impl QueryLanguageParser { // also report rfc3339 error if float parsing fails .map_err(|_| rfc3339_result.unwrap_err())?; - let duration = - Duration::try_from_secs_f64(secs).context(TryIntoDurationSnafu { raw: timestamp })?; - SystemTime::UNIX_EPOCH - .checked_add(duration) - .context(AddSystemTimeOverflowSnafu { duration }) + // Prometheus accepts timestamps before the Unix epoch. Follow the arithmetic of its + // `parseTime` and `timestamp.FromTime` step by step: the float fraction is rounded to + // milliseconds before any conversion to nanoseconds, which decides boundary values + // like 1.0005, whose fraction is just below 0.0005. + let magnitude = Duration::try_from_secs_f64(secs.abs()) + .context(TryIntoDurationSnafu { raw: timestamp })?; + let whole = secs.trunc(); + let fraction = ((secs - whole) * 1000.0).round() / 1000.0; + let nanos = whole as i128 * 1_000_000_000 + (fraction * 1e9) as i128; + let millis = nanos.div_euclid(1_000_000); + let duration = u64::try_from(millis.unsigned_abs()) + .ok() + .map(Duration::from_millis) + .context(AddSystemTimeOverflowSnafu { + duration: magnitude, + })?; + if millis < 0 { + SystemTime::UNIX_EPOCH + .checked_sub(duration) + .context(AddSystemTimeOverflowSnafu { duration }) + } else { + SystemTime::UNIX_EPOCH + .checked_add(duration) + .context(AddSystemTimeOverflowSnafu { duration }) + } } } @@ -380,6 +400,31 @@ mod test { } } + #[test] + fn parse_promql_timestamp_matches_prometheus_millis() { + // Expected values are Prometheus' `timestamp.FromTime(parseTime(input))`. + for (input, millis) in [ + ("-30000.5", -30_000_500), + ("-0.001", -1), + ("-0.0006", -1), + ("-1.2346", -1235), + ("1.2346", 1235), + ("0.0004", 0), + ("1.0005", 1000), + ("-1.0005", -1000), + ("0.0004999999", 0), + ("1969-12-31T23:59:59.9999Z", -1), + ("1970-01-01T00:00:00.0009Z", 0), + ] { + let result = QueryLanguageParser::parse_promql_timestamp(input).unwrap(); + assert_eq!( + crate::promql::label_values::signed_millis_since_epoch(result).unwrap(), + millis, + "{input}" + ); + } + } + #[test] fn parse_promql_timestamp() { let cases = vec![ diff --git a/src/query/src/promql/label_values.rs b/src/query/src/promql/label_values.rs index 841666758e4..cae53c64853 100644 --- a/src/query/src/promql/label_values.rs +++ b/src/query/src/promql/label_values.rs @@ -33,13 +33,20 @@ use crate::promql::error::{ /// a negative duration, and an RFC3339 query parameter can name one, so the sign /// is recovered here instead of unwrapping. fn millis_since_epoch(time: SystemTime) -> Result { + Ok(Timestamp::new_millisecond(signed_millis_since_epoch(time)?)) +} + +/// Converts a [`SystemTime`] to signed milliseconds since the Unix epoch, rejecting instants +/// outside the `i64` millisecond range. +/// +/// Sub-millisecond parts round down, also before the epoch, like Prometheus' +/// `timestamp.FromTime`. +pub(crate) fn signed_millis_since_epoch(time: SystemTime) -> Result { let (millis, before_epoch) = match time.duration_since(UNIX_EPOCH) { Ok(duration) => (duration.as_millis(), false), - Err(earlier) => (earlier.duration().as_millis(), true), + Err(earlier) => (earlier.duration().as_nanos().div_ceil(1_000_000), true), }; - let millis = - signed_millis(millis, before_epoch).with_context(|| SystemTimeOutOfRangeSnafu { time })?; - Ok(Timestamp::new_millisecond(millis)) + signed_millis(millis, before_epoch).with_context(|| SystemTimeOutOfRangeSnafu { time }) } /// Converts a millisecond magnitude and epoch direction without platform time limits. diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index a0ace9b9db6..5d8fc50edbd 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -21,7 +21,6 @@ mod set_operator; use std::collections::{BTreeSet, HashSet, VecDeque}; use std::sync::Arc; -use std::time::UNIX_EPOCH; use arrow::datatypes::IntervalDayTime; use async_recursion::async_recursion; @@ -112,6 +111,7 @@ use crate::promql::error::{ UnknownTableSnafu, UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, }; +use crate::promql::label_values::signed_millis_since_epoch; use crate::query_engine::QueryEngineState; /// `time()` function in PromQL. @@ -241,10 +241,11 @@ impl BinaryResultLabels { } impl PromPlannerContext { - fn from_eval_stmt(stmt: &EvalStmt) -> Self { - let start = stmt.start.duration_since(UNIX_EPOCH).unwrap().as_millis() as Millisecond; - let end = stmt.end.duration_since(UNIX_EPOCH).unwrap().as_millis() as Millisecond; - Self { + fn from_eval_stmt(stmt: &EvalStmt) -> Result { + // Evaluation times before the Unix epoch are valid in Prometheus. + let start = signed_millis_since_epoch(stmt.start)?; + let end = signed_millis_since_epoch(stmt.end)?; + Ok(Self { start, end, interval: stmt.interval.as_millis() as _, @@ -252,7 +253,7 @@ impl PromPlannerContext { stmt_start: start, stmt_end: end, ..Default::default() - } + }) } /// Reset all planner states @@ -309,7 +310,7 @@ impl PromPlanner { ) -> Result { let mut planner = Self { table_provider, - ctx: PromPlannerContext::from_eval_stmt(stmt), + ctx: PromPlannerContext::from_eval_stmt(stmt)?, promql_annotations, }; diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 642255532bd..19de1098288 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::collections::HashMap; -use std::time::{Duration, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use catalog::RegisterTableRequest; use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; @@ -3297,7 +3297,7 @@ async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { .await; let mut planner = PromPlanner { table_provider, - ctx: PromPlannerContext::from_eval_stmt(&eval_stmt), + ctx: PromPlannerContext::from_eval_stmt(&eval_stmt).unwrap(), promql_annotations: None, }; let plan = planner @@ -6266,7 +6266,7 @@ async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { .await; let mut planner = PromPlanner { table_provider, - ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), + ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")).unwrap(), promql_annotations: None, }; planner.ctx.time_index_column = Some("timestamp".to_string()); @@ -8872,3 +8872,31 @@ async fn test_count_values_groups_by_formatted_value_for_bigint_input() { vec![("9007199254740992", 2.0)] ); } + +#[tokio::test] +async fn eval_times_before_epoch() { + let plan = |start: SystemTime, end: SystemTime| async move { + let mut eval_stmt = build_eval_stmt("some_metric"); + eval_stmt.start = start; + eval_stmt.end = end; + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()).await + }; + + plan( + UNIX_EPOCH - Duration::from_secs(30_000), + UNIX_EPOCH - Duration::from_secs(29_000), + ) + .await + .unwrap(); + // Beyond the i64 millisecond range, which must not wrap into the future. Windows cannot + // represent this SystemTime. + if let Some(too_early) = UNIX_EPOCH.checked_sub(Duration::from_secs(10_000_000_000_000_000)) { + assert!(plan(too_early, too_early).await.is_err()); + } +} diff --git a/tests/cases/distributed/flow-tql/flow_tql.result b/tests/cases/distributed/flow-tql/flow_tql.result index 30158a176f8..73321d040e1 100644 --- a/tests/cases/distributed/flow-tql/flow_tql.result +++ b/tests/cases/distributed/flow-tql/flow_tql.result @@ -155,7 +155,7 @@ Error: 2000(InvalidSyntax), Invalid TQL syntax: sql parser error: Expected expre CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-(now()+'15s'::interval), '5s') count_values("status_code", http_requests); -Error: 3001(EngineExecuteQuery), Failed to convert float seconds to duration, raw: -15 +Error: 3001(EngineExecuteQuery), The end time must be greater than start time, start: 0, end: -15000 CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-now()+'15s'::interval, '5s') count_values("status_code", http_requests); diff --git a/tests/cases/standalone/flow-tql/flow_tql.result b/tests/cases/standalone/flow-tql/flow_tql.result index 13239bf4e1f..73321d040e1 100644 --- a/tests/cases/standalone/flow-tql/flow_tql.result +++ b/tests/cases/standalone/flow-tql/flow_tql.result @@ -155,7 +155,7 @@ Error: 2000(InvalidSyntax), Invalid TQL syntax: sql parser error: Expected expre CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-(now()+'15s'::interval), '5s') count_values("status_code", http_requests); -Error: 3001(EngineExecuteQuery), Failed to convert float seconds to duration, raw: -15: cannot convert float seconds to Duration: value is negative +Error: 3001(EngineExecuteQuery), The end time must be greater than start time, start: 0, end: -15000 CREATE FLOW calc_reqs SINK TO cnt_reqs EVAL INTERVAL '1m' AS TQL EVAL (now() - now(), now()-now()+'15s'::interval, '5s') count_values("status_code", http_requests); From 178a0021b8041db96647c5dc5a959178218abea3 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 06:45:11 +0000 Subject: [PATCH 131/133] fix: unify unknown flow option rejection across OSS and enterprise builds (#9319) * fix(operator): unify unknown flow option rejection across OSS and enterprise builds Remove the enterprise-gated passthrough that let unknown flow options through SQL-layer validation. Both OSS and enterprise binaries now reject unknown flow options with a single 1004(InvalidArguments) error at the SQL layer, producing identical sqlness output. The metasrv-side validate_flow_options remains as a defensive check. Signed-off-by: discord9 * fix(flow): use DISTINCT output columns as sink key for streaming flows When a streaming flow's source table has no primary key, sink key inference produced an empty key. For SELECT DISTINCT v, the DISTINCT establishes v as the output key, but the inference only looked at source primary keys. Combined with the zero-valued placeholder timestamp, all output rows shared the same storage identity and were deduplicated into a single row. Use all non-time-index output columns as the sink key when the plan is a DISTINCT and no source primary key was inherited. Signed-off-by: discord9 * fix(flow): preserve extension options and complete DISTINCT sink identity Keep enterprise CREATE FLOW options available to the existing handler and report unknown ordinary options as InvalidArguments with the OSS SQL error text. Infer DISTINCT sink keys from every selected non-time-index output column, including when source key columns survive. Wait for asynchronous mirrors in show_create_flow and exercise 1000 distinct sink rows in the SQL regression. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): align meta unknown option status assertions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): use valid aliases in DISTINCT sink regression SQL Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(flow): record complete DISTINCT sink SQL results Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/meta/src/ddl/create_flow.rs | 6 +- src/common/meta/src/ddl/tests/create_flow.rs | 29 +++- src/common/meta/src/error.rs | 8 ++ src/flow/src/adapter.rs | 19 ++- src/flow/src/adapter/tests.rs | 129 +++++++++++++++++- .../common/flow/flow_advance_ttl.result | 107 +++++++++++++++ .../common/flow/flow_advance_ttl.sql | 55 ++++++++ 7 files changed, 341 insertions(+), 12 deletions(-) diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index e36c89fbb34..89ade1b5cbe 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -42,7 +42,7 @@ use crate::cache_invalidator::Context; use crate::ddl::DdlContext; use crate::ddl::event::flow::{CREATE_FLOW_EVENT_TYPE, CreateFlowEventIntent, FlowDdlEvent}; use crate::ddl::utils::{add_peer_context_if_needed, map_to_procedure_error}; -use crate::error::{self, Result, UnexpectedSnafu}; +use crate::error::{self, InvalidFlowOptionSnafu, Result, UnexpectedSnafu}; use crate::instruction::{CacheIdent, CreateFlow, DropFlow}; use crate::key::flow::flow_info::{FlowInfoValue, FlowScheduleConfig, FlowStatus}; use crate::key::flow::flow_route::FlowRouteValue; @@ -558,9 +558,9 @@ pub fn validate_flow_options(flow_task: &CreateFlowTask) -> Result<()> { | FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY | FlowType::FLOW_TYPE_KEY => {} unknown => { - return UnexpectedSnafu { + return InvalidFlowOptionSnafu { err_msg: format!( - "Unknown flow option '{unknown}', supported user options: {DEFER_ON_MISSING_SOURCE_KEY}, {FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY}" + "unknown flow option '{unknown}', supported options: {DEFER_ON_MISSING_SOURCE_KEY}, {FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY}" ), } .fail(); diff --git a/src/common/meta/src/ddl/tests/create_flow.rs b/src/common/meta/src/ddl/tests/create_flow.rs index e7004656fe9..ad334e29863 100644 --- a/src/common/meta/src/ddl/tests/create_flow.rs +++ b/src/common/meta/src/ddl/tests/create_flow.rs @@ -18,6 +18,8 @@ use std::sync::Arc; use api::v1::flow::CreateRequest; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; use common_procedure::{Procedure, Status}; use common_procedure_test::execute_procedure_until_done; use table::table_name::TableName; @@ -326,6 +328,26 @@ fn test_validate_flow_options_allows_incremental_read_option() { validate_flow_options(&task).unwrap(); } +#[test] +fn test_validate_flow_options_unknown_option_matches_invalid_sql_error() { + let mut task = test_create_flow_task( + "my_flow", + vec![], + TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "my_sink_table"), + false, + ); + task.flow_options + .insert("foo".to_string(), "sensitive-value".to_string()); + + let err = validate_flow_options(&task).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); + assert_eq!( + "Invalid SQL, error: unknown flow option 'foo', supported options: defer_on_missing_source, experimental_enable_incremental_read", + err.output_msg() + ); + assert!(!err.output_msg().contains("sensitive-value")); +} + #[test] fn test_validate_flow_options_rejects_schedule_and_internal_keys_as_unknown() { for key in [ @@ -346,9 +368,10 @@ fn test_validate_flow_options_rejects_schedule_and_internal_keys_as_unknown() { .insert(key.to_string(), "value".to_string()); let err = validate_flow_options(&task).unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); assert!( err.to_string() - .contains(&format!("Unknown flow option '{key}'")), + .contains(&format!("unknown flow option '{key}'")), "unexpected error for {key}: {err}" ); } @@ -884,10 +907,10 @@ async fn test_create_flow_rejects_unknown_option_in_meta_task() { let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); - assert_matches!(err, error::Error::Unexpected { .. }); + assert_eq!(StatusCode::InvalidArguments, err.status_code()); assert!( err.to_string() - .contains("Unknown flow option 'unknown_option'") + .contains("unknown flow option 'unknown_option'") ); } diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 62559fd6f1b..5c1febb9b4d 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -515,6 +515,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid SQL, error: {err_msg}"))] + InvalidFlowOption { + err_msg: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid flow request body: {:?}", body))] InvalidFlowRequestBody { body: Box>, @@ -1247,6 +1254,7 @@ impl ErrorExt for Error { ProcedureNotFound { .. } | InvalidViewInfo { .. } + | InvalidFlowOption { .. } | PrimaryKeyNotFound { .. } | EmptyKey { .. } | AlterLogicalTablesInvalidArguments { .. } diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index c98510bd7d5..5b15ad0a93c 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -141,8 +141,9 @@ fn relation_desc_from_output( columns: &[ColumnSchema], lineage: &[Option], source_primary_key_indices: &[usize], + plan: &LogicalPlan, ) -> RelationDesc { - let keys = source_primary_key_indices + let mut keys = source_primary_key_indices .iter() .filter_map(|source_index| { lineage @@ -150,6 +151,14 @@ fn relation_desc_from_output( .position(|index| index == &Some(*source_index)) }) .collect_vec(); + // DISTINCT defines identity from every output tuple; the source timestamp participates + // separately as the relation time index, so it is excluded from the tag key. + if matches!(plan, LogicalPlan::Distinct(Distinct::All(_))) { + keys = columns + .iter() + .positions(|column| !column.is_time_index()) + .collect(); + } let time_index = columns.iter().position(ColumnSchema::is_time_index); RelationDesc { typ: crate::repr::RelationType { @@ -849,8 +858,12 @@ impl StreamingEngine { let source_primary_key_indices = source_meta.primary_key_indices; stateless::validate_source_scan(&flow_plan, source_table_id, &source_schema)?; let (inferred_schema, lineage) = output_column_schemas(&flow_plan, &source_schema)?; - let inferred_relation = - relation_desc_from_output(&inferred_schema, &lineage, &source_primary_key_indices); + let inferred_relation = relation_desc_from_output( + &inferred_schema, + &lineage, + &source_primary_key_indices, + &flow_plan, + ); let sink_exists = self.fetch_table_pk_schema(sink_table_name).await?.is_some(); if !sink_exists { validate_auto_column_names(&inferred_schema)?; diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 6bc02f60786..6a2bf5ed816 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -20,7 +20,7 @@ use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as use datafusion::catalog::MemTable; use datafusion::datasource::provider_as_source; use datafusion_common::TableReference; -use datafusion_expr::LogicalPlanBuilder; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema, SchemaBuilder}; use store_api::storage::{ConcreteDataType, TableId}; use table::metadata::{TableInfo, TableInfoBuilder, TableMetaBuilder}; @@ -194,10 +194,133 @@ fn stateless_distinct_preserves_direct_column_lineage() { let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); assert_eq!(output[0].name, "dis"); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); } +#[test] +fn stateless_distinct_output_is_always_the_sink_key() { + let source = Arc::new(Schema::new(vec![ + ColumnSchema::new("v", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("k", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new("other", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + // All cases share one provider so the key assertions exercise only plan/output metadata. + let provider = Arc::new( + MemTable::try_new( + Arc::new(datafusion::arrow::datatypes::Schema::new(vec![ + Field::new("v", ArrowDataType::Int32, true), + Field::new("k", ArrowDataType::Int32, true), + Field::new("other", ArrowDataType::Int32, true), + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + ])), + vec![vec![]], + ) + .unwrap(), + ); + let make_plan = |expressions: Vec, distinct: bool| { + let builder = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(provider.clone()), + None, + ) + .unwrap() + .project(expressions) + .unwrap(); + if distinct { + builder.distinct().unwrap().build().unwrap() + } else { + builder.build().unwrap() + } + }; + let assert_key = |plan: &LogicalPlan, source_keys: &[usize], expected: &[usize]| { + let (output, lineage) = super::output_column_schemas(plan, &source).unwrap(); + let relation = super::relation_desc_from_output(&output, &lineage, source_keys, plan); + assert_eq!( + relation + .typ + .keys + .iter() + .map(|key| key.column_indices.clone()) + .collect::>(), + if expected.is_empty() { + vec![] + } else { + vec![expected.to_vec()] + } + ); + relation + }; + + // No source PK: all distinct output fields form the key. + assert_key(&make_plan(vec![datafusion_expr::col("v")], true), &[], &[0]); + // A surviving subset of a composite source PK cannot identify projected DISTINCT tuples. + assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + true, + ), + &[0, 1], + &[0, 1], + ); + // Even a retained source PK needs the additional selected fields in the DISTINCT key. + assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + true, + ), + &[0], + &[0, 1], + ); + // Aliases and computed output expressions are included by output position. + let alias_plan = make_plan( + vec![ + datafusion_expr::col("v").alias("renamed"), + (datafusion_expr::col("v") + datafusion_expr::col("other")).alias("computed"), + ], + true, + ); + let (output, lineage) = super::output_column_schemas(&alias_plan, &source).unwrap(); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &alias_plan); + assert_eq!( + output + .iter() + .map(|column| column.name.as_str()) + .collect_vec(), + ["renamed", "computed"] + ); + assert_eq!(relation.typ.keys[0].column_indices, vec![0, 1]); + + // A selected time index remains the relation time index, not part of the tag key. + let relation = assert_key( + &make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("ts")], + true, + ), + &[0], + &[0], + ); + assert_eq!(relation.typ.time_index, Some(1)); + + // Non-DISTINCT plans retain only inherited source-key lineage (or no key if none survives). + let ordinary = make_plan( + vec![datafusion_expr::col("v"), datafusion_expr::col("other")], + false, + ); + assert_key(&ordinary, &[0, 1], &[0]); + assert_key(&ordinary, &[], &[]); +} + #[test] fn stateless_normalizes_dictionary_output_type() { let field = Field::new_dictionary("host", ArrowDataType::UInt32, ArrowDataType::Utf8, true); @@ -220,7 +343,7 @@ fn stateless_normalizes_dictionary_output_type() { assert_eq!(output[0].data_type, ConcreteDataType::string_datatype()); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); } diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.result b/tests/cases/standalone/common/flow/flow_advance_ttl.result index eae65328a6b..27a28babebf 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.result +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.result @@ -414,3 +414,110 @@ DROP TABLE out_distinct_basic; Affected Rows: 0 +-- Streaming DISTINCT auto-sink retains all 1000 output tuples. +CREATE TABLE distinct_auto_sink ( + v INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX +) WITH ('ttl' = 'instant'); + +Affected Rows: 0 + +CREATE FLOW test_distinct_auto_sink SINK TO out_distinct_auto_sink AS +SELECT DISTINCT v FROM distinct_auto_sink; + +Affected Rows: 0 + +INSERT INTO distinct_auto_sink (v) +SELECT number FROM numbers LIMIT 1000; + +Affected Rows: 0 + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink'); + ++---------------------------------------------+ +| ADMIN FLUSH_FLOW('test_distinct_auto_sink') | ++---------------------------------------------+ +| FLOW_FLUSHED | ++---------------------------------------------+ + +SELECT count(*) AS rows, count(DISTINCT v) AS distinct_values, min(v) AS min_v, max(v) AS max_v +FROM out_distinct_auto_sink; + ++------+-----------------+-------+-------+ +| rows | distinct_values | min_v | max_v | ++------+-----------------+-------+-------+ +| 1000 | 1000 | 0 | 999 | ++------+-----------------+-------+-------+ + +DROP FLOW test_distinct_auto_sink; + +Affected Rows: 0 + +DROP TABLE distinct_auto_sink; + +Affected Rows: 0 + +DROP TABLE out_distinct_auto_sink; + +Affected Rows: 0 + +CREATE TABLE distinct_auto_sink_pk ( + v INT, + k INT, + extra_value INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX, + PRIMARY KEY (v, k) +) WITH ('ttl' = 'instant'); + +Affected Rows: 0 + +CREATE FLOW test_distinct_auto_sink_pk SINK TO out_distinct_auto_sink_pk AS +SELECT DISTINCT v, extra_value FROM distinct_auto_sink_pk; + +Affected Rows: 0 + +INSERT INTO distinct_auto_sink_pk (v, k, extra_value) +SELECT number % 10 AS v, number AS k, number AS extra_value FROM numbers LIMIT 1000; + +Affected Rows: 0 + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk'); + ++------------------------------------------------+ +| ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk') | ++------------------------------------------------+ +| FLOW_FLUSHED | ++------------------------------------------------+ + +SELECT + count(*) AS rows, + count(DISTINCT extra_value) AS distinct_values, + min(extra_value) AS min_value, + max(extra_value) AS max_value, + sum(CASE WHEN v = extra_value % 10 THEN 0 ELSE 1 END) AS mismatched_keys +FROM out_distinct_auto_sink_pk; + ++------+-----------------+-----------+-----------+-----------------+ +| rows | distinct_values | min_value | max_value | mismatched_keys | ++------+-----------------+-----------+-----------+-----------------+ +| 1000 | 1000 | 0 | 999 | 0 | ++------+-----------------+-----------+-----------+-----------------+ + +DROP FLOW test_distinct_auto_sink_pk; + +Affected Rows: 0 + +DROP TABLE distinct_auto_sink_pk; + +Affected Rows: 0 + +DROP TABLE out_distinct_auto_sink_pk; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/flow/flow_advance_ttl.sql b/tests/cases/standalone/common/flow/flow_advance_ttl.sql index 2aacc0f9458..72356d96372 100644 --- a/tests/cases/standalone/common/flow/flow_advance_ttl.sql +++ b/tests/cases/standalone/common/flow/flow_advance_ttl.sql @@ -179,3 +179,58 @@ SELECT number FROM distinct_basic; DROP FLOW test_distinct_basic; DROP TABLE distinct_basic; DROP TABLE out_distinct_basic; + +-- Streaming DISTINCT auto-sink retains all 1000 output tuples. +CREATE TABLE distinct_auto_sink ( + v INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX +) WITH ('ttl' = 'instant'); + +CREATE FLOW test_distinct_auto_sink SINK TO out_distinct_auto_sink AS +SELECT DISTINCT v FROM distinct_auto_sink; + +INSERT INTO distinct_auto_sink (v) +SELECT number FROM numbers LIMIT 1000; + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink'); + +SELECT count(*) AS rows, count(DISTINCT v) AS distinct_values, min(v) AS min_v, max(v) AS max_v +FROM out_distinct_auto_sink; + +DROP FLOW test_distinct_auto_sink; +DROP TABLE distinct_auto_sink; +DROP TABLE out_distinct_auto_sink; + +CREATE TABLE distinct_auto_sink_pk ( + v INT, + k INT, + extra_value INT, + ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP TIME INDEX, + PRIMARY KEY (v, k) +) WITH ('ttl' = 'instant'); + +CREATE FLOW test_distinct_auto_sink_pk SINK TO out_distinct_auto_sink_pk AS +SELECT DISTINCT v, extra_value FROM distinct_auto_sink_pk; + +INSERT INTO distinct_auto_sink_pk (v, k, extra_value) +SELECT number % 10 AS v, number AS k, number AS extra_value FROM numbers LIMIT 1000; + +-- Mirror inserts reach the flownode asynchronously; wait before flushing. +-- SQLNESS SLEEP 3s +-- SQLNESS REPLACE (ADMIN\sFLUSH_FLOW\('\w+'\)\s+\|\n\+-+\+\n\|\s+)[0-9]+\s+\| $1 FLOW_FLUSHED | +ADMIN FLUSH_FLOW('test_distinct_auto_sink_pk'); + +SELECT + count(*) AS rows, + count(DISTINCT extra_value) AS distinct_values, + min(extra_value) AS min_value, + max(extra_value) AS max_value, + sum(CASE WHEN v = extra_value % 10 THEN 0 ELSE 1 END) AS mismatched_keys +FROM out_distinct_auto_sink_pk; + +DROP FLOW test_distinct_auto_sink_pk; +DROP TABLE distinct_auto_sink_pk; +DROP TABLE out_distinct_auto_sink_pk; From 28697c9e550aa8dba3b39bbc873c68a96d96b047 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 06:49:38 +0000 Subject: [PATCH 132/133] feat(datanode): execute nested MergeScan plans on the datanode (#9206) * feat(datanode): execute nested MergeScan plans on the datanode Give a datanode the ability to decode and execute a plan whose MergeScan carries a nested plan (e.g. an inner MergeScan over a build table), so a distributed sub-plan pushed to a datanode can itself fan out to other datanodes. This is the execution/codec capability only; it is not wired to any user-facing switch and produces no plan by itself. Codec (src/query/src/query_engine/default_serializer.rs): - MergeScanAwareSerializer keeps the request's SessionState and the engine catalog manager, and decodes a MergeScan payload's tables by (catalog, schema, table) name through the engine catalog rather than the region-aware request catalog list (which binds every table to the request's region and breaks on mismatched columns). Datanode (src/datanode): - New DatanodeRegionQueryHandler (region_query.rs) serves region queries for the MergeScan nodes of a received plan, resolving region leaders via PartitionRuleManager and dispatching to peer datanodes. - Wire KvBackendCatalogManager + PartitionRuleManager + NodeClients into the datanode's catalog/cache registry (with_dist_planner stays false: a datanode never rewrites its own plan). Bug fixes uncovered while wiring this up: - cache registry: register the table/table_info/table_name/table_route/ partition_info caches the datanode's catalog manager and partition rule manager require (previously missing -> panic on lookup). - default_serializer: decode_sub_plan bridge no longer drives a future on a current-thread runtime from within another executor (spawn a dedicated thread + runtime when not on a multi-thread runtime). Tested by query_engine::default_serializer unit tests: nested MergeScan round-trip and engine-catalog resolution with mismatched column names. No user-facing behavior change. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep GreptimeDB function bindings after rebuilding decode state SessionStateBuilder::new_from_existing collects the function registry (HashMap) into a Vec and build() re-registers each entry; register_udf registers aliases before the function's own name and the Vec order is unstable. Rebuilding the session state to install MergeScanAwareSerializer after the GreptimeDB functions were registered could therefore rebind a name/alias (e.g. date_format, an alias of DataFusion's to_char) back to the built-in function. This path runs for every decoded plan, not only nested MergeScan ones. Extract the GreptimeDB function registration into register_greptime_functions and re-apply it after the final build of both the top-level decode state and the MergeScan payload state. Add regression tests asserting date_format stays bound to the GreptimeDB implementation on both paths. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): address review on nested MergeScan capability Three follow-ups from review of the nested MergeScan execution capability: 1. Layer the datanode caches (P2b). table_cache derives from table_info_cache + table_name_cache, and partition_info_cache derives from table_route_cache, but they were all registered in a single flat CacheRegistry whose invalidate() runs concurrently with no ordering. A derived cache could be invalidated and refilled from a base cache that had not been invalidated yet, leaving stale metadata. Build the base caches and the derived caches as separate layers of a LayeredCacheRegistry (reusing the existing layering mechanism) so base caches finish invalidating before derived ones. Both the production path (cmd) and the test scaffolding now build the same layered registry; the flat build_datanode_cache_registry is removed. 2. Don't let a nested MergeScan's internal region query compete for the concurrency permit (P2a). An outer query holds its permit until its stream is fully consumed, and a nested MergeScan makes the datanode issue region queries to peer datanodes as execution stages of that same query. If such an inner stage also took a permit, the outer query waiting on it would deadlock against its own held permit (or time out when max_concurrent_queries > 0). Mark inner DN->DN region queries as internal stages: QueryRequest gains an flag carried over the wire in the request header's query_context extensions (reserved key, not injectable via hints), and the region server skips acquiring a concurrency permit for internal-stage requests on both acquire sites. Outer queries still acquire normally. 3. Add an end-to-end capability test (tests-integration/.../nested_merge_scan_capability_test.rs): hand-build Join(, MergeScan(build regions on both datanodes)) with mismatched column names and duplicate build keys, ship it to a datanode over a real TCP Flight endpoint, consume the stream fully and compare the row multiset against the frontend's reference join result, assert a real cross-datanode DoGet via an RPC counter, and assert an inner region failure fails the whole query. Also migrate the with_real_datanode_grpc_addr + DatanodeRpcStats test scaffolding. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): adapt date_format regression test to the updated DataFusion fork ScalarUDFImpl's trait bounds changed so its impl can no longer be downcast via as_any(); assert the function equals the GreptimeDB implementation instead, which to_char can never satisfy. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover the internal stage under the concurrency limit Add an end-to-end regression test for the nested MergeScan capability under max_concurrent_queries = 1: both datanodes' only permit is held by the outer queries' streams while their inner MergeScans fetch the build regions from the same datanodes over real Flight connections. The result consumption is bounded by a timeout, so a broken link anywhere in the chain (marking the inner request internal, carrying the marker in the request header, honoring it at the receiving acquire sites) fails the test instead of deadlocking. Also add a minimal with_datanode_options_override hook to the cluster builder, and fix the test module doc: the outer request is handed to RegionServer::handle_remote_read in process; only the inner region queries go over the wire. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover cold route loads under unrelated invalidations The table route cache uses InitStrategy::VersionChecked, whose version counter belongs to the whole CacheContainer rather than to a single table: any TableId invalidation bumps the shared version, so a cold load of one table's route retries when an unrelated table is invalidated concurrently. Document this as a known limitation on the region query handler, and pin the behavior with an integration test: run the nested plans while a background task invalidates an unrelated table id on every datanode, assert the results still match the expected multiset (retries only add work, they never return a stale route), and report the observed load amplification via the cache miss counter (2 loads without invalidations, 8 over three runs with 24 unrelated invalidations in this run). Also expose the datanode cache registries on the test cluster so a test can drive the local invalidation path that the heartbeat handler drives after a metasrv broadcast. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(query): trim nested plan decoder commentary Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor: remove redundant nested scan support plumbing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): decode nested MergeScan payloads asynchronously Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): avoid unnecessary extension name allocation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): retain extension decoding for manual session states Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(datanode): keep nested scan query admission unchanged Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): honor query handler factories and trim nested scan fixtures Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): retain required cache error import Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): reuse registered state for nested scan decoding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): verify payload information schema table functions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): retain nested scan runtime and cold-route guards Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): address nested scan review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(session): preserve read preference across region query hops Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(session): verify strict context rejection at Flow ingress Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): remove redundant serializer error case Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 4 + src/cache/src/lib.rs | 65 +- src/cmd/src/datanode/builder.rs | 11 +- src/common/substrait/src/df_substrait.rs | 19 +- src/common/substrait/src/lib.rs | 1 + src/datanode/Cargo.toml | 5 +- src/datanode/src/datanode.rs | 96 +- src/datanode/src/error.rs | 10 +- src/datanode/src/lib.rs | 1 + src/datanode/src/region_query.rs | 109 ++ src/datanode/src/region_server.rs | 22 +- src/flow/src/adapter/flownode_impl.rs | 11 +- .../src/query_engine/default_serializer.rs | 1371 +++++++++++++++-- src/servers/src/grpc/greptime_handler.rs | 9 +- src/servers/src/http/hints.rs | 9 +- src/session/Cargo.toml | 3 + src/session/src/context.rs | 176 ++- src/session/src/hints.rs | 5 +- tests-integration/src/cluster.rs | 237 ++- tests-integration/src/standalone.rs | 9 +- tests-integration/src/tests.rs | 1 + .../nested_merge_scan_capability_test.rs | 749 +++++++++ tests/compatibility/README.md | 13 +- .../wire/read_preference/README.md | 45 + .../read_preference/legacy_query_context.pb | 3 + 25 files changed, 2790 insertions(+), 194 deletions(-) create mode 100644 src/datanode/src/region_query.rs create mode 100644 tests-integration/src/tests/nested_merge_scan_capability_test.rs create mode 100644 tests/compatibility/wire/read_preference/README.md create mode 100644 tests/compatibility/wire/read_preference/legacy_query_context.pb diff --git a/Cargo.lock b/Cargo.lock index 9b447169fbc..055bf712394 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4473,6 +4473,8 @@ dependencies = [ "async-trait", "bytes", "cache", + "catalog", + "client", "common-base", "common-catalog", "common-config", @@ -4480,6 +4482,7 @@ dependencies = [ "common-error", "common-function", "common-greptimedb-telemetry", + "common-grpc", "common-macro", "common-meta", "common-options", @@ -13663,6 +13666,7 @@ dependencies = [ "datafusion-common", "derive_builder 0.20.2", "derive_more", + "prost 0.14.1", "snafu 0.8.6", "sql", "uuid", diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index 1a234aa802e..c7304ae6009 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -20,8 +20,8 @@ use std::time::Duration; use catalog::kvbackend::new_table_cache; use common_meta::cache::{ - CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_schema_cache, - new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, + CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder, + new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, new_table_route_cache, new_table_schema_cache, new_view_info_cache, }; use common_meta::kv_backend::KvBackendRef; @@ -53,11 +53,11 @@ pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache"; pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache"; pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; -/// Builds cache registry for datanode, including: -/// - Schema cache. -/// - Table id to schema name cache. -pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry { - // Builds table id schema name cache that never expires. +/// Builds the datanode cache registry with base metadata caches before derived caches. +/// +/// Invalidating the base layer first prevents a concurrent query from refilling a derived cache +/// from stale metadata after that derived cache has been invalidated. +pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry { let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build(); let table_id_schema_cache = Arc::new(new_table_schema_cache( TABLE_SCHEMA_NAME_CACHE_NAME.to_string(), @@ -65,7 +65,6 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - // Builds schema cache let cache = default_cache(); let schema_cache = Arc::new(new_schema_cache( SCHEMA_CACHE_NAME.to_string(), @@ -73,9 +72,57 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - CacheRegistryBuilder::default() + let cache = default_cache(); + let table_info_cache = Arc::new(new_table_info_cache( + TABLE_INFO_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_name_cache = Arc::new(new_table_name_cache( + TABLE_NAME_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_route_cache = Arc::new(new_table_route_cache( + TABLE_ROUTE_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_cache = Arc::new(new_table_cache( + TABLE_CACHE_NAME.to_string(), + cache, + table_info_cache.clone(), + table_name_cache.clone(), + )); + + let cache = default_cache(); + let partition_info_cache = Arc::new(new_partition_info_cache( + PARTITION_INFO_CACHE_NAME.to_string(), + cache, + table_route_cache.clone(), + )); + + let base_registry = CacheRegistryBuilder::default() .add_cache(table_id_schema_cache) .add_cache(schema_cache) + .add_cache(table_info_cache) + .add_cache(table_name_cache) + .add_cache(table_route_cache) + .build(); + let derived_registry = CacheRegistryBuilder::default() + .add_cache(table_cache) + .add_cache(partition_info_cache) + .build(); + + LayeredCacheRegistryBuilder::default() + .add_cache_registry(base_registry) + .add_cache_registry(derived_registry) .build() } diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index 810dfeedc7b..6df7ac728e8 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/src/cmd/src/datanode/builder.rs @@ -14,10 +14,9 @@ use std::sync::Arc; -use cache::build_datanode_cache_registry; +use cache::build_datanode_layered_cache_registry; use catalog::kvbackend::new_read_only_meta_kv_backend; use common_base::Plugins; -use common_meta::cache::LayeredCacheRegistryBuilder; use common_telemetry::info; use common_version::{short_version, verbose_version}; use datanode::datanode::DatanodeBuilder; @@ -111,11 +110,9 @@ impl InstanceBuilder { let backend = new_read_only_meta_kv_backend(client.clone()); let mut builder = DatanodeBuilder::new(dn_opts.clone(), plugins.clone(), backend.clone()); - let registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(backend)) - .build(), - ); + // The derived caches (table/partition caches) are put into the layer after the caches + // they are derived from, so they are invalidated only after them. + let registry = Arc::new(build_datanode_layered_cache_registry(backend)); builder .with_cache_registry(registry) .with_meta_client(client.clone()); diff --git a/src/common/substrait/src/df_substrait.rs b/src/common/substrait/src/df_substrait.rs index 08eeedd8316..9570076ee7a 100644 --- a/src/common/substrait/src/df_substrait.rs +++ b/src/common/substrait/src/df_substrait.rs @@ -25,7 +25,9 @@ use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::optimizer::AnalyzerRule; use datafusion::prelude::SessionConfig; use datafusion_expr::LogicalPlan; -use datafusion_substrait::logical_plan::consumer::from_substrait_plan; +use datafusion_substrait::logical_plan::consumer::{ + SubstraitConsumer, from_substrait_plan, from_substrait_plan_with_consumer, +}; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::substrait::proto::Plan; use prost::Message; @@ -74,6 +76,21 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor { } impl DFLogicalSubstraitConvertor { + /// Decode with a custom consumer, retaining Greptime's state-aggregate ordering correction. + pub async fn decode_with_consumer( + &self, + plan: &Plan, + state: &SessionState, + consumer: &impl SubstraitConsumer, + ) -> Result { + let df_plan = from_substrait_plan_with_consumer(consumer, plan) + .await + .context(DecodeDfPlanSnafu)?; + FixStateUdafOrderingAnalyzer {} + .analyze(df_plan, state.config_options()) + .context(DecodeDfPlanSnafu) + } + pub fn to_sub_plan( &self, plan: &LogicalPlan, diff --git a/src/common/substrait/src/lib.rs b/src/common/substrait/src/lib.rs index 3eeb3c570e1..a6932451a62 100644 --- a/src/common/substrait/src/lib.rs +++ b/src/common/substrait/src/lib.rs @@ -20,6 +20,7 @@ use async_trait::async_trait; use bytes::{Buf, Bytes}; use datafusion::execution::context::SessionState; pub use datafusion::execution::registry::SerializerRegistry; +pub use datafusion_substrait::extensions::Extensions; /// Re-export the Substrait module of datafusion, pub use datafusion_substrait::substrait as substrait_proto_df; pub use datafusion_substrait::{logical_plan as df_logical_plan, variation_const}; diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 7c126764723..4d4e470e665 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -16,6 +16,8 @@ api.workspace = true arrow-flight.workspace = true async-trait.workspace = true bytes.workspace = true +catalog.workspace = true +client.workspace = true common-base.workspace = true common-catalog.workspace = true common-config.workspace = true @@ -23,6 +25,7 @@ common-datasource.workspace = true common-error.workspace = true common-function.workspace = true common-greptimedb-telemetry.workspace = true +common-grpc.workspace = true common-macro.workspace = true common-meta.workspace = true common-options.workspace = true @@ -54,6 +57,7 @@ metric-engine.workspace = true mito2.workspace = true num_cpus.workspace = true object-store.workspace = true +partition.workspace = true prometheus.workspace = true prost.workspace = true query.workspace = true @@ -78,5 +82,4 @@ common-query.workspace = true common-test-util.workspace = true datafusion-common.workspace = true mito2 = { workspace = true, features = ["test"] } -partition.workspace = true session.workspace = true diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index fdac0597572..cfbcf948b96 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -19,17 +19,22 @@ use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; +use catalog::information_schema::NoopInformationExtension; +use catalog::kvbackend::KvBackendCatalogManagerBuilder; +use client::client_manager::NodeClients; use common_base::Plugins; use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE}; use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; +use common_grpc::channel_manager::ChannelConfig; use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::datanode::TopicStatsReporter; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef}; use common_meta::kv_backend::KvBackendRef; +use common_meta::node_manager::NodeManagerRef; pub use common_procedure::options::ProcedureConfig; use common_query::prelude::set_default_prefix; use common_stat::ResourceStatImpl; @@ -54,7 +59,8 @@ use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; use query::QueryEngineFactory; -use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef}; +use query::dummy_catalog::TableProviderFactoryRef; +use query::region_query::RegionQueryHandlerFactoryRef; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::LogStore; @@ -80,6 +86,7 @@ use crate::event_listener::{ use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::heartbeat::HeartbeatTask; use crate::partition_expr_fetcher::MetaPartitionExprFetcher; +use crate::region_query::DatanodeRegionQueryHandler; use crate::region_server::{DummyTableProviderFactory, RegionServer}; use crate::store::{self, new_object_store_without_cache}; use crate::utils::{RegionOpenRequests, build_region_open_requests}; @@ -357,6 +364,7 @@ impl DatanodeBuilder { schema_metadata_manager, region_event_listener, file_ref_manager, + cache_registry.clone(), ) .await?; @@ -497,14 +505,41 @@ impl DatanodeBuilder { schema_metadata_manager: SchemaMetadataManagerRef, event_listener: RegionServerEventListenerRef, file_ref_manager: FileReferenceManagerRef, + cache_registry: Arc, ) -> Result { let opts: &DatanodeOptions = &self.opts; + // Keep `with_dist_planner` false: datanodes execute received MergeScan plans. + let catalog_manager = KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + self.kv_backend.clone(), + cache_registry.clone(), + ) + .build(); + let partition_manager = catalog_manager.partition_manager(); + + // MergeScan region requests stream until completion, so they have no client timeout. + let mut channel_config = ChannelConfig { + timeout: None, + ..Default::default() + }; + if opts.grpc.flight_compression.transport_compression() { + channel_config.accept_compression = true; + channel_config.send_compression = true; + } + let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config)); + let region_query_handler = + if let Some(factory) = self.plugins.get::() { + factory.build(partition_manager.clone(), node_manager) + } else { + DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager) + }; + let query_engine_factory = QueryEngineFactory::try_new_with_plugins( // query engine in datanode only executes plan with resolved table source. - DummyCatalogManager::arc(), - None, - None, + catalog_manager, + Some(partition_manager), + Some(region_query_handler), None, None, None, @@ -1060,17 +1095,17 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; - use cache::build_datanode_cache_registry; + use cache::build_datanode_layered_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; - use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::key::RegionRoleSet; use common_meta::key::datanode_table::DatanodeTableManager; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::node_manager::NodeManagerRef; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; @@ -1080,6 +1115,8 @@ mod tests { use mito2::engine::MITO_ENGINE_NAME; use object_store::ObjectStore; use object_store::manager::ObjectStoreManager; + use partition::manager::PartitionRuleManagerRef; + use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef}; use store_api::logstore::LogStore; use store_api::region_request::RegionRequest; use store_api::storage::RegionId; @@ -1090,6 +1127,7 @@ mod tests { wal_object_store, }; use crate::error::{self, Error}; + use crate::region_query::DatanodeRegionQueryHandler; use crate::tests::{MockRegionEngine, mock_region_server}; async fn setup_table_datanode(kv: &KvBackendRef) { @@ -1118,11 +1156,8 @@ mod tests { mock_region_server.register_engine(mock_region.clone()); let kv_backend = Arc::new(MemoryKvBackend::new()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new( DatanodeOptions { @@ -1164,11 +1199,8 @@ mod tests { wal: DatanodeWalConfig, kv_backend: KvBackendRef, ) -> DatanodeBuilder { - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut opts = DatanodeOptions { node_id: Some(0), wal, @@ -1446,6 +1478,38 @@ mod tests { ); } + struct CountingRegionQueryHandlerFactory(Arc); + + impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory { + fn build( + &self, + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> query::region_query::RegionQueryHandlerRef { + self.0.fetch_add(1, Ordering::SeqCst); + DatanodeRegionQueryHandler::arc(partition_manager, node_manager) + } + } + + #[tokio::test] + async fn test_build_uses_installed_region_query_handler_factory() { + let data_home = create_temp_dir("region-query-handler-factory"); + let builds = Arc::new(AtomicUsize::new(0)); + let plugins = Plugins::new(); + plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone())) + as RegionQueryHandlerFactoryRef); + let mut builder = datanode_builder( + data_home.path().to_str().unwrap(), + DatanodeWalConfig::default(), + Arc::new(MemoryKvBackend::new()), + ); + builder.set_plugins(plugins); + + let mut datanode = builder.build().await.unwrap(); + assert_eq!(builds.load(Ordering::SeqCst), 1); + datanode.shutdown().await.unwrap(); + } + #[tokio::test] async fn test_build_and_shutdown_with_raft_engine_wal() { let data_home = create_temp_dir("raft-engine-wal"); diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 539712711b7..e9b5d5ec1dc 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -145,6 +145,13 @@ pub enum Error { source: Box, }, + #[snafu(display("Invalid query context: {}", reason))] + InvalidQueryContext { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid SQL, error: {}", msg))] InvalidSql { msg: String }, @@ -485,7 +492,8 @@ impl ErrorExt for Error { Delete { source, .. } => source.status_code(), - InvalidSql { .. } + InvalidQueryContext { .. } + | InvalidSql { .. } | IllegalPrimaryKeysDef { .. } | MissingTimestampColumn { .. } | SchemaNotFound { .. } diff --git a/src/datanode/src/lib.rs b/src/datanode/src/lib.rs index c14ec886fac..5a8948e2c84 100644 --- a/src/datanode/src/lib.rs +++ b/src/datanode/src/lib.rs @@ -22,6 +22,7 @@ pub mod heartbeat; pub mod metrics; mod partition_expr_fetcher; pub mod query_stream; +pub mod region_query; pub mod region_server; pub mod service; pub mod store; diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs new file mode 100644 index 00000000000..1642362ee1a --- /dev/null +++ b/src/datanode/src/region_query.rs @@ -0,0 +1,109 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The datanode's [`RegionQueryHandler`] implementation for querying remote regions in a +//! `MergeScan` plan. + +use std::sync::Arc; + +use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; +use async_trait::async_trait; +use common_error::ext::BoxedError; +use common_meta::node_manager::NodeManagerRef; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManagerRef; +use query::error::{RegionQuerySnafu, Result as QueryResult, UnimplementedSnafu}; +use query::region_query::{RegionQueryHandler, RegionQueryTarget}; +use session::ReadPreference; +use snafu::ResultExt; +use store_api::storage::RegionId; + +/// Serves region queries by resolving their leader and forwarding requests to its datanode. +pub(crate) struct DatanodeRegionQueryHandler { + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, +} + +impl DatanodeRegionQueryHandler { + pub fn arc( + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> Arc { + Arc::new(Self { + partition_manager, + node_manager, + }) + } +} + +#[async_trait] +impl RegionQueryHandler for DatanodeRegionQueryHandler { + async fn select_target( + &self, + _read_preference: ReadPreference, + region_id: RegionId, + ) -> QueryResult { + // The leader of a region is the only peer that can serve a region query. + // + let peer = self + .partition_manager + .find_region_leader(region_id) + .await + .map_err(BoxedError::new) + .context(RegionQuerySnafu)?; + + Ok(RegionQueryTarget::new(peer)) + } + + async fn do_get( + &self, + target: &RegionQueryTarget, + request: QueryRequest, + ) -> QueryResult { + self.node_manager + .datanode(target.peer()) + .await + .handle_query(request) + .await + .map_err(BoxedError::new) + .context(RegionQuerySnafu) + } + + /// Datanodes forward received region queries and do not produce remote dynamic filters. + async fn handle_remote_dyn_filter_update( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _update: RemoteDynFilterUpdate, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter update from a datanode", + } + .fail() + } + + /// See [`RegionQueryHandler::handle_remote_dyn_filter_update`]. + async fn handle_remote_dyn_filter_unregister( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _unregister: RemoteDynFilterUnregister, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter unregister from a datanode", + } + .fail() + } +} diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 7d5218f30b7..ef333cff04d 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -354,7 +354,17 @@ impl RegionServer { None }; - let ctx = request.header.as_ref().map(|h| h.into()); + let ctx = request + .header + .as_ref() + .map(QueryContext::try_from) + .transpose() + .map_err(|err| { + crate::error::InvalidQueryContextSnafu { + reason: err.to_string(), + } + .build() + })?; let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build())); let region_id = request.region_id; @@ -977,7 +987,15 @@ impl FlightCraft for RegionServer { let query_ctx = request .header .as_ref() - .map(|h| Arc::new(QueryContext::from(h))) + .map(QueryContext::try_from) + .transpose() + .map_err(|err| { + servers_error::InvalidParameterSnafu { + reason: err.to_string(), + } + .build() + })? + .map(Arc::new) .unwrap_or(QueryContext::arc()); let region_server = self.clone(); diff --git a/src/flow/src/adapter/flownode_impl.rs b/src/flow/src/adapter/flownode_impl.rs index 62e8597dbb6..cf4b6fd2d34 100644 --- a/src/flow/src/adapter/flownode_impl.rs +++ b/src/flow/src/adapter/flownode_impl.rs @@ -826,7 +826,16 @@ impl common_meta::node_manager::Flownode for FlowDualEngine { let query_ctx = request .header .and_then(|h| h.query_context) - .map(|ctx| ctx.into()); + .map(session::context::QueryContext::try_from) + .transpose() + .map_err(|err| { + to_meta_err(snafu::location!())( + crate::error::InvalidQuerySnafu { + reason: err.to_string(), + } + .build(), + ) + })?; match request.body { Some(flow_request::Body::Create(CreateRequest { flow_id: Some(task_id), diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index f4d6b29fa2c..c924a57a68e 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -14,19 +14,21 @@ use std::sync::Arc; +use bytes::Bytes; +use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; use common_query::error::RegisterUdfSnafu; use common_query::logical_plan::SubstraitPlanDecoder; -use datafusion::catalog::CatalogProviderList; -use datafusion::common::DataFusionError; +use datafusion::catalog::{CatalogProviderList, TableProvider}; +use datafusion::common::{DFSchema, DataFusionError, TableReference, not_impl_err, substrait_err}; use datafusion::error::Result; use datafusion::execution::context::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::execution::{FunctionRegistry, SessionStateBuilder}; -use datafusion::logical_expr::LogicalPlan; -use datafusion_expr::UserDefinedLogicalNode; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion_expr::{Expr, UserDefinedLogicalNode}; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, @@ -44,12 +46,18 @@ use promql::functions::{ Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; -use session::context::QueryContextRef; +use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; +use substrait::df_logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer}; +use substrait::error::{DecodeDfPlanSnafu, DecodeRelSnafu}; use substrait::extension_serializer::ExtensionSerializer; -use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use substrait::substrait_proto_df::proto::{ + ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, Plan, Type, +}; +use substrait::{DFLogicalSubstraitConvertor, Extensions, SubstraitPlan}; use crate::dist_plan::MergeScanLogicalPlan; +use crate::query_engine::QueryEngineState; /// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization. #[derive(Debug)] @@ -86,8 +94,7 @@ impl SerializerRegistry for DefaultSerializer { bytes: &[u8], ) -> Result> { if name == MergeScanLogicalPlan::name() { - // TODO(dennis): missing `session_state` to decode the logical plan in `MergeScanLogicalPlan`, - // so we only save the unoptimized logical plan for view currently. + // This registry has no session state; the query engine decodes `MergeScan` payloads. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -97,6 +104,196 @@ impl SerializerRegistry for DefaultSerializer { } } +/// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer. +struct MergeScanSubstraitConsumer<'a> { + inner: DefaultSubstraitConsumer<'a>, + session_state: &'a SessionState, + catalog_manager: Option, + payload: bool, +} + +impl MergeScanSubstraitConsumer<'_> { + fn query_ctx(&self) -> QueryContextRef { + self.session_state + .config() + .get_extension() + .unwrap_or_else(QueryContext::arc) + } + + async fn decode_payload(&self, payload: Vec) -> Result { + decode_plan(Bytes::from(payload), self.session_state, true) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + } +} + +#[async_trait::async_trait] +impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { + async fn resolve_table_ref( + &self, + table_ref: &TableReference, + ) -> Result>> { + if !self.payload { + return self.inner.resolve_table_ref(table_ref).await; + } + + let query_ctx = self.query_ctx(); + let resolved = table_ref + .clone() + .resolve(query_ctx.current_catalog(), &query_ctx.current_schema()); + // Manual states retain DataFusion's own catalog and information-schema behavior. + let Some(manager) = &self.catalog_manager else { + return self + .inner + .resolve_table_ref(&TableReference::from(resolved)) + .await; + }; + let catalog_list: Arc = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager.clone(), + query_ctx, + ), + ); + let schema_name = resolved.schema.as_ref(); + let catalog_name = resolved.catalog.as_ref(); + let schema = catalog_list + .catalog(catalog_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}")) + })? + .schema(schema_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve schema: {schema_name}")) + })?; + schema.table(&resolved.table).await + } + + fn get_extensions(&self) -> &Extensions { + self.inner.get_extensions() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + self.inner.get_function_registry() + } + + fn push_outer_schema(&self, schema: Arc) { + self.inner.push_outer_schema(schema) + } + + fn pop_outer_schema(&self) { + self.inner.pop_outer_schema() + } + + fn get_outer_schema(&self, steps_out: usize) -> Option> { + self.inner.get_outer_schema(steps_out) + } + + fn push_lambda_parameters( + &self, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> Result> { + self.inner + .push_lambda_parameters(lambda_parameters, input_schema) + } + + fn pop_lambda_parameters(&self) { + self.inner.pop_lambda_parameters() + } + + fn lambda_variable(&self, steps_out: usize, field_idx: usize) -> Result { + self.inner.lambda_variable(steps_out, field_idx) + } + + async fn consume_extension_leaf(&self, rel: &ExtensionLeafRel) -> Result { + // Unknown and detail-less leaves keep the default consumer's errors. + let Some(detail) = rel.detail.as_ref() else { + return self.inner.consume_extension_leaf(rel).await; + }; + if detail.type_url != MergeScanLogicalPlan::name() { + return self.inner.consume_extension_leaf(rel).await; + } + + let merge_scan = PbMergeScan::decode(detail.value.as_ref()).map_err(|e| { + DataFusionError::Substrait(format!("Failed to decode the MergeScan plan node: {e}")) + })?; + + let input = self.decode_payload(merge_scan.input).await?; + + // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. + Ok( + MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default()) + .into_logical_plan(), + ) + } + + // Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding. + async fn consume_extension_single(&self, rel: &ExtensionSingleRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let Some(input_rel) = &rel.input else { + return substrait_err!( + "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" + ); + }; + let input_plan = self.consume_rel(input_rel).await?; + let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + } + + // Multi-input extension children likewise need this consumer's dispatch. + async fn consume_extension_multi(&self, rel: &ExtensionMultiRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let mut inputs = Vec::with_capacity(rel.inputs.len()); + for input in &rel.inputs { + let input_plan = self.consume_rel(input).await?; + inputs.push(input_plan); + } + let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + } +} + +/// Decode with async MergeScan handling, preserving the default consumer's extension checks. +async fn decode_plan( + bytes: Bytes, + state: &SessionState, + payload: bool, +) -> substrait::error::Result { + let plan = Plan::decode(bytes).context(DecodeRelSnafu)?; + let extensions = Extensions::try_from(&plan.extensions).context(DecodeDfPlanSnafu)?; + if !extensions.type_variations.is_empty() { + return not_impl_err!("Type variation extensions are not supported") + .context(DecodeDfPlanSnafu); + } + + let catalog_manager = state + .config() + .get_extension::() + .map(|engine_state| engine_state.catalog_manager().clone()); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, state), + session_state: state, + catalog_manager, + payload, + }; + + DFLogicalSubstraitConvertor + .decode_with_consumer(&plan, state, &consumer) + .await +} + /// The datafusion `[LogicalPlan]` decoder. pub struct DefaultPlanDecoder { session_state: SessionState, @@ -115,6 +312,108 @@ impl DefaultPlanDecoder { } } +/// Registers Greptime functions after building the request decoder state. +/// +/// DataFusion aliases can rebind Greptime `date_format` to built-in `to_char`; registration must +/// happen after `SessionStateBuilder::build` on the state that decodes the request plan. +fn register_greptime_functions( + session_state: &mut SessionState, + query_ctx: &QueryContextRef, +) -> common_query::error::Result<()> { + for func in FUNCTION_REGISTRY.scalar_functions() { + let udf = func.provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }); + session_state + .register_udf(Arc::new(udf)) + .context(RegisterUdfSnafu { name: func.name() })?; + } + + for func in FUNCTION_REGISTRY.aggregate_functions() { + let name = func.name().to_string(); + session_state + .register_udaf(Arc::new(func)) + .context(RegisterUdfSnafu { name })?; + } + + let _ = session_state.register_udaf(quantile_udaf()); + + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); + let double_exponential_smoothing_udf = + DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); + let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); + + for udf in [ + NativeHistogramAbsentOverTime::scalar_udf(), + NativeHistogramAdd::scalar_udf(), + NativeHistogramAvg::scalar_udf(), + NativeHistogramAvgOverTime::scalar_udf(), + NativeHistogramChanges::scalar_udf(), + NativeHistogramCount::scalar_udf(), + NativeHistogramCountOverTime::scalar_udf(), + NativeHistogramDelta::scalar_udf(), + NativeHistogramDivScalar::scalar_udf(), + NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), + NativeHistogramDrop::float_null_udf(String::new(), None), + NativeHistogramEq::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramIDelta::scalar_udf(), + NativeHistogramIRate::scalar_udf(), + NativeHistogramIncrease::scalar_udf(), + NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), + NativeHistogramMulScalar::scalar_udf(), + NativeHistogramNeg::scalar_udf(), + NativeHistogramNotEq::scalar_udf(), + NativeHistogramPresentOverTime::scalar_udf(), + NativeHistogramQuantile::scalar_udf(), + NativeHistogramRate::scalar_udf(), + NativeHistogramResets::scalar_udf(), + NativeHistogramScalarMul::scalar_udf(), + NativeHistogramStddev::scalar_udf(), + NativeHistogramStdvar::scalar_udf(), + NativeHistogramSub::scalar_udf(), + NativeHistogramSum::scalar_udf(), + NativeHistogramSumOverTime::scalar_udf(), + NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), + ] { + let _ = session_state.register_udf(Arc::new(udf)); + } + + for udaf in [ + NativeHistogramAggAvg::aggregate_udf(), + NativeHistogramAggSum::aggregate_udf(), + ] { + let _ = session_state.register_udaf(Arc::new(udaf)); + } + + Ok(()) +} + #[async_trait::async_trait] impl SubstraitPlanDecoder for DefaultPlanDecoder { async fn decode( @@ -123,108 +422,14 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_serializer_registry(Arc::new(DefaultSerializer)) .with_catalog_list(catalog_list) .build(); - // Substrait decoder will look up the UDFs in SessionState, so we need to register them - // Note: the query context must be passed to set the timezone - // We MUST register the UDFs after we build the session state, otherwise the UDFs will be lost - // if they have the same name as the default UDFs or their alias. - // e.g. The default UDF `to_char()` has an alias `date_format()`, if we register a UDF with the name `date_format()` - // before we build the session state, the UDF will be lost. - for func in FUNCTION_REGISTRY.scalar_functions() { - let udf = func.provide(FunctionContext { - query_ctx: self.query_ctx.clone(), - state: Default::default(), - }); - session_state - .register_udf(Arc::new(udf)) - .context(RegisterUdfSnafu { name: func.name() })?; - } + // Re-register after the build to avoid Greptime UDF alias collisions. + register_greptime_functions(&mut session_state, &self.query_ctx)?; - for func in FUNCTION_REGISTRY.aggregate_functions() { - let name = func.name().to_string(); - session_state - .register_udaf(Arc::new(func)) - .context(RegisterUdfSnafu { name })?; - } - - let _ = session_state.register_udaf(quantile_udaf()); - - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); - let double_exponential_smoothing_udf = - DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); - let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); - - for udf in [ - NativeHistogramAbsentOverTime::scalar_udf(), - NativeHistogramAdd::scalar_udf(), - NativeHistogramAvg::scalar_udf(), - NativeHistogramAvgOverTime::scalar_udf(), - NativeHistogramChanges::scalar_udf(), - NativeHistogramCount::scalar_udf(), - NativeHistogramCountOverTime::scalar_udf(), - NativeHistogramDelta::scalar_udf(), - NativeHistogramDivScalar::scalar_udf(), - NativeHistogramDrop::bool_false_udf(String::new(), None), - NativeHistogramDrop::bool_true_udf(String::new(), None), - NativeHistogramDrop::float_null_udf(String::new(), None), - NativeHistogramEq::scalar_udf(), - NativeHistogramFraction::scalar_udf(), - NativeHistogramIDelta::scalar_udf(), - NativeHistogramIRate::scalar_udf(), - NativeHistogramIncrease::scalar_udf(), - NativeHistogramLastOverTime::scalar_udf(), - MixedRange::float_udf(None), - MixedRange::histogram_udf(None), - NativeHistogramMulScalar::scalar_udf(), - NativeHistogramNeg::scalar_udf(), - NativeHistogramNotEq::scalar_udf(), - NativeHistogramPresentOverTime::scalar_udf(), - NativeHistogramQuantile::scalar_udf(), - NativeHistogramRate::scalar_udf(), - NativeHistogramResets::scalar_udf(), - NativeHistogramScalarMul::scalar_udf(), - NativeHistogramStddev::scalar_udf(), - NativeHistogramStdvar::scalar_udf(), - NativeHistogramSub::scalar_udf(), - NativeHistogramSum::scalar_udf(), - NativeHistogramSumOverTime::scalar_udf(), - NativeHistogramToString::scalar_udf(), - PromqlFloatToString::scalar_udf(), - ] { - let _ = session_state.register_udf(Arc::new(udf)); - } - for udaf in [ - NativeHistogramAggAvg::aggregate_udf(), - NativeHistogramAggSum::aggregate_udf(), - ] { - let _ = session_state.register_udaf(Arc::new(udaf)); - } - - let logical_plan = DFLogicalSubstraitConvertor - .decode(message, session_state) + let logical_plan = decode_plan(message, &session_state, false) .await .map_err(BoxedError::new) .context(common_query::error::DecodePlanSnafu)?; @@ -241,18 +446,35 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, Ordering}; + + use catalog::error::Result as CatalogResult; + use catalog::{CatalogManager, RegisterTableRequest}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; + use common_function::aggrs::aggr_wrapper::StateWrapper; + use common_function::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer; + use common_query::logical_plan::SubstraitPlanDecoderRef; use common_query::native_histogram::native_histogram_value_type; + use common_time::Timezone; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; use datafusion::logical_expr::Extension; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit}; + use datafusion::optimizer::AnalyzerRule; + use datafusion_expr::expr::{AggregateFunction, Cast, InSubquery, ScalarFunction, Sort}; + use datafusion_expr::{ + AggregateUDF, Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, Subquery, col, lit, + }; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, }; use datatypes::data_type::DataType; - use promql::extension_plan::RangeManipulate; + use futures::stream::BoxStream; + use promql::extension_plan::{RangeManipulate, SeriesNormalize, UnionDistinctOn}; use session::context::QueryContext; + use table::TableRef; + use table::metadata::{TableId, TableInfoRef}; + use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use super::*; use crate::QueryEngineFactory; @@ -488,4 +710,935 @@ mod tests { assert!(decoded.contains("prom_mixed_range_float")); assert!(decoded.contains("prom_mixed_range_histogram")); } + + #[tokio::test] + async fn test_serializer_decode_merge_scan() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan(numbers_scan(), true); + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + assert!(merge_scan.is_placeholder()); + assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string()); + // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. + assert!(merge_scan.partition_cols().is_empty()); + } + + /// Payload table binding uses the engine catalog, not the request catalog. + #[tokio::test] + async fn test_serializer_decode_merge_scan_with_engine_catalog() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let decoded = plan_decoder(engine.as_ref()) + .decode( + encode_plan(&merge_scan(numbers_scan(), false)), + request_catalog_list(), + false, + ) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + let LogicalPlan::TableScan(scan) = merge_scan.input() else { + panic!("Expect a table scan, got: {}", merge_scan.input()); + }; + assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME); + assert_eq!(scan.source.schema().fields().len(), 1); + assert_eq!(scan.source.schema().field(0).name(), "number"); + } + + /// Nested `MergeScan` payloads decode recursively. + #[tokio::test] + async fn test_serializer_decode_nested_merge_scan() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let inner = merge_scan(numbers_scan(), false); + let input = LogicalPlanBuilder::from(inner) + .filter(col("number").lt(lit(10u32))) + .unwrap() + .build() + .unwrap(); + let plan = merge_scan(input, false); + + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } + + fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc { + FUNCTION_REGISTRY + .get_function("date_format") + .expect("`date_format` is a function of GreptimeDB") + .provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }) + .into() + } + + /// Uses a non-default timezone to distinguish the Greptime implementation. + fn date_format_query_ctx() -> QueryContextRef { + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); + Arc::new(query_ctx) + } + + fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr { + Expr::ScalarFunction(ScalarFunction { + func: greptime_date_format_udf(query_ctx), + args: vec![timestamp, lit("%Y-%m-%d %H:%i:%S")], + }) + } + + /// Includes `MergeScan` payload expressions, which are not node inputs. + fn scalar_functions_of(plan: &LogicalPlan) -> Vec> { + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + + let mut udfs = Vec::new(); + for expr in plan.expressions() { + expr.apply(|expr| { + if let Expr::ScalarFunction(func) = expr { + udfs.push(Arc::clone(&func.func)); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + } + + if let LogicalPlan::Extension(extension) = plan + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + udfs.extend(scalar_functions_of(merge_scan.input())); + } + + for input in plan.inputs() { + udfs.extend(scalar_functions_of(input)); + } + + udfs + } + + /// Asserts the Greptime `date_format` binding rather than DataFusion's `to_char` alias. + fn assert_only_function_is_greptime_date_format( + plan: &LogicalPlan, + query_ctx: &QueryContextRef, + ) { + let udfs = scalar_functions_of(plan); + assert_eq!(udfs.len(), 1, "expected one function in: {plan}"); + + let udf = &udfs[0]; + assert_eq!(udf.name(), "date_format", "got: {plan}"); + assert_eq!( + udf.as_ref(), + greptime_date_format_udf(query_ctx).as_ref(), + "`date_format` is not bound to the implementation of GreptimeDB: {plan}" + ); + } + + /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases. + #[tokio::test] + async fn test_serializer_decode_keeps_greptime_date_format() { + // Repeat to expose nondeterministic function-registration order. + for _ in 0..8 { + let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); + let factory = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + let table_provider = Arc::new(mock_table_provider(1.into())); + let query_ctx = date_format_query_ctx(); + let plan = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(table_provider.schema().clone())), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) + .unwrap() + .build() + .unwrap(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap(); + + assert_only_function_is_greptime_date_format(&decoded, &query_ctx); + } + } + + /// Nested payload decoding retains the Greptime `date_format` binding. + #[tokio::test] + async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + + for _ in 0..8 { + let factory = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + let query_ctx = date_format_query_ctx(); + + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr( + &query_ctx, + Expr::Cast(Cast::new( + Box::new(col("number")), + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + )), + )]) + .unwrap() + .build() + .unwrap(); + let plan = + MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); + + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + + assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); + } + } + + /// Registers a `numbers` table in a memory catalog manager. + fn numbers_catalog_manager() -> CatalogManagerRef { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + catalog_manager + } + + /// Scan plan over the `numbers` schema shared by the fixtures below. + fn numbers_scan() -> LogicalPlan { + LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan { + MergeScanLogicalPlan::new(input, is_placeholder, Default::default()).into_logical_plan() + } + + fn encode_plan(plan: &LogicalPlan) -> Bytes { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .unwrap() + } + + fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef { + engine + .engine_context(QueryContext::arc()) + .new_plan_decoder() + .unwrap() + } + + /// Request catalog that resolves any table to a provider payload resolution must not use. + fn request_catalog_list() -> Arc { + Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))) + } + + /// Payload resolution preserves context and explicit catalog/schema components. + #[tokio::test] + async fn test_serializer_decode_payload_observes_catalog_and_query_context() { + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let inner = catalog::memory::new_memory_catalog_manager().unwrap(); + let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner, + resolution: TableResolution::ObserveContext(observed.clone()), + }); + let mut resolver_ctx = QueryContext::arc(); + Arc::make_mut(&mut resolver_ctx).set_current_catalog("request_cat"); + Arc::make_mut(&mut resolver_ctx).set_current_schema("request_schema"); + let engine = QueryEngineFactory::new( + manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let mut state = engine.engine_context(resolver_ctx.clone()).state().clone(); + state.config_mut().set_extension(resolver_ctx); + let extensions = Extensions::default(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: state + .config() + .get_extension::() + .map(|s| s.catalog_manager().clone()), + payload: true, + }; + for reference in [ + TableReference::bare(NUMBERS_TABLE_NAME), + TableReference::partial("explicit_schema", NUMBERS_TABLE_NAME), + TableReference::full( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + NUMBERS_TABLE_NAME, + ), + ] { + consumer + .resolve_table_ref(&reference) + .await + .unwrap() + .unwrap(); + } + assert_eq!( + *observed.lock().unwrap(), + vec![ + [ + "request_cat".to_string(), + "request_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + "request_cat".to_string(), + "explicit_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + ], + ); + } + + /// How [`InterceptingCatalogManager`] resolves payload tables. + enum TableResolution { + /// Records catalog lookup and request context values. + ObserveContext(Arc>>), + /// Resolves once another task on the caller runtime signals. + AwaitSignal { + notify: Arc, + entered: Arc, + }, + } + + /// Engine catalog manager that takes over `table` resolution while delegating the rest. + struct InterceptingCatalogManager { + inner: CatalogManagerRef, + resolution: TableResolution, + } + + #[async_trait::async_trait] + impl CatalogManager for InterceptingCatalogManager { + fn as_any(&self) -> &dyn Any { + self.inner.as_any() + } + + async fn catalog_names(&self) -> CatalogResult> { + self.inner.catalog_names().await + } + + async fn schema_names( + &self, + catalog: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.schema_names(catalog, query_ctx).await + } + + async fn table_names( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.table_names(catalog, schema, query_ctx).await + } + + async fn catalog_exists(&self, catalog: &str) -> CatalogResult { + self.inner.catalog_exists(catalog).await + } + + async fn schema_exists( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner.schema_exists(catalog, schema, query_ctx).await + } + + async fn table_exists( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner + .table_exists(catalog, schema, table, query_ctx) + .await + } + + async fn table( + &self, + catalog: &str, + schema: &str, + table_name: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + match &self.resolution { + TableResolution::ObserveContext(observed) => { + let query_ctx = query_ctx.unwrap(); + observed.lock().unwrap().push([ + catalog.to_string(), + schema.to_string(), + table_name.to_string(), + query_ctx.current_catalog().to_string(), + query_ctx.current_schema(), + ]); + Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID))) + } + + TableResolution::AwaitSignal { notify, entered } => { + entered.store(true, Ordering::SeqCst); + notify.notified().await; + self.inner + .table(catalog, schema, table_name, query_ctx) + .await + } + } + } + + async fn table_info_by_id(&self, table_id: TableId) -> CatalogResult> { + self.inner.table_info_by_id(table_id).await + } + + async fn tables_by_ids( + &self, + catalog: &str, + schema: &str, + table_ids: &[TableId], + ) -> CatalogResult> { + self.inner.tables_by_ids(catalog, schema, table_ids).await + } + + fn tables<'a>( + &'a self, + catalog: &'a str, + schema: &'a str, + query_ctx: Option<&'a QueryContext>, + ) -> BoxStream<'a, CatalogResult> { + self.inner.tables(catalog, schema, query_ctx) + } + } + + fn engine_with_intercepting_catalog( + resolution: TableResolution, + ) -> Arc { + let catalog_manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner: numbers_catalog_manager(), + resolution, + }); + QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + } + + /// A payload lookup parked on a caller-runtime task must keep the decode async. + /// + /// The engine catalog resolves only once a task spawned on this runtime signals it, so a + /// blocking bridge would deadlock the test instead of decoding. + #[tokio::test(flavor = "current_thread")] + async fn test_serializer_decode_merge_scan_awaits_caller_runtime_task() { + let notify = Arc::new(tokio::sync::Notify::new()); + let entered = Arc::new(AtomicBool::new(false)); + let engine = engine_with_intercepting_catalog(TableResolution::AwaitSignal { + notify: notify.clone(), + entered: entered.clone(), + }); + let plan = merge_scan(numbers_scan(), false); + let bytes = encode_plan(&plan); + + let decoder = plan_decoder(engine.as_ref()); + let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false)); + // First poll reaches the payload table lookup, which waits for the signal below. + assert!(futures::poll!(decode.as_mut()).is_pending()); + assert!(entered.load(Ordering::SeqCst)); + + let signal = tokio::spawn(async move { + notify.notify_waiters(); + }); + tokio::task::yield_now().await; + let decoded = decode.await.unwrap(); + signal.await.unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + } + + /// A missing engine-catalog table does not fall back to the usable request catalog. + #[tokio::test] + async fn test_serializer_decode_missing_engine_catalog_table() { + let missing_manager: CatalogManagerRef = + catalog::memory::new_memory_catalog_manager().unwrap(); + let engine = QueryEngineFactory::new( + missing_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Table not found: greptime.public.numbers"), + "unexpected error: {err}" + ); + } + + /// Manual states use their own catalog and query-context defaults for payloads. + #[tokio::test] + async fn test_serializer_decode_manual_state_keeps_extension_registry() { + let manager = catalog::memory::new_memory_catalog_manager().unwrap(); + manager.register_catalog_sync("manual_cat").unwrap(); + manager + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + }) + .unwrap(); + manager + .register_table_sync(RegisterTableRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); + let mut config = datafusion::execution::context::SessionConfig::new() + .with_information_schema(true) + .with_default_catalog_and_schema("manual_cat", "manual_schema"); + config.set_extension(query_ctx.clone()); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .build(); + let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + false, + Vec::new(), + merge_scan(numbers_scan(), false), + )), + }); + let request_catalog = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager, query_ctx, + ), + ); + let decoded = decoder + .decode(encode_plan(&plan), request_catalog, false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); + } + + /// `MergeScan` decodes under an ordinary single-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_series_normalize() { + let outer = merge_scan(numbers_scan(), false); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + false, + Vec::new(), + outer.clone(), + )), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a SeriesNormalize plan, got: {decoded}"); + }; + let normalize = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a SeriesNormalize plan node"); + assert_eq!(normalize.inputs()[0].to_string(), outer.to_string()); + } + + /// `MergeScan` decodes under an ordinary multi-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_union_distinct_on() { + let left = merge_scan(numbers_scan(), false); + let right = merge_scan(numbers_scan(), true); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new( + UnionDistinctOn::try_new(left.clone(), right.clone(), vec![], 0).unwrap(), + ), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a UnionDistinctOn plan, got: {decoded}"); + }; + let union = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a UnionDistinctOn plan node"); + assert_eq!(union.inputs()[0].to_string(), left.to_string()); + assert_eq!(union.inputs()[1].to_string(), right.to_string()); + } + + /// `MergeScan` inside a subquery still decodes through the wrapper consumer. + /// + /// Correlated outer references cannot be covered here: the in-repo Substrait producer + /// rejects `Expr::OuterReferenceColumn`, so no supported helper can build such a plan. + #[tokio::test] + async fn test_serializer_decode_merge_scan_inside_subquery() { + let inner = merge_scan(numbers_scan(), false); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .filter(Expr::InSubquery(InSubquery::new( + Box::new(col("number")), + Subquery { + subquery: Arc::new(inner.clone()), + outer_ref_columns: vec![], + spans: Default::default(), + }, + false, + ))) + .unwrap() + .build() + .unwrap(); + let plan = merge_scan(payload.clone(), false); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + assert!( + decoded.to_string().contains("IN ()"), + "expect the decoded subquery filter, got: {decoded}" + ); + } + + /// Recursive payload decoding must restore ordered state aggregates, including their schema. + #[tokio::test] + async fn test_serializer_decode_payload_keeps_state_udaf_ordering() { + let last_value = (*datafusion::functions_aggregate::first_last::last_value_udaf()).clone(); + let state_udaf = AggregateUDF::new_from_impl(StateWrapper::new(last_value).unwrap()); + let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::new(state_udaf), + vec![col("number")], + false, + None, + vec![Sort::new(col("number"), true, true)], + None, + )); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .aggregate(Vec::::new(), vec![aggr_expr]) + .unwrap() + .build() + .unwrap(); + let bytes = encode_plan(&payload); + let encoded = Plan::decode(bytes.clone()).unwrap(); + let extensions = Extensions::try_from(&encoded.extensions).unwrap(); + let catalog_manager = numbers_catalog_manager(); + let engine = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let mut state = SessionStateBuilder::new_from_existing( + engine.engine_context(QueryContext::arc()).state().clone(), + ) + .with_serializer_registry(Arc::new(DefaultSerializer)) + .build(); + register_greptime_functions(&mut state, &QueryContext::arc()).unwrap(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: Some(catalog_manager), + payload: false, + }; + let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap(); + + let fixed_payload = FixStateUdafOrderingAnalyzer {} + .analyze(payload.clone(), &Default::default()) + .unwrap(); + assert_ne!( + fixed_payload.schema().as_arrow(), + payload.schema().as_arrow() + ); + assert_eq!( + decoded.schema().as_arrow(), + fixed_payload.schema().as_arrow() + ); + let LogicalPlan::Aggregate(aggregate) = &decoded else { + panic!("Expected an aggregate payload, got: {decoded}"); + }; + let Expr::AggregateFunction(function) = &aggregate.aggr_expr[0] else { + panic!("Expected a state aggregate, got: {:?}", aggregate.aggr_expr); + }; + assert!(function.func.inner().is::()); + assert_eq!(function.params.order_by.len(), 1); + let ArrowDataType::Struct(fields) = decoded.schema().field(0).data_type() else { + panic!("Expected the ordered aggregate state type"); + }; + assert_eq!(fields.len(), 3); + assert_eq!(fields[1].name(), "number"); + assert!(fields[1].is_nullable()); + } + + /// Type variation extensions keep being rejected while decoding. + #[tokio::test] + async fn test_serializer_decode_rejects_type_variations() { + use substrait::substrait_proto_df::proto::extensions::SimpleExtensionDeclaration; + use substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::{ + ExtensionTypeVariation, MappingType, + }; + + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let encoded = encode_plan(&merge_scan(numbers_scan(), false)); + let mut substrait_plan = Plan::decode(encoded).unwrap(); + substrait_plan.extensions.push(SimpleExtensionDeclaration { + mapping_type: Some(MappingType::ExtensionTypeVariation( + ExtensionTypeVariation { + extension_urn_reference: u32::MAX, + type_variation_anchor: 1, + name: "u!variation".to_string(), + }, + )), + }); + let bytes = Bytes::from(substrait_plan.encode_to_vec()); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Type variation extensions are not supported"), + "unexpected error: {err}" + ); + } } diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad0a693ac59..c29ecc6493a 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -322,7 +322,8 @@ mod tests { use common_time::Timezone; use query::options::FLOW_SCHEDULED_TIME_MILLIS; use session::hints::{ - INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, + INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, }; use snafu::ResultExt; use tonic::Code; @@ -418,6 +419,10 @@ mod tests { vec![ ("auto_create_table".to_string(), "true".to_string()), ("read_preference".to_string(), "leader".to_string()), + ( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "follower".to_string(), + ), ( REMOTE_QUERY_ID_EXTENSION_KEY.to_string(), "spoofed".to_string(), @@ -446,6 +451,8 @@ mod tests { ReadPreference::Leader )); assert_eq!(query_context.extension("auto_create_table"), Some("true")); + assert_eq!(query_context.extension(READ_PREFERENCE_HINT), None); + assert_eq!(query_context.extension(READ_PREFERENCE_EXTENSION_KEY), None); assert_ne!(query_context.remote_query_id(), Some("spoofed")); assert!( query_context diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs index 09e4879dd50..7ef9401cf18 100644 --- a/src/servers/src/http/hints.rs +++ b/src/servers/src/http/hints.rs @@ -49,7 +49,8 @@ mod tests { use query::options::FLOW_SCHEDULED_TIME_MILLIS; use session::context::{QueryContextBuilder, generate_remote_query_id}; use session::hints::{ - INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, + INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, }; use super::apply_hints; @@ -75,6 +76,10 @@ mod tests { INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(), "spoofed-regs".to_string(), ), + ( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "follower".to_string(), + ), ( FLOW_SCHEDULED_TIME_MILLIS.to_string(), "1700000000000".to_string(), @@ -92,6 +97,8 @@ mod tests { .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY) .is_none() ); + assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader); + assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none()); assert_eq!( query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS), Some("1700000000000") diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 5485198c266..93c46ff7ff7 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -27,6 +27,9 @@ snafu.workspace = true sql.workspace = true uuid.workspace = true +[dev-dependencies] +prost.workspace = true + [package.metadata.cargo-shear] # `common-error` is required by `common-macro`'s `stack_trace_debug`: the # generated code references `::common_error` paths, so it must stay a direct diff --git a/src/session/src/context.rs b/src/session/src/context.rs index c1448053f35..2f81dfabffb 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -15,6 +15,7 @@ use std::collections::HashMap; use std::fmt::{Display, Formatter}; use std::net::SocketAddr; +use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Duration; @@ -34,12 +35,14 @@ use derive_builder::Builder; use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; pub use crate::hints::{ - LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, - SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, + REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, }; use crate::protocol_ctx::ProtocolCtx; use crate::query_id::QueryId; -use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle}; +use crate::session_config::{ + InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle, +}; use crate::{MutableInner, ReadPreference}; pub type QueryContextRef = Arc; @@ -180,25 +183,41 @@ impl QueryContextBuilder { } } -impl From<&RegionRequestHeader> for QueryContext { - fn from(value: &RegionRequestHeader) -> Self { +impl TryFrom<&RegionRequestHeader> for QueryContext { + type Error = crate::session_config::Error; + + fn try_from(value: &RegionRequestHeader) -> Result { if let Some(ctx) = &value.query_context { - ctx.clone().into() + ctx.clone().try_into() } else { - QueryContextBuilder::default() + Ok(QueryContextBuilder::default() .set_extension( REMOTE_QUERY_ID_EXTENSION_KEY.to_string(), generate_remote_query_id(), ) - .build() + .build()) } } } -impl From for QueryContext { - fn from(ctx: api::v1::QueryContext) -> Self { +impl TryFrom for QueryContext { + type Error = crate::session_config::Error; + + fn try_from(mut ctx: api::v1::QueryContext) -> Result { + let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY); + let read_preference = match read_preference { + Some(value) => ReadPreference::from_str(&value).map_err(|_| { + InvalidConfigValueSnafu { + name: "read_preference", + value, + hint: "Expected a supported read preference", + } + .build() + })?, + None => ReadPreference::default(), + }; let sequences = ctx.snapshot_seqs.as_ref(); - QueryContextBuilder::default() + Ok(QueryContextBuilder::default() .current_catalog(ctx.current_catalog) .current_schema(ctx.current_schema) .timezone(parse_timezone(Some(&ctx.timezone))) @@ -215,7 +234,8 @@ impl From for QueryContext { .unwrap_or_default(), ))) .explain_options(ctx.explain) - .build() + .read_preference(read_preference) + .build()) } } @@ -234,6 +254,17 @@ impl From for api::v1::QueryContext { ) -> Self { let explain = mutable_query_context_data.read().unwrap().explain_options; let mutable_inner = mutable_inner.read().unwrap(); + let mut extensions = extensions; + extensions.remove(READ_PREFERENCE_EXTENSION_KEY); + if mutable_inner.read_preference != ReadPreference::Leader { + extensions.insert( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + mutable_inner + .read_preference + .to_string() + .to_ascii_lowercase(), + ); + } api::v1::QueryContext { current_catalog, current_schema: mutable_inner.schema.clone(), @@ -843,7 +874,7 @@ mod test { 0 ); let wire: api::v1::QueryContext = admitted.into(); - let restored = QueryContext::from(wire); + let restored = QueryContext::try_from(wire).unwrap(); assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100); } @@ -856,7 +887,7 @@ mod test { .extensions .contains_key(crate::hints::INSERT_SKIP_WAL_HINT) ); - let restored: QueryContext = api_context.into(); + let restored: QueryContext = api_context.try_into().unwrap(); assert!(!restored.skip_wal()); } @@ -875,7 +906,7 @@ mod test { explain: None, }; - let session_ctx: QueryContext = api_ctx.clone().into(); + let session_ctx: QueryContext = api_ctx.clone().try_into().unwrap(); let roundtrip_api: api::v1::QueryContext = session_ctx.into(); assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog); @@ -910,7 +941,7 @@ mod test { assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id); let proto: api::v1::QueryContext = (&ctx).into(); - let restored = QueryContext::from(proto); + let restored = QueryContext::try_from(proto).unwrap(); assert_eq!(restored.remote_query_id(), Some(query_id)); assert_eq!( restored.remote_query_id_value().unwrap().to_string(), @@ -918,6 +949,117 @@ mod test { ); } + #[test] + fn test_query_context_read_preference_wire_values() { + let wire = |value: Option<&str>| api::v1::QueryContext { + current_catalog: DEFAULT_CATALOG_NAME.to_string(), + current_schema: DEFAULT_SCHEMA_NAME.to_string(), + timezone: "UTC".to_string(), + extensions: value + .map(|value| { + HashMap::from([(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string())]) + }) + .unwrap_or_default(), + channel: Channel::Internal as u32, + snapshot_seqs: None, + explain: None, + }; + + let header = RegionRequestHeader::default(); + let header_context = QueryContext::try_from(&header).unwrap(); + assert_eq!(header_context.read_preference(), ReadPreference::Leader); + + for value in [None, Some("leader"), Some("LEADER")] { + let ctx = QueryContext::try_from(wire(value)).unwrap(); + assert_eq!(ctx.read_preference(), ReadPreference::Leader); + assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY)); + } + for value in ["", "unknown", "follower", "follower_preferred", "random"] { + let error = QueryContext::try_from(wire(Some(value))).unwrap_err(); + assert!(matches!( + error, + crate::session_config::Error::InvalidConfigValue { + name, + value: actual, + .. + } if name == "read_preference" && actual == value + )); + } + } + + #[test] + fn test_legacy_query_context_wire_defaults_to_leader() { + use prost::Message; + + let legacy = api::v1::QueryContext::decode( + include_bytes!( + "../../../tests/compatibility/wire/read_preference/legacy_query_context.pb" + ) + .as_slice(), + ) + .unwrap(); + assert_eq!(legacy.current_catalog, "c1"); + assert_eq!(legacy.current_schema, "s1"); + assert_eq!(legacy.timezone, "UTC"); + assert_eq!(legacy.channel, Channel::Internal as u32); + assert_eq!( + legacy + .extensions + .get("flow.return_region_seq") + .map(String::as_str), + Some("true") + ); + assert!( + !legacy + .extensions + .contains_key(READ_PREFERENCE_EXTENSION_KEY) + ); + + let restored = QueryContext::try_from(legacy).unwrap(); + assert_eq!(restored.read_preference(), ReadPreference::Leader); + assert_eq!(restored.current_catalog(), "c1"); + assert_eq!(restored.current_schema(), "s1"); + assert_eq!(restored.timezone().to_string(), "UTC"); + assert_eq!(restored.channel(), Channel::Internal); + assert_eq!(restored.extension("flow.return_region_seq"), Some("true")); + + let outbound: api::v1::QueryContext = restored.into(); + let roundtrip = api::v1::QueryContext::decode(outbound.encode_to_vec().as_slice()).unwrap(); + assert_eq!(roundtrip.current_catalog, "c1"); + assert_eq!(roundtrip.current_schema, "s1"); + assert_eq!(roundtrip.timezone, "UTC"); + assert_eq!(roundtrip.channel, Channel::Internal as u32); + assert_eq!( + roundtrip + .extensions + .get("flow.return_region_seq") + .map(String::as_str), + Some("true") + ); + assert!( + !roundtrip + .extensions + .contains_key(READ_PREFERENCE_EXTENSION_KEY) + ); + } + + #[test] + fn test_outbound_read_preference_clears_stale_extension() { + let context = QueryContextBuilder::default() + .set_extension( + READ_PREFERENCE_EXTENSION_KEY.to_string(), + "stale".to_string(), + ) + .set_extension("unrelated".to_string(), "preserved".to_string()) + .build(); + let wire: api::v1::QueryContext = context.into(); + assert!(!wire.extensions.contains_key(READ_PREFERENCE_EXTENSION_KEY)); + assert_eq!( + wire.extensions.get("unrelated").map(String::as_str), + Some("preserved") + ); + } + #[test] fn test_live_analyze_metrics_requires_matching_remote_query_id() { let mut ctx = QueryContext::arc().as_ref().clone(); @@ -944,7 +1086,7 @@ mod test { assert!(ctx.clone().batching_enabled()); assert!(ctx.fork().batching_enabled()); let wire: api::v1::QueryContext = ctx.into(); - let restored = QueryContext::from(wire); + let restored = QueryContext::try_from(wire).unwrap(); assert!(!restored.batching_enabled()); assert!(!restored.logical_batching_enabled()); let ctx = QueryContextBuilder::default() diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 971426f4cc0..3ed4fd8dc2e 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -27,11 +27,13 @@ pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal"; pub const READ_PREFERENCE_HINT: &str = "read_preference"; -pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [ +pub const READ_PREFERENCE_EXTENSION_KEY: &str = "query.read_preference"; +pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [ REMOTE_QUERY_ID_EXTENSION_KEY, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, LIVE_ANALYZE_METRICS_EXTENSION_KEY, + READ_PREFERENCE_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. @@ -65,6 +67,7 @@ mod tests { assert!(is_reserved_extension_key( LIVE_ANALYZE_METRICS_EXTENSION_KEY )); + assert!(is_reserved_extension_key(READ_PREFERENCE_EXTENSION_KEY)); assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT)); } } diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 76f198c9707..bd357033a7e 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -17,13 +17,15 @@ use std::env; use std::net::TcpListener; use std::ops::RangeInclusive; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; use std::time::Duration; use api::v1::health_check_server::HealthCheckServer; use api::v1::region::region_server::RegionServer; use arrow_flight::flight_service_server::FlightServiceServer; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_extension::DistributedInformationExtension; @@ -74,9 +76,11 @@ use servers::grpc::{GrpcOptions, HealthCheckHandler}; use servers::server::ServerHandlers; use store_api::storage::RegionId; use tempfile::TempDir; +use tokio::io::{AsyncRead, AsyncWrite}; use tonic::codec::CompressionEncoding; use tonic::transport::Server; -use tower::service_fn; +use tonic::transport::server::Connected; +use tower::{Layer, Service, service_fn}; use uuid::Uuid; use crate::test_util::{ @@ -84,11 +88,31 @@ use crate::test_util::{ create_tmp_dir_and_datanode_opts, }; +/// Flight `DoGet` requests another node sent to a datanode over its network address (see +/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). +/// +/// The frontend and datanodes inside the test process use in-process clients, so counted requests +/// come from remote nodes querying regions owned by this datanode. +#[derive(Debug, Default)] +pub struct DatanodeRpcStats { + do_get_requests: AtomicUsize, +} + +impl DatanodeRpcStats { + /// The number of Flight `DoGet` requests another node sent to the datanode. + pub fn requests(&self) -> usize { + self.do_get_requests.load(Ordering::Relaxed) + } +} + pub struct GreptimeDbCluster { pub guards: Vec, pub datanode_options: Vec, pub datanode_instances: HashMap, + /// The stats of the network address of every datanode, empty unless the cluster serves the + /// datanodes at their own addresses. + pub datanode_rpc_stats: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -99,6 +123,13 @@ impl GreptimeDbCluster { &self.frontend.instance } + /// The stats of the network address of the datanode `datanode_id`, if the cluster serves the + /// datanodes at their own addresses (see + /// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). + pub fn datanode_rpc_stats(&self, datanode_id: DatanodeId) -> Option<&Arc> { + self.datanode_rpc_stats.get(&datanode_id) + } + /// List all SST files from all datanodes. pub async fn list_sst_files_from_all_datanodes(&self) -> BTreeSet { let mut sst_files = BTreeSet::new(); @@ -177,6 +208,7 @@ pub struct GreptimeDbClusterBuilder { plugins: Plugins, local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, + real_datanode_grpc_addr: bool, } impl GreptimeDbClusterBuilder { @@ -214,6 +246,7 @@ impl GreptimeDbClusterBuilder { plugins: Plugins::default(), local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), + real_datanode_grpc_addr: false, } } @@ -299,6 +332,21 @@ impl GreptimeDbClusterBuilder { self } + /// Serves every datanode on its own real localhost TCP port and registers that + /// address in metasrv, instead of the default in-process duplex placeholder address + /// that all datanodes share. + /// + /// Datanodes serve the region queries of the plans they execute by querying the + /// leader of every region of the plan (see `DatanodeRegionQueryHandler`), so a test + /// that exercises a datanode-to-datanode region query needs each datanode to be + /// reachable at the address recorded in the table route, i.e. the address the + /// datanode registers with its heartbeat. + #[must_use] + pub fn with_real_datanode_grpc_addr(mut self, real_datanode_grpc_addr: bool) -> Self { + self.real_datanode_grpc_addr = real_datanode_grpc_addr; + self + } + pub async fn build_with( &self, datanode_options: Vec, @@ -346,7 +394,22 @@ impl GreptimeDbClusterBuilder { .build_datanodes_with_options(&metasrv, &datanode_options) .await; - build_datanode_clients(datanode_clients.clone(), &datanode_instances).await; + let datanode_rpc_stats = if self.real_datanode_grpc_addr { + datanode_options + .iter() + .map(|opts| (opts.node_id.unwrap(), Arc::new(DatanodeRpcStats::default()))) + .collect::>() + } else { + HashMap::new() + }; + + build_datanode_clients( + datanode_clients.clone(), + &datanode_instances, + &datanode_options, + &datanode_rpc_stats, + ) + .await; self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes) .await; @@ -365,6 +428,7 @@ impl GreptimeDbClusterBuilder { datanode_options, guards, datanode_instances, + datanode_rpc_stats, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -421,12 +485,28 @@ impl GreptimeDbClusterBuilder { opts }; opts.node_id = Some(datanode_id); - + if self.real_datanode_grpc_addr { + // Give every datanode its own address, so that the other datanodes can + // reach it (see `with_real_datanode_grpc_addr`). + let addr = self.choose_datanode_grpc_addr(); + opts.grpc.bind_addr = addr.clone(); + opts.grpc.server_addr = addr; + } options.push(opts); } (options, guards) } + /// Chooses a free localhost port for a datanode. + /// + /// Uses a port range disjoint from the one of the frontend servers + /// (see [`Self::build_frontend_options`]). + fn choose_datanode_grpc_addr(&self) -> String { + let localhost = "127.0.0.1"; + let port = self.choose_random_unused_port(24000..=34000, 20, localhost); + format!("{}:{}", localhost, port) + } + async fn build_datanodes_with_options( &self, metasrv: &MockInfo, @@ -436,7 +516,8 @@ impl GreptimeDbClusterBuilder { for opts in options { let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; - instances.insert(opts.node_id.unwrap(), datanode); + let datanode_id = opts.node_id.unwrap(); + instances.insert(datanode_id, datanode); } instances @@ -474,15 +555,12 @@ impl GreptimeDbClusterBuilder { let meta_backend = new_read_only_meta_kv_backend(meta_client.clone()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(meta_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(meta_backend.clone())); let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend); builder - .with_cache_registry(layered_cache_registry) + .with_cache_registry(layered_cache_registry.clone()) .with_meta_client(meta_client); let mut datanode = builder.build().await.unwrap(); @@ -630,18 +708,34 @@ impl GreptimeDbClusterBuilder { async fn build_datanode_clients( clients: Arc, instances: &HashMap, + options: &[DatanodeOptions], + rpc_stats: &HashMap>, ) { - for (&datanode_id, instance) in instances { - let (addr, client) = create_datanode_client(instance).await; + for opts in options { + let datanode_id = opts.node_id.unwrap(); + let instance = &instances[&datanode_id]; + let (addr, client) = create_datanode_client( + instance, + &opts.grpc.server_addr, + rpc_stats.get(&datanode_id).cloned(), + ) + .await; clients .insert_client(Peer::new(datanode_id, addr), client) .await; } } -async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { - let (client, server) = tokio::io::duplex(1024); - +/// Runs a gRPC server that serves the region server of `datanode` over `incoming` until the +/// cluster is dropped. +fn serve_datanode_region_server( + datanode: &Datanode, + incoming: I, + rpc_stats: Option>, +) where + I: futures::Stream> + Send + 'static, + IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static, +{ let runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("grpc-handlers") @@ -652,8 +746,9 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { let region_server_handler = RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime); - let _handle = tokio::spawn(async move { + tokio::spawn(async move { Server::builder() + .layer(CountRequestsLayer { rpc_stats }) .add_service( FlightServiceServer::new(flight_handler) .accept_compressed(CompressionEncoding::Gzip) @@ -669,12 +764,40 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { .send_compressed(CompressionEncoding::Zstd), ) .add_service(HealthCheckServer::new(HealthCheckHandler)) - .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) + .serve_with_incoming(incoming) .await }); +} + +/// Serves the region server of `datanode` and returns the address and the client of the address. +/// +/// The returned client is always an in-process client: the frontend and the datanodes of the test +/// process reach the region server through a duplex stream, no matter which address they use. This +/// keeps the in-process traffic out of `rpc_stats`. +/// +/// [`test_util::PEER_PLACEHOLDER_ADDR`] is not a real address, so it cannot be served over the +/// network at all. Any other address is a real one: the region server is additionally served on +/// its local TCP port (`rpc_stats` counts the requests of the remote nodes that reach it, e.g. the +/// region query that a datanode sends for the regions of another datanode). +async fn create_datanode_client( + datanode: &Datanode, + addr: &str, + rpc_stats: Option>, +) -> (String, Client) { + let (client, server) = tokio::io::duplex(1024); + serve_datanode_region_server( + datanode, + futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]), + None, + ); + + if addr != test_util::PEER_PLACEHOLDER_ADDR { + let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + serve_datanode_region_server(datanode, incoming, rpc_stats); + } let mut client = Some(client); - let addr = test_util::PEER_PLACEHOLDER_ADDR; let channel_manager = ChannelManager::new(); channel_manager .reset_with_connector( @@ -699,6 +822,52 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ) } +/// Counts Flight `DoGet` requests received by the wrapped server. +#[derive(Clone)] +struct CountRequests { + inner: S, + rpc_stats: Option>, +} + +impl Service> for CountRequests +where + S: Service>, +{ + type Response = S::Response; + type Error = S::Error; + type Future = S::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + if request.uri().path() == "/arrow.flight.protocol.FlightService/DoGet" + && let Some(rpc_stats) = &self.rpc_stats + { + rpc_stats.do_get_requests.fetch_add(1, Ordering::Relaxed); + } + + self.inner.call(request) + } +} + +#[derive(Clone)] +struct CountRequestsLayer { + rpc_stats: Option>, +} + +impl Layer for CountRequestsLayer { + type Service = CountRequests; + + fn layer(&self, inner: S) -> Self::Service { + CountRequests { + inner, + rpc_stats: self.rpc_stats.clone(), + } + } +} + // Mock connectors are registered on a shared channel manager. #[cfg(test)] #[allow(deprecated)] @@ -993,13 +1162,41 @@ mod tests { // Assert: the inactive peer was rejected and the production Flow service handled the RPC. assert_eq!( - (vec![active_peer.addr], vec![inactive_peer.addr]), + (vec![active_peer.addr.clone()], vec![inactive_peer.addr]), client.peer_addresses_by_state() ); assert_eq!( StatusCode::InvalidArguments, response.unwrap_err().status_code() ); + + for value in ["", "unknown", "follower"] { + let mut query_context = + api::v1::QueryContext::from(session::context::QueryContext::arc().as_ref()); + query_context.extensions.insert( + session::hints::READ_PREFERENCE_EXTENSION_KEY.to_string(), + value.to_string(), + ); + let error = clients + .flownode(&active_peer) + .await + .handle(FlowRequest { + header: Some(api::v1::flow::FlowRequestHeader { + tracing_context: Default::default(), + query_context: Some(query_context), + }), + body: None, + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + let message = error.output_msg(); + assert!(message.contains("read_preference"), "{message}"); + if !value.is_empty() { + assert!(message.contains(value), "{message}"); + } + } + flownode.shutdown().await.unwrap(); } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 98eb336d9cc..241d4daf5b5 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use std::time::Duration; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_schema::NoopInformationExtension; @@ -209,11 +209,8 @@ impl GreptimeDbStandaloneBuilder { ) -> GreptimeDbStandalone { let plugins = self.plugin.clone().unwrap_or_default(); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new(opts.datanode_options(), plugins.clone(), kv_backend.clone()); diff --git a/tests-integration/src/tests.rs b/tests-integration/src/tests.rs index d70af021651..4a5b2f0c6c5 100644 --- a/tests-integration/src/tests.rs +++ b/tests-integration/src/tests.rs @@ -16,6 +16,7 @@ mod gc; mod instance_kafka_wal_test; mod instance_noop_wal_test; mod instance_test; +mod nested_merge_scan_capability_test; mod promql_test; mod reconcile_table; mod remote_dyn_filter_test; diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs new file mode 100644 index 00000000000..2c86589112d --- /dev/null +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -0,0 +1,749 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Exercises hand-built nested `MergeScan` plans at datanode execution, not optimizer behavior. +//! The outer `handle_remote_read` request is in-process; inner region reads use real TCP Flight. +//! The probe scan resolves against the request's region catalog, while the nested payload resolves +//! against the engine catalog; distinct column names ensure these scopes cannot be confused. +//! Success compares the complete duplicate-sensitive result multiset with frontend SQL and known +//! rows. Missing inner tables and closed remote inner regions must each fail the whole query. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use api::v1::region::{QueryRequest, RegionRequestHeader}; +use client::Client; +use client::region::RegionRequester; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_grpc::channel_manager::ChannelManager; +use common_query::Output; +use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; +use common_telemetry::info; +use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datanode::region_server::RegionServer; +use frontend::instance::Instance; +use query::datafusion::QUERY_PARALLELISM_HINT; +use query::dist_plan::MergeScanLogicalPlan; +use query::parser::QueryLanguageParser; +use query::query_engine::DefaultSerializer; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::{QueryContext, QueryContextRef}; +use store_api::region_request::{RegionCloseRequest, RegionRequest}; +use store_api::storage::RegionId; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; + +use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; +use crate::test_util::execute_sql; + +/// Partitioned probe table. +const PROBE_TABLE: &str = "nested_cap_probe"; +/// Partitioned build table with duplicate join keys. +const BUILD_TABLE: &str = "nested_cap_build"; +/// Missing inner table used to verify plan-resolution failure. +const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; + +/// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key. +const PROBE_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 100), + (2, 20, 200), + // The second partition of the probe table. + (120, 30, 300), + (121, 10, 400), +]; + +/// `(b_id, build_key, build_v)` fixture rows; key `10` duplicates and `40` is unmatched. +const BUILD_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 1), + (2, 10, 2), + (3, 20, 3), + // The second partition of the build table. + (120, 30, 4), + (121, 40, 5), +]; + +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_join_across_datanodes() { + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_capability_join_across_datanodes").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + assert_eq!( + 2, + datanodes(&probe_leaders).len(), + "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ + leaders: {probe_leaders:?}" + ); + assert_eq!( + 2, + datanodes(&build_leaders).len(), + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ + leaders: {build_leaders:?}" + ); + + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + assert_eq!( + expected_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + let query_ctx = query_ctx(); + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + + let remote_datanodes = datanodes(&build_leaders); + let baseline_requests = remote_datanodes + .iter() + .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) + .collect::>(); + + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; + + let actual_rows = multiset(rows); + assert_eq!( + expected_join_rows(), + actual_rows, + "the nested plans of the datanodes returned an unexpected multiset of rows" + ); + assert_eq!( + multiset(table_cells(&reference)), + actual_rows, + "frontend result differs:\n{reference}" + ); + assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); +} + +/// Internal read-preference wire metadata is rejected before nested execution on both inbound paths. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { + use common_meta::node_manager::Datanode; + use common_query::request::QueryRequest as LogicalQueryRequest; + use futures::TryStreamExt; + use session::hints::READ_PREFERENCE_EXTENSION_KEY; + + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_internal_read_preference").await; + let frontend = cluster.fe_instance().clone(); + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let query_ctx = query_ctx(); + let plan = nested_plan(&frontend, &query_ctx).await; + let (region_id, (datanode_id, address)) = probe_leaders.iter().next().unwrap(); + let target = region_server(&cluster, *datanode_id); + // Confirm the same plan completes on the valid default-Leader path before asserting rejection. + let valid_rows = run_nested_queries(&cluster, &probe_leaders, &encode(&plan), &query_ctx).await; + assert_eq!(expected_join_rows(), multiset(valid_rows)); + let mut baseline = cluster + .datanode_instances + .keys() + .map(|id| (*id, rpc_requests(&cluster, *id))) + .collect::>(); + + for value in ["", "invalid", "follower"] { + let mut header = region_request_header(&query_ctx); + header + .query_context + .as_mut() + .unwrap() + .extensions + .insert(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string()); + let header = Some(header); + let error = target + .handle_read(LogicalQueryRequest { + header: header.clone(), + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }) + .await + .err() + .expect("internal read preference must be rejected at logical ingress"); + assert_eq!( + StatusCode::InvalidArguments, + error.status_code(), + "value={value:?}" + ); + let logical_message = error.output_msg().to_ascii_lowercase(); + assert!( + logical_message.contains("read_preference"), + "expected an invalid query-context read-preference error, got: {logical_message}" + ); + if !value.is_empty() { + assert!( + logical_message.contains(value), + "unexpected logical error: {logical_message}" + ); + } + for id in cluster.datanode_instances.keys() { + assert_eq!( + baseline[id], + rpc_requests(&cluster, *id), + "logical invalid context caused DoGet at datanode {id}" + ); + } + + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [address.as_str()], + ); + let requester = RegionRequester::new(client, false, false); + let request = LogicalQueryRequest { + header, + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }; + let flight_message = match Datanode::handle_query(&requester, request).await { + Err(error) => { + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.output_msg().to_ascii_lowercase() + } + Ok(stream) => { + let error = stream + .try_collect::>() + .await + .expect_err("Flight must reject invalid internal read preference"); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.output_msg().to_ascii_lowercase() + } + }; + assert!( + flight_message.contains("read_preference"), + "expected an invalid-parameter read-preference error, got: {flight_message}" + ); + if !value.is_empty() { + assert!( + flight_message.contains(value), + "unexpected Flight error: {flight_message}" + ); + } + for id in cluster.datanode_instances.keys() { + let expected = baseline[id] + u8::from(*id == *datanode_id) as usize; + assert_eq!( + expected, + rpc_requests(&cluster, *id), + "invalid Flight context caused nested DoGet at datanode {id}" + ); + } + for id in cluster.datanode_instances.keys() { + baseline.insert(*id, rpc_requests(&cluster, *id)); + } + } +} + +/// A failing inner region makes the whole query fail: the datanode must not return the rows of the +/// inner regions that succeed. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_inner_failure_fails_query() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_merge_scan_capability_inner_failure_fails_query").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + + let (probe_region, (probe_datanode, _)) = probe_leaders + .iter() + .find(|(_, (datanode, _))| *datanode == 1) + .unwrap_or_else(|| { + panic!( + "expected a region of {PROBE_TABLE} on the datanode 1, actual region leaders: \ + {probe_leaders:?}" + ) + }); + let (build_region, (build_datanode, _)) = build_leaders + .iter() + .find(|(_, (datanode, _))| *datanode != *probe_datanode) + .unwrap_or_else(|| { + panic!( + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual \ + region leaders: {build_leaders:?}" + ) + }); + + // The nested payload resolves against the engine catalog, not the request's region catalog. + let query_ctx = query_ctx(); + let missing_inner_plan = + nested_plan_with_inner_table(&frontend, &query_ctx, MISSING_INNER_TABLE).await; + let plan = encode(&missing_inner_plan); + let error = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .err() + .expect("a nested plan whose inner relation does not exist must not be executed"); + // The top level error of the datanode only names the failing step (`Failed to decode logical + // plan`): the reason is the source of it, which names the table that could not be resolved. + let chain = error_chain(&error); + assert!( + chain.contains(MISSING_INNER_TABLE), + "expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \ + actual error: {chain}" + ); + + // A closed remote inner region must fail the whole query, not return a partial result. + region_server(&cluster, *build_datanode) + .handle_request( + RegionId::from_u64(*build_region), + RegionRequest::Close(RegionCloseRequest { + flush_on_close: false, + }), + ) + .await + .expect("the datanode must close the region of the build table"); + + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + let baseline = rpc_requests(&cluster, *build_datanode); + let stream = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .expect("the request itself must be accepted by the datanode"); + let error = collect_error(stream) + .await + .expect("the query must fail when one of its inner regions fails"); + info!("the nested plan failed as a whole: {error}"); + + // The increased DoGet count proves execution reached the remote region owner. + let actual = rpc_requests(&cluster, *build_datanode); + assert!( + actual > baseline, + "expected the failed nested plan to reach datanode {build_datanode} for region {build_region} over Flight DoGet (count {actual}, baseline {baseline})" + ); +} + +async fn build_cluster(test_name: &str) -> GreptimeDbCluster { + GreptimeDbClusterBuilder::new(test_name) + .await + // A datanode serves the inner `MergeScan` of a nested plan by querying the datanodes that + // own the regions of the inner table, so every datanode must be reachable at the address it + // registers in metasrv, i.e. the address of the table routes. The cluster then also counts + // the requests of the remote nodes, see [`assert_cross_datanode_region_query`]. + .with_real_datanode_grpc_addr(true) + .with_datanodes(2) + .build(false) + .await +} + +fn assert_cross_datanode_region_query( + cluster: &GreptimeDbCluster, + remote_datanodes: &BTreeSet, + baseline_requests: &BTreeMap, +) { + assert_eq!( + 2, + remote_datanodes.len(), + "build table must span both datanodes" + ); + for datanode in remote_datanodes { + let actual = rpc_requests(cluster, *datanode); + let baseline = baseline_requests[datanode]; + assert!( + actual > baseline, + "expected a nested Flight DoGet at datanode {datanode}, count {actual}, baseline {baseline}" + ); + } +} + +fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize { + cluster + .datanode_rpc_stats(datanode_id) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode_id}")) + .requests() +} + +fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer { + cluster + .datanode_instances + .get(&datanode_id) + .unwrap_or_else(|| panic!("expected a datanode {datanode_id}")) + .region_server() +} + +async fn run_nested_queries( + cluster: &GreptimeDbCluster, + probe_leaders: &BTreeMap, + plan: &[u8], + query_ctx: &QueryContextRef, +) -> Vec> { + let mut rows = Vec::new(); + for (probe_region, (probe_datanode, _)) in probe_leaders { + let stream = region_server(cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(query_ctx)), + region_id: *probe_region, + plan: plan.to_vec(), + }, + query_ctx.clone(), + ) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region}: {e}" + ) + }); + let batches = RecordBatches::try_collect(stream) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region} completely: {e}" + ) + }); + let actual = batches + .pretty_print() + .expect("the result of the datanode must be printable"); + rows.extend(table_cells(&actual)); + } + + rows +} + +/// Creates the probe and build tables of the tests and inserts the rows. +async fn prepare_tables(frontend: &Arc) { + execute_sql( + frontend, + &format!( + "CREATE TABLE {PROBE_TABLE} ( + a_id INT, + probe_key INT, + probe_v INT, + a_ts TIMESTAMP, + TIME INDEX (a_ts), + PRIMARY KEY(a_id) + ) + PARTITION ON COLUMNS (a_id) ( + a_id < 100, + a_id >= 100 + ) + engine=mito" + ), + ) + .await; + execute_sql( + frontend, + &format!( + "CREATE TABLE {BUILD_TABLE} ( + b_id INT, + build_key INT, + build_v INT, + b_ts TIMESTAMP, + TIME INDEX (b_ts), + PRIMARY KEY(b_id) + ) + PARTITION ON COLUMNS (b_id) ( + b_id < 100, + b_id >= 100 + ) + engine=mito" + ), + ) + .await; + + execute_sql( + frontend, + &format!( + "INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES {}", + PROBE_ROWS + .iter() + .enumerate() + .map(|(i, (a_id, probe_key, probe_v))| { + format!("({a_id}, {probe_key}, {probe_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; + execute_sql( + frontend, + &format!( + "INSERT INTO {BUILD_TABLE}(b_id, build_key, build_v, b_ts) VALUES {}", + BUILD_ROWS + .iter() + .enumerate() + .map(|(i, (b_id, build_key, build_v))| { + format!("({b_id}, {build_key}, {build_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; +} + +fn encode(plan: &LogicalPlan) -> Vec { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .expect("the plan must be encodable") + .to_vec() +} + +/// Plans the hand-built nested join sent in the region request. +async fn nested_plan(frontend: &Arc, query_ctx: &QueryContextRef) -> LogicalPlan { + nested_plan_with_inner_table(frontend, query_ctx, BUILD_TABLE).await +} + +async fn nested_plan_with_inner_table( + frontend: &Arc, + query_ctx: &QueryContextRef, + inner_table: &str, +) -> LogicalPlan { + let probe = plan_sql( + frontend, + &format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"), + query_ctx, + ) + .await; + let build = plan_sql( + frontend, + &format!("SELECT build_key, build_v FROM {BUILD_TABLE}"), + query_ctx, + ) + .await; + let build = rename_scans(build, inner_table); + + nested_plan_of(&probe, build) +} + +fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { + info!("probe side of the nested plan:\n{probe}"); + info!("build side of the nested plan:\n{build}"); + + let inner = MergeScanLogicalPlan::new(build, false, Default::default()).into_logical_plan(); + let plan = LogicalPlanBuilder::from(probe.clone()) + .join( + inner, + JoinType::Inner, + (vec!["probe_key"], vec!["build_key"]), + None, + ) + .expect("the join keys must be resolvable in the sides of the join") + .build() + .expect("the nested plan must be a valid logical plan"); + info!("nested plan sent to the datanode:\n{plan}"); + + plan +} + +fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan { + use datafusion::common::TableReference; + use datafusion::common::tree_node::{Transformed, TreeNode}; + + plan.transform_up(|plan| match plan { + LogicalPlan::TableScan(mut scan) => { + info!( + "renaming the scan of the nested plan from {} to {table_name}", + scan.table_name + ); + scan.table_name = TableReference::bare(table_name); + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + } + other => Ok(Transformed::no(other)), + }) + .expect("the plan must be rewritable") + .data +} + +async fn plan_sql(frontend: &Arc, sql: &str, query_ctx: &QueryContextRef) -> LogicalPlan { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx) + .unwrap_or_else(|e| panic!("failed to parse `{sql}`: {e}")); + frontend + .statement_executor() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap_or_else(|e| panic!("failed to plan `{sql}`: {e}")) +} + +/// Frontend reference join over the nested plan's output columns. +fn join_sql() -> String { + format!( + "SELECT p.a_id, p.probe_key, p.probe_v, b.build_key, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_key" + ) +} + +/// Expected complete join result, retaining duplicate rows. +fn expected_join_rows() -> Vec> { + let mut rows = Vec::new(); + for (a_id, probe_key, probe_v) in PROBE_ROWS { + for (_, build_key, build_v) in BUILD_ROWS { + if probe_key == build_key { + rows.push(vec![ + a_id.to_string(), + probe_key.to_string(), + probe_v.to_string(), + build_key.to_string(), + build_v.to_string(), + ]); + } + } + } + multiset(rows) +} + +async fn region_leaders(frontend: &Arc, table: &str) -> BTreeMap { + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::UInt64Type; + + let sql = format!( + "SELECT region_id, peer_id, peer_addr FROM information_schema.region_peers + WHERE table_schema = 'public' AND table_name = '{table}' AND is_leader = 'Yes' + ORDER BY region_id" + ); + let batches = output_batches(run_sql(frontend, &sql, query_ctx()).await).await; + + let mut leaders = BTreeMap::new(); + for batch in batches.iter() { + let batch = batch.df_record_batch(); + let region_ids = batch + .column_by_name("region_id") + .expect("region_id column") + .as_primitive::(); + let peer_ids = batch + .column_by_name("peer_id") + .expect("peer_id column") + .as_primitive::(); + let peer_addrs = batch + .column_by_name("peer_addr") + .expect("peer_addr column") + .as_string::(); + for row in 0..batch.num_rows() { + leaders.insert( + region_ids.value(row), + (peer_ids.value(row), peer_addrs.value(row).to_string()), + ); + } + } + + leaders +} + +fn datanodes(leaders: &BTreeMap) -> BTreeSet { + leaders.values().map(|(id, _)| *id).collect() +} + +async fn output_batches(output: Output) -> RecordBatches { + match output.data { + common_query::OutputData::Stream(stream) => RecordBatches::try_collect(stream) + .await + .expect("the output stream must be collectable"), + common_query::OutputData::RecordBatches(batches) => batches, + common_query::OutputData::AffectedRows(rows) => { + panic!("expected a query output, got {rows} affected rows") + } + } +} + +/// Consumes the full stream to detect errors after any partial rows. +async fn collect_error(stream: SendableRecordBatchStream) -> Option { + use futures::TryStreamExt; + + let result = stream.try_collect::>().await; + match result { + Ok(rows) => { + info!( + "the query returned {} batches instead of failing", + rows.len() + ); + None + } + Err(error) => Some(error_chain(&error)), + } +} + +fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { + let mut messages = vec![error.to_string()]; + let mut source = error.source(); + while let Some(inner) = source { + messages.push(inner.to_string()); + source = inner.source(); + } + messages.join(": ") +} + +async fn query_pretty(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> String { + run_sql(frontend, sql, query_ctx) + .await + .data + .pretty_print() + .await +} + +async fn run_sql(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> Output { + SqlQueryHandler::do_query(frontend.as_ref(), sql, query_ctx) + .await + .remove(0) + .unwrap() +} + +/// Uses one target partition to keep the plans stable. +fn query_ctx() -> QueryContextRef { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, "1"); + Arc::new(query_ctx) +} + +fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { + RegionRequestHeader { + query_context: Some(query_ctx.as_ref().into()), + ..Default::default() + } +} + +fn table_cells(pretty: &str) -> Vec> { + let mut lines = pretty.lines().filter(|line| line.starts_with('|')); + // The first `|` line is the header of the table. + lines.next(); + + lines + .map(|line| { + line.trim_matches('|') + .split('|') + .map(|cell| cell.trim().to_string()) + .collect() + }) + .collect() +} + +fn multiset(mut rows: Vec>) -> Vec> { + rows.sort(); + rows +} diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 542f10debfd..404b6332b94 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -2,7 +2,18 @@ Compatibility tests verify that one GreptimeDB version can restart on state written by another version. -Tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. +SQL restart tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. + +## Wire fixtures + +Frozen protobuf fixtures are exercised by the Rust tests that own their context +conversion, not by SQLness case discovery. The +[legacy query-context fixture](wire/read_preference/README.md) checks omitted +read-preference defaulting and unrelated metadata preservation: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` ## Quick Start diff --git a/tests/compatibility/wire/read_preference/README.md b/tests/compatibility/wire/read_preference/README.md new file mode 100644 index 00000000000..9e7b9ba74c1 --- /dev/null +++ b/tests/compatibility/wire/read_preference/README.md @@ -0,0 +1,45 @@ +# Legacy query-context wire compatibility + +`legacy_query_context.pb` is a generated legacy-shaped protobuf fixture, not a +capture from a released binary. It uses the unchanged `greptime.v1.QueryContext` +schema at greptime-proto revision +`9533078d6fd4bdaddbca0750cb9bf7ad64b2e568` and omits the internal +`query.read_preference` extension. + +Logical contents: + +- Catalog `c1`, schema `s1`, timezone `UTC`. +- Unrelated extension `flow.return_region_seq = true`. +- Channel `Internal` (`255`); no snapshot sequences or explain options. + +Run the real protobuf decode and session conversion regression: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` + +The test checks that the omitted preference defaults to Leader, that the original +fields and unrelated extension survive conversion and re-encoding, and that +Leader does not introduce a preference carrier. Existing session tests cover +present invalid or unsupported values. This fixture does not demonstrate that +an old receiver honors a non-Leader preference sent by an updated node. + +## Reproduce the fixture + +With `PROTO_ROOT` pointing to the `proto` directory of the revision above: + +```shell +printf '%s\n' \ + 'current_catalog: "c1"' \ + 'current_schema: "s1"' \ + 'timezone: "UTC"' \ + 'extensions { key: "flow.return_region_seq" value: "true" }' \ + 'channel: 255' | + protoc --proto_path="$PROTO_ROOT" --encode=greptime.v1.QueryContext \ + greptime/v1/common.proto > legacy_query_context.pb +``` + +Generated with `libprotoc 25.1`. The 48-byte fixture has SHA-256 +`4127ebbee8cdb1375ef8ad52754be56ed206c12513e5f61c4354c04587cba638`. +The test never rewrites it and does not require protobuf map ordering to remain +stable. diff --git a/tests/compatibility/wire/read_preference/legacy_query_context.pb b/tests/compatibility/wire/read_preference/legacy_query_context.pb new file mode 100644 index 00000000000..1990ede9411 --- /dev/null +++ b/tests/compatibility/wire/read_preference/legacy_query_context.pb @@ -0,0 +1,3 @@ + +c1s1"UTC* +flow.return_region_seqtrue0ÿ \ No newline at end of file From b913be2af77a5a936bcccb3d5b50358d1313a777 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sat, 10 Oct 2026 07:04:02 +0000 Subject: [PATCH 133/133] fix(promql): return range-vector results for instant queries (#9477) * fix(promql): return range-vector results for instant queries An instant query of a range selector or subquery (e.g. `up[5m]`) returned the internal range arrays, which panicked when the frontend decoded them as dictionaries. Plan such queries as the selected samples with their own timestamps, and reject range-vector expressions in range queries like Prometheus does. Signed-off-by: Dennis Zhuang * test(promql): keep encode_substrait on a valid range query A range query of a range-vector expression now fails like in Prometheus. Wrap the subquery in max_over_time so the case still encodes both RangeManipulate nodes through substrait. Signed-off-by: Dennis Zhuang * fix(promql): plan wrapped range-vector roots and reject a zero subquery step TQL EXPLAIN, ANALYZE and AS alias wrap the root in an extension, which the range-vector path did not unwrap. A zero subquery step divided by zero. Signed-off-by: Dennis Zhuang * test(promql): cover wrapped range-vector roots and a zero subquery step Signed-off-by: Dennis Zhuang * fix(promql): default subquery step and reject range vectors in range queries A root subquery without a step uses the one-minute default instead of the query step. The HTTP range query endpoint rejects range-vector expressions even when start equals end. Signed-off-by: Dennis Zhuang * test: cover root subquery default step Signed-off-by: Dennis Zhuang * refactor(promql): align root subqueries with subquery_child_window Signed-off-by: Dennis Zhuang * fix(promql): return millisecond timestamps for range-vector results Signed-off-by: Dennis Zhuang * refactor(promql): keep the time index timezone and name its table in errors Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/promql/error.rs | 16 ++ src/query/src/promql/planner.rs | 179 +++++++++++++++++- src/query/src/promql/planner/at_modifier.rs | 2 +- src/servers/src/http/prometheus.rs | 45 +++++ tests-integration/tests/http.rs | 19 ++ .../common/promql/encode_substrait.result | 56 +++--- .../common/promql/encode_substrait.sql | 8 +- .../common/promql/range_vector_result.result | 143 ++++++++++++++ .../common/promql/range_vector_result.sql | 49 +++++ 9 files changed, 479 insertions(+), 38 deletions(-) create mode 100644 tests/cases/standalone/common/promql/range_vector_result.result create mode 100644 tests/cases/standalone/common/promql/range_vector_result.sql diff --git a/src/query/src/promql/error.rs b/src/query/src/promql/error.rs index 121b5e81a83..78e7289b479 100644 --- a/src/query/src/promql/error.rs +++ b/src/query/src/promql/error.rs @@ -130,6 +130,20 @@ pub enum Error { location: Location, }, + #[snafu(display("Subquery step must be positive"))] + ZeroSubqueryStep { + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "invalid expression type \"range vector\" for range query, must be Scalar or instant Vector" + ))] + RangeVectorInRangeQuery { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display( "The end time must be greater than start time, start: {:?}, end: {:?}", start, @@ -253,6 +267,8 @@ impl ErrorExt for Error { | MultipleVector { .. } | ExpectRangeSelector { .. } | ZeroRangeSelector { .. } + | RangeVectorInRangeQuery { .. } + | ZeroSubqueryStep { .. } | InvalidTimeRange { .. } | SubqueryTimestampOutOfRange { .. } | ColumnNotFound { .. } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5d8fc50edbd..5867301f972 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -106,14 +106,17 @@ use crate::promql::error::{ CatalogSnafu, ColumnNotFoundSnafu, DataFusionPlanningSnafu, ExpectRangeSelectorSnafu, FunctionInvalidArgumentSnafu, InvalidDestinationLabelNameSnafu, InvalidRegularExpressionSnafu, InvalidTimeRangeSnafu, MultiFieldsNotSupportedSnafu, MultipleMetricMatchersSnafu, - MultipleVectorSnafu, NoMetricMatcherSnafu, Result, SubqueryTimestampOutOfRangeSnafu, - TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, - UnknownTableSnafu, UnsupportedExprSnafu, UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, - ZeroRangeSelectorSnafu, + MultipleVectorSnafu, NoMetricMatcherSnafu, RangeVectorInRangeQuerySnafu, Result, + SubqueryTimestampOutOfRangeSnafu, TableNameNotFoundSnafu, TimeIndexNotFoundSnafu, + UnexpectedPlanExprSnafu, UnexpectedTokenSnafu, UnknownTableSnafu, UnsupportedExprSnafu, + UnsupportedMatcherOpSnafu, ValueNotFoundSnafu, ZeroRangeSelectorSnafu, ZeroSubqueryStepSnafu, }; use crate::promql::label_values::signed_millis_since_epoch; use crate::query_engine::QueryEngineState; +/// Step of a subquery without an explicit one (`x[5m:]`). Prometheus uses the global +/// evaluation interval, whose default is one minute, rather than the query step. +const DEFAULT_SUBQUERY_STEP_MS: Millisecond = 60_000; /// `time()` function in PromQL. const SPECIAL_TIME_FUNCTION: &str = "time"; /// `scalar()` function in PromQL. @@ -315,7 +318,7 @@ impl PromPlanner { }; let plan = planner - .prom_expr_to_plan(&stmt.expr, query_engine_state) + .prom_root_expr_to_plan(&stmt.expr, query_engine_state) .await?; // Never leak internal series identifier to output. @@ -399,6 +402,170 @@ impl PromPlanner { Ok(res) } + /// Plans the root expression of a statement, or the expression wrapped by a root extension + /// such as `EXPLAIN` or an alias. + /// + /// Prometheus only accepts a range-vector expression in an instant query, whose result is + /// the matrix of samples the expression selects. + async fn prom_root_expr_to_plan( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + if matches!(prom_expr, PromExpr::Extension(_)) + || prom_expr.value_type() != ValueType::Matrix + { + return self.prom_expr_to_plan(prom_expr, query_engine_state).await; + } + ensure!(self.ctx.start == self.ctx.end, RangeVectorInRangeQuerySnafu); + self.prom_range_vector_result_to_plan(prom_expr, query_engine_state) + .await + } + + /// Plans the result of a range-vector expression in an instant query: the raw samples of a + /// range selector, or the evaluation points of a subquery, with their own timestamps. + /// + /// Range functions consume the same expressions through [`RangeManipulate`], whose range + /// arrays are an internal representation that must not reach the query output. + async fn prom_range_vector_result_to_plan( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + let mut expr = prom_expr; + while let PromExpr::Paren(ParenExpr { expr: inner }) = expr { + expr = inner; + } + match expr { + PromExpr::MatrixSelector(selector) => { + self.prom_range_selector_samples_to_plan(selector).await + } + PromExpr::Subquery(subquery) => { + self.prom_subquery_samples_to_plan(subquery, query_engine_state) + .await + } + _ => UnsupportedExprSnafu { + name: format!("range vector result of {expr:?}"), + } + .fail(), + } + } + + /// Selects the samples of `selector` in `(t - offset - range, t - offset]`, where `t` is the + /// evaluation time or the `@` anchor. Stale markers are dropped and timestamps are kept. + async fn prom_range_selector_samples_to_plan( + &mut self, + selector: &MatrixSelector, + ) -> Result { + let MatrixSelector { vs, range } = selector; + let VectorSelector { + name, + offset, + matchers, + at, + } = vs; + let matchers = self.preprocess_label_matchers(matchers, name)?; + ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); + self.ctx.range = Some(range.as_millis() as _); + let offset_ms = match self.at_modifier_offset(at, offset)? { + Some(at_offset) => at_offset, + None => Self::offset_millis(offset), + }; + if let Some(empty_plan) = self.setup_context().await? { + return Ok(empty_plan); + } + let plan = self + .selector_to_series_normalize_plan(offset_ms, matchers, true) + .await?; + self.millisecond_time_index_plan(plan) + } + + /// Casts a microsecond or nanosecond time index to milliseconds, the precision of every + /// other PromQL result and the only one the Prometheus HTTP API encodes. + fn millisecond_time_index_plan(&self, plan: LogicalPlan) -> Result { + let time_index = + self.ctx + .time_index_column + .as_deref() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let (_, field) = plan + .schema() + .qualified_field_with_unqualified_name(time_index) + .context(DataFusionPlanningSnafu)?; + let millis = match field.data_type() { + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => return Ok(plan), + ArrowDataType::Timestamp(_, timezone) => { + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, timezone.clone()) + } + other => unreachable!("time index is a timestamp, got {other}"), + }; + let exprs = plan + .schema() + .iter() + .map(|(qualifier, field)| { + let column = DfExpr::Column(Column::new(qualifier.cloned(), field.name())); + if field.name() == time_index { + column + .cast_to(&millis, plan.schema()) + .map(|cast| cast.alias_qualified(qualifier.cloned(), field.name().clone())) + } else { + Ok(column) + } + }) + .collect::>>() + .context(DataFusionPlanningSnafu)?; + LogicalPlanBuilder::from(plan) + .project(exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + + /// Evaluates the inner expression of `subquery` at the multiples of its step in + /// `(t - offset - range, t - offset]`, where `t` is the evaluation time or the `@` anchor. + async fn prom_subquery_samples_to_plan( + &mut self, + subquery: &SubqueryExpr, + query_engine_state: &QueryEngineState, + ) -> Result { + let SubqueryExpr { + expr, + range, + step, + offset, + at, + .. + } = subquery; + ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); + let end = match self.at_ref_time(at, offset)? { + Some(anchor) => anchor, + None => Self::anchor_sub(self.ctx.end, Self::offset_millis(offset))?, + }; + self.ctx.interval = match step { + Some(step) => step.as_millis() as _, + None => DEFAULT_SUBQUERY_STEP_MS, + }; + let interval = self.ctx.interval; + ensure!(interval > 0, ZeroSubqueryStepSnafu); + let window = Self::subquery_child_window(end, end, 0, 0, range.as_millis() as _, interval)?; + let Some((start, end)) = window else { + // No step multiple falls in the window; plan the inner expression only for its + // schema. + self.ctx.start = end; + self.ctx.end = end; + let plan = self.prom_expr_to_plan(expr, query_engine_state).await?; + return Ok(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: plan.schema().clone(), + })); + }; + self.ctx.start = start; + self.ctx.end = end; + self.prom_expr_to_plan(expr, query_engine_state).await + } + async fn prom_subquery_expr_to_plan( &mut self, query_engine_state: &QueryEngineState, @@ -2226,7 +2393,7 @@ impl PromPlanner { let expr = &ext_expr.expr; let children = expr.children(); let plan = self - .prom_expr_to_plan(&children[0], query_engine_state) + .prom_root_expr_to_plan(&children[0], query_engine_state) .await?; // Wrapper for the explanation/analyze of the existing plan // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain diff --git a/src/query/src/promql/planner/at_modifier.rs b/src/query/src/promql/planner/at_modifier.rs index bde7c51bcb1..545f1714a22 100644 --- a/src/query/src/promql/planner/at_modifier.rs +++ b/src/query/src/promql/planner/at_modifier.rs @@ -41,7 +41,7 @@ impl PromPlanner { /// - `offset` shifts the anchor backwards: the window ends at `anchor - offset`. /// /// Returns `None` when the selector has no `@` modifier. - fn at_ref_time( + pub(crate) fn at_ref_time( &self, at: &Option, offset: &Option, diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index da89d2a7b6c..e417044ef43 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -547,6 +547,12 @@ pub async fn range_query( StatusCode::InvalidArguments ); let promql_expr = prom_query.expr(); + // Prometheus rejects a range-vector expression in a range query even when start == end, + // which the planner cannot tell apart from an instant query. + if promql_expr.value_type() == ValueType::Matrix { + let err = query::promql::error::RangeVectorInRangeQuerySnafu.build(); + return PrometheusJsonResponse::error(err.status_code(), err.output_msg()); + } let metric_name_discovery = try_call_return_response!(find_metric_name_not_equal_matchers(promql_expr)); @@ -2830,6 +2836,45 @@ mod tests { ); } + #[tokio::test] + async fn range_query_rejects_range_vector_with_equal_endpoints() { + let handler = Arc::new(TestPrometheusHandler { + catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + label_metric_names: Vec::new(), + label_lookups: Mutex::new(Vec::new()), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }); + let response = range_query( + State(handler.clone() as PrometheusHandlerRef), + Query(RangeQuery { + query: Some("vector(1)[5m:]".to_string()), + start: Some("300".to_string()), + end: Some("300".to_string()), + step: Some("1".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(RangeQuery::default()), + ) + .await; + + assert_eq!(Some(StatusCode::InvalidArguments), response.status_code); + assert_eq!( + response.error.as_deref(), + Some( + "invalid expression type \"range vector\" for range query, must be Scalar or instant Vector" + ) + ); + assert!(handler.queries.lock().unwrap().is_empty()); + } + #[tokio::test] async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index b3293bf6421..91979ddbf59 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2475,6 +2475,7 @@ pub async fn test_promql_over_non_millisecond_physical_tables(store_type: Storag "query_range?query=unit_gauge&start=0&end=5&step=1", "query?query=avg_over_time(unit_gauge[1m])&time=2", "query?query=rate(unit_gauge[1m])&time=2", + "query?query=unit_gauge[1m]&time=2", ]; for path in query_paths { let mut baseline: Option<(String, PrometheusResponse)> = None; @@ -2542,6 +2543,24 @@ pub async fn test_promql_over_non_millisecond_physical_tables(store_type: Storag .unwrap() ); + let res = client + .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=unit_gauge[1m]&time=2") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + body.data, + serde_json::from_value::(json!({ + "resultType": "matrix", + "result": [{ + "metric": {"__name__": "unit_gauge", "job": "demo"}, + "values": [[1.0, "1.0"], [2.0, "2.0"]] + }] + })) + .unwrap() + ); + let res = client .get("/v1/prometheus/api/v1/query?db=promql_units_ms&query=avg_over_time(unit_gauge[1m])&time=2") .send() diff --git a/tests/cases/standalone/common/promql/encode_substrait.result b/tests/cases/standalone/common/promql/encode_substrait.result index 230eadb02f7..53bd8a94b1e 100644 --- a/tests/cases/standalone/common/promql/encode_substrait.result +++ b/tests/cases/standalone/common/promql/encode_substrait.result @@ -12,40 +12,42 @@ Affected Rows: 0 -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED tql explain (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| plan_type | plan | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | -| | PromRangeManipulate: req range=[0..100000], interval=[1000], eval range=[43200000], time index=[ts], values=["prom_increase(ts_range,val,ts,Int64(3600000))"] | -| | PromSeriesDivide: tags=["tag_a", "tag_b"] | -| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | -| | Filter: prom_increase(ts_range,val,ts,Int64(3600000)) IS NOT NULL | -| | Projection: count_total.ts, prom_increase(ts_range, val, count_total.ts, Int64(3600000)) AS prom_increase(ts_range,val,ts,Int64(3600000)), count_total.tag_a, count_total.tag_b | -| | PromRangeManipulate: req range=[-39600000..100000], interval=[3600000], eval range=[3600000], time index=[ts], values=["val"] | -| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | -| | PromSeriesDivide: tags=["tag_a", "tag_b"] | -| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | -| | Filter: count_total.tag_a = Utf8("ffa") AND count_total.ts >= TimestampMillisecond(-43199999, None) AND count_total.ts <= TimestampMillisecond(100000, None) | -| | TableScan: count_total, partial_filters=[count_total.tag_a = Utf8("ffa"), count_total.ts >= TimestampMillisecond(-43199999, None), count_total.ts <= TimestampMillisecond(100000, None)] | -| | ]] | -| physical_plan | CooperativeExec | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Filter: prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))) IS NOT NULL | +| | Projection: count_total.ts, prom_max_over_time(ts_range, prom_increase(ts_range,val,ts,Int64(3600000))) AS prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))), count_total.tag_a, count_total.tag_b | +| | PromRangeManipulate: req range=[0..100000], interval=[1000], eval range=[43200000], time index=[ts], values=["prom_increase(ts_range,val,ts,Int64(3600000))"] | +| | PromSeriesDivide: tags=["tag_a", "tag_b"] | +| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | +| | Filter: prom_increase(ts_range,val,ts,Int64(3600000)) IS NOT NULL | +| | Projection: count_total.ts, prom_increase(ts_range, val, count_total.ts, Int64(3600000)) AS prom_increase(ts_range,val,ts,Int64(3600000)), count_total.tag_a, count_total.tag_b | +| | PromRangeManipulate: req range=[-39600000..100000], interval=[3600000], eval range=[3600000], time index=[ts], values=["val"] | +| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | +| | PromSeriesDivide: tags=["tag_a", "tag_b"] | +| | Sort: count_total.tag_a ASC NULLS FIRST, count_total.tag_b ASC NULLS FIRST, count_total.ts ASC NULLS FIRST | +| | Filter: count_total.tag_a = Utf8("ffa") AND count_total.ts >= TimestampMillisecond(-43199999, None) AND count_total.ts <= TimestampMillisecond(100000, None) | +| | TableScan: count_total, partial_filters=[count_total.tag_a = Utf8("ffa"), count_total.ts >= TimestampMillisecond(-43199999, None), count_total.ts <= TimestampMillisecond(100000, None)] | +| | ]] | +| physical_plan | CooperativeExec | | | MergeScanExec: REDACTED -| | | -+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ tql eval (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); -+----+-----------------------------------------------+-------+-------+----------+ -| ts | prom_increase(ts_range,val,ts,Int64(3600000)) | tag_a | tag_b | ts_range | -+----+-----------------------------------------------+-------+-------+----------+ -+----+-----------------------------------------------+-------+-------+----------+ ++----+----------------------------------------------------------------------------+-------+-------+ +| ts | prom_max_over_time(ts_range,prom_increase(ts_range,val,ts,Int64(3600000))) | tag_a | tag_b | ++----+----------------------------------------------------------------------------+-------+-------+ ++----+----------------------------------------------------------------------------+-------+-------+ drop table count_total; diff --git a/tests/cases/standalone/common/promql/encode_substrait.sql b/tests/cases/standalone/common/promql/encode_substrait.sql index 195c383ebd2..6443f3e1959 100644 --- a/tests/cases/standalone/common/promql/encode_substrait.sql +++ b/tests/cases/standalone/common/promql/encode_substrait.sql @@ -10,13 +10,13 @@ create table count_total ( -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (partitioning.*) REDACTED tql explain (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); tql eval (0, 100, '1s') - increase(count_total{ + max_over_time(increase(count_total{ tag_a="ffa", - }[1h])[12h:1h]; + }[1h])[12h:1h]); drop table count_total; diff --git a/tests/cases/standalone/common/promql/range_vector_result.result b/tests/cases/standalone/common/promql/range_vector_result.result new file mode 100644 index 00000000000..20d60c19480 --- /dev/null +++ b/tests/cases/standalone/common/promql/range_vector_result.result @@ -0,0 +1,143 @@ +-- An instant query of a range-vector expression returns the selected samples +-- with their own timestamps. Expected results are from Prometheus 3. +create table range_vector_result ( + ts timestamp time index, + val double, + host string primary key +); + +Affected Rows: 0 + +insert into range_vector_result values + (0, 1.0, 'a'), + (10000, 2.0, 'a'), + (20000, 3.0, 'a'), + (30000, 4.0, 'a'), + (0, 10.0, 'b'), + (20000, 30.0, 'b'); + +Affected Rows: 6 + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s]; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] offset 10s; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[15s] @ 20; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[30s:10s]; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:10 | 10.0 | b | +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | +| 1970-01-01T00:00:30 | 30.0 | b | +| 1970-01-01T00:00:30 | 4.0 | a | ++---------------------+------+------+ + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[35s:10s] offset 5s; + ++---------------------+------+------+ +| ts | val | host | ++---------------------+------+------+ +| 1970-01-01T00:00:00 | 1.0 | a | +| 1970-01-01T00:00:00 | 10.0 | b | +| 1970-01-01T00:00:10 | 10.0 | b | +| 1970-01-01T00:00:10 | 2.0 | a | +| 1970-01-01T00:00:20 | 3.0 | a | +| 1970-01-01T00:00:20 | 30.0 | b | ++---------------------+------+------+ + +tql eval (35, 35, '1s') range_vector_result[4s:10s]; + ++----+-----+------+ +| ts | val | host | ++----+-----+------+ ++----+-----+------+ + +tql eval (0, 30, '10s') range_vector_result[20s]; + +Error: 1004(InvalidArguments), invalid expression type "range vector" for range query, must be Scalar or instant Vector + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] as v; + ++------+------+---------------------+ +| v | host | ts | ++------+------+---------------------+ +| 3.0 | a | 1970-01-01T00:00:20 | +| 30.0 | b | 1970-01-01T00:00:20 | +| 4.0 | a | 1970-01-01T00:00:30 | ++------+------+---------------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +tql explain (30, 30, '1s') range_vector_result[20s]; + ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | +| | PromSeriesDivide: tags=["host"] | +| | Sort: range_vector_result.host ASC NULLS FIRST, range_vector_result.ts ASC NULLS FIRST | +| | Filter: range_vector_result.ts >= TimestampMillisecond(10001, None) AND range_vector_result.ts <= TimestampMillisecond(30000, None) | +| | TableScan: range_vector_result, partial_filters=[range_vector_result.ts >= TimestampMillisecond(10001, None), range_vector_result.ts <= TimestampMillisecond(30000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +tql eval (30, 30, '1s') range_vector_result[30s:0.0001]; + +Error: 1004(InvalidArguments), Subquery step must be positive + +-- A subquery without a step uses the one-minute default, not the query step. +tql eval (300, 300, '1s') vector(1)[5m:]; + ++---------------------+----------------+ +| time | greptime_value | ++---------------------+----------------+ +| 1970-01-01T00:01:00 | 1.0 | +| 1970-01-01T00:02:00 | 1.0 | +| 1970-01-01T00:03:00 | 1.0 | +| 1970-01-01T00:04:00 | 1.0 | +| 1970-01-01T00:05:00 | 1.0 | ++---------------------+----------------+ + +drop table range_vector_result; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/range_vector_result.sql b/tests/cases/standalone/common/promql/range_vector_result.sql new file mode 100644 index 00000000000..78fdbd81b4f --- /dev/null +++ b/tests/cases/standalone/common/promql/range_vector_result.sql @@ -0,0 +1,49 @@ +-- An instant query of a range-vector expression returns the selected samples +-- with their own timestamps. Expected results are from Prometheus 3. + +create table range_vector_result ( + ts timestamp time index, + val double, + host string primary key +); + +insert into range_vector_result values + (0, 1.0, 'a'), + (10000, 2.0, 'a'), + (20000, 3.0, 'a'), + (30000, 4.0, 'a'), + (0, 10.0, 'b'), + (20000, 30.0, 'b'); + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] offset 10s; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[15s] @ 20; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[30s:10s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[35s:10s] offset 5s; + +tql eval (35, 35, '1s') range_vector_result[4s:10s]; + +tql eval (0, 30, '10s') range_vector_result[20s]; + +-- SQLNESS SORT_RESULT 3 1 +tql eval (30, 30, '1s') range_vector_result[20s] as v; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +tql explain (30, 30, '1s') range_vector_result[20s]; + +tql eval (30, 30, '1s') range_vector_result[30s:0.0001]; + +-- A subquery without a step uses the one-minute default, not the query step. +tql eval (300, 300, '1s') vector(1)[5m:]; + +drop table range_vector_result;