feat: add support for limitk and limit_ratio (#141)

Author: sunng87

src/parser/parse.rs

@@ -1256,6 +1256,32 @@ mod tests {
                 let ex = Expr::from(VectorSelector::from("sum"));
                 Expr::new_aggregate_expr(token::T_SUM, None, FunctionArgs::new_args(ex))
             }),
+            ("limitk by (group) (-1, http_requests)", {
+                let ex = Expr::from(VectorSelector::from("http_requests"));
+                let modifier = LabelModifier::include(vec!["group"]);
+                let param = Expr::from(-1.0);
+                let args = FunctionArgs::new_args(param).append_args(ex);
+                Expr::new_aggregate_expr(token::T_LIMITK, Some(modifier), args)
+            }),
+            ("limitk by (group) (1, http_requests)", {
+                let ex = Expr::from(VectorSelector::from("http_requests"));
+                let modifier = LabelModifier::include(vec!["group"]);
+                let param = Expr::from(1.0);
+                let args = FunctionArgs::new_args(param).append_args(ex);
+                Expr::new_aggregate_expr(token::T_LIMITK, Some(modifier), args)
+            }),
+            ("limit_ratio(-1.1, http_requests)", {
+                let ex = Expr::from(VectorSelector::from("http_requests"));
+                let param = Expr::from(-1.1);
+                let args = FunctionArgs::new_args(param).append_args(ex);
+                Expr::new_aggregate_expr(token::T_LIMIT_RATIO, None, args)
+            }),
+            ("limit_ratio(0.2, http_requests)", {
+                let ex = Expr::from(VectorSelector::from("http_requests"));
+                let param = Expr::from(0.2);
+                let args = FunctionArgs::new_args(param).append_args(ex);
+                Expr::new_aggregate_expr(token::T_LIMIT_RATIO, None, args)
+            }),
         ];
         assert_cases(Case::new_result_cases(cases));
 

src/parser/promql.y

@@ -118,7 +118,7 @@ START_METRIC_SELECTOR
 
 %expect-unused 'BLANK' 'COMMENT' 'ERROR' 'SEMICOLON' 'SPACE' 'TIMES' 'OPEN_HIST' 'CLOSE_HIST'
 %expect-unused 'OPERATORS_START' 'OPERATORS_END' 'AGGREGATORS_START' 'AGGREGATORS_END'
-%expect-unused 'LIMITK' 'LIMIT_RATIO' 'KEYWORDS_START' 'KEYWORDS_END' 'PREPROCESSOR_START' 'PREPROCESSOR_END'
+%expect-unused 'KEYWORDS_START' 'KEYWORDS_END' 'PREPROCESSOR_START' 'PREPROCESSOR_END'
 %expect-unused 'STEP' 'STARTSYMBOLS_START'
 %expect-unused 'START_METRIC' 'START_SERIES_DESCRIPTION' 'START_EXPRESSION' 'START_METRIC_SELECTOR' 'STARTSYMBOLS_END'
 %expect-unused 'SMOOTHED' 'ANCHORED'
@@ -572,6 +572,8 @@ metric_identifier -> Result<Token, String>:
         |       STDVAR { lexeme_to_token($lexer, $1) }
         |       SUM { lexeme_to_token($lexer, $1) }
         |       TOPK { lexeme_to_token($lexer, $1) }
+        |       LIMITK { lexeme_to_token($lexer, $1) }
+        |       LIMIT_RATIO { lexeme_to_token($lexer, $1) }
         |       WITHOUT { lexeme_to_token($lexer, $1) }
         |       START { lexeme_to_token($lexer, $1) }
         |       END { lexeme_to_token($lexer, $1) }
@@ -598,6 +600,8 @@ aggregate_op -> Result<Token, String>:
         |       STDVAR { lexeme_to_token($lexer, $1) }
         |       SUM { lexeme_to_token($lexer, $1) }
         |       TOPK { lexeme_to_token($lexer, $1) }
+        |       LIMITK { lexeme_to_token($lexer, $1) }
+        |       LIMIT_RATIO { lexeme_to_token($lexer, $1) }
 ;
 
 // inside of grouping options label names can be recognized as keywords by the lexer.
@@ -630,6 +634,8 @@ maybe_label -> Result<Token, String>:
         |       STDVAR { lexeme_to_token($lexer, $1) }
         |       SUM { lexeme_to_token($lexer, $1) }
         |       TOPK { lexeme_to_token($lexer, $1) }
+        |       LIMITK { lexeme_to_token($lexer, $1) }
+        |       LIMIT_RATIO { lexeme_to_token($lexer, $1) }
         |       START { lexeme_to_token($lexer, $1) }
         |       END { lexeme_to_token($lexer, $1) }
         |       ATAN2 { lexeme_to_token($lexer, $1) }

src/parser/token.rs

@@ -230,7 +230,10 @@ impl TokenType {
     }
 
     pub fn is_aggregator_with_param(&self) -> bool {
-        matches!(self.0, T_TOPK | T_BOTTOMK | T_COUNT_VALUES | T_QUANTILE)
+        matches!(
+            self.0,
+            T_TOPK | T_BOTTOMK | T_COUNT_VALUES | T_QUANTILE | T_LIMITK | T_LIMIT_RATIO
+        )
     }
 
     pub fn is_comparison_operator(&self) -> bool {
@@ -408,6 +411,8 @@ mod tests {
         assert!(TokenType(T_BOTTOMK).is_aggregator_with_param());
         assert!(TokenType(T_COUNT_VALUES).is_aggregator_with_param());
         assert!(TokenType(T_QUANTILE).is_aggregator_with_param());
+        assert!(TokenType(T_LIMITK).is_aggregator_with_param());
+        assert!(TokenType(T_LIMIT_RATIO).is_aggregator_with_param());
 
         assert!(!TokenType(T_MAX).is_aggregator_with_param());
         assert!(!TokenType(T_MIN).is_aggregator_with_param());
@@ -479,6 +484,8 @@ mod tests {
         assert!(TokenType(T_STDVAR).is_aggregator());
         assert!(TokenType(T_SUM).is_aggregator());
         assert!(TokenType(T_TOPK).is_aggregator());
+        assert!(TokenType(T_LIMITK).is_aggregator());
+        assert!(TokenType(T_LIMIT_RATIO).is_aggregator());
 
         assert!(!TokenType(T_LOR).is_aggregator());
         assert!(!TokenType(T_ADD).is_aggregator());