From aada5c1706d16e939aadffb959b8a936c819c8eb Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 19 Feb 2025 12:51:27 +0900 Subject: [PATCH] fix(promql-planner): remove le tag in ctx (#5560) * fix(promql-planner): remove le tag in ctx * test: add sqlness test * chore: apply suggestions from CR --- src/query/src/promql/planner.rs | 94 +++++++++++++++++++ .../common/promql/simple_histogram.result | 10 ++ .../common/promql/simple_histogram.sql | 3 + 3 files changed, 107 insertions(+) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 1104dfa8c6..45ace2501e 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1803,6 +1803,8 @@ impl PromPlanner { fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), })? .clone(); + // remove le column from tag columns + self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); Ok(LogicalPlan::Extension(Extension { node: Arc::new( @@ -2620,6 +2622,68 @@ mod test { ) } + 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::default() + .schema(schema) + .primary_key_indices((0..num_tag).collect()) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name(table_name.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: 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, + ) + } + // { // input: `abs(some_metric{foo!="bar"})`, // expected: &Call{ @@ -3198,6 +3262,36 @@ mod test { indie_query_plan_compare(query, expected).await; } + #[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", "$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_session_state()) + .await + .unwrap(); + } + #[tokio::test] async fn value_matcher() { // template diff --git a/tests/cases/standalone/common/promql/simple_histogram.result b/tests/cases/standalone/common/promql/simple_histogram.result index be462cc4bb..4ea820b9e6 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.result +++ b/tests/cases/standalone/common/promql/simple_histogram.result @@ -115,6 +115,16 @@ tql eval (3000, 3000, '1s') histogram_quantile(0.8, histogram_bucket); | 1970-01-01T00:50:00 | positive | 0.72 | +---------------------+----------+------+ +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') label_replace(histogram_quantile(0.8, histogram_bucket), "s", "$1", "s", "(.*)tive"); + ++---------------------+------+------+ +| ts | val | s | ++---------------------+------+------+ +| 1970-01-01T00:50:00 | 0.3 | nega | +| 1970-01-01T00:50:00 | 0.72 | posi | ++---------------------+------+------+ + -- More realistic with rates. -- This case doesn't contains value because other point are not inserted. -- quantile with rate is covered in other cases diff --git a/tests/cases/standalone/common/promql/simple_histogram.sql b/tests/cases/standalone/common/promql/simple_histogram.sql index 3cd341b33b..8def0216d4 100644 --- a/tests/cases/standalone/common/promql/simple_histogram.sql +++ b/tests/cases/standalone/common/promql/simple_histogram.sql @@ -51,6 +51,9 @@ tql eval (3000, 3000, '1s') histogram_quantile(0.5, histogram_bucket); -- SQLNESS SORT_RESULT 3 1 tql eval (3000, 3000, '1s') histogram_quantile(0.8, histogram_bucket); +-- SQLNESS SORT_RESULT 3 1 +tql eval (3000, 3000, '1s') label_replace(histogram_quantile(0.8, histogram_bucket), "s", "$1", "s", "(.*)tive"); + -- More realistic with rates. -- This case doesn't contains value because other point are not inserted. -- quantile with rate is covered in other cases