[query-engine] Relax restrictions in KQL summarize expressions (open-telemetry#1262)

## Changes

* Allow aggregate and group by expression without explicit assignment in
KQL summarize expressions

## Details

Today you can do:
* `summarize by Timestamp = bin(Timestamp, 1m)`
* `summarize Count = count()`

But you will get an error if you do:
* `summarize by bin(Timestamp, 1m)`
* `summarize count()`

This PR makes the second two examples work. In the first case the name
`Timestamp` will be used and in the second case `count`.

This is done by giving the KQL parser some logic to determine the
identifier for any given scalar expression. Note: We could probably use
this in extend\project* cases too but I plan to look at that in a
follow-up.

Fixes open-telemetry#1243

---------

Co-authored-by: Drew Relmas <[email protected]>

Author: CodeBlanch

rust/experimental/query_engine/engine-recordset/src/summary/summary_data_expression.rs

@@ -61,7 +61,7 @@ pub fn execute_summary_data_expression<'a, TRecord: Record>(
                     None => {
                         execution_context.add_diagnostic_if_enabled(
                             RecordSetEngineDiagnosticLevel::Warn,
-                            expression.get_value_expression().as_ref().unwrap(),
+                            expression.get_value_expression().unwrap(),
                             || {
                                 format!(
                                     "Value expression value of '{:?}' type could not be converted to integer or double",
@@ -104,7 +104,7 @@ pub fn execute_summary_data_expression<'a, TRecord: Record>(
                     None => {
                         execution_context.add_diagnostic_if_enabled(
                             RecordSetEngineDiagnosticLevel::Warn,
-                            expression.get_value_expression().as_ref().unwrap(),
+                            expression.get_value_expression().unwrap(),
                             || {
                                 format!(
                                     "Value expression value of '{:?}' type could not be converted to integer or double",

rust/experimental/query_engine/expressions/src/summary/summary_data_expression.rs

@@ -108,8 +108,8 @@ impl AggregationExpression {
         self.aggregation_function.clone()
     }
 
-    pub fn get_value_expression(&self) -> &Option<ScalarExpression> {
-        &self.value_expression
+    pub fn get_value_expression(&self) -> Option<&ScalarExpression> {
+        self.value_expression.as_ref()
     }
 
     pub(crate) fn try_fold(

rust/experimental/query_engine/kql-parser/src/aggregate_expressions.rs

@@ -5,67 +5,107 @@ use data_engine_expressions::*;
 use data_engine_parser_abstractions::*;
 use pest::iterators::Pair;
 
-use crate::{Rule, scalar_expression::parse_scalar_expression};
+use crate::{
+    Rule,
+    scalar_expression::{parse_scalar_expression, try_resolve_identifier},
+};
 
-pub(crate) fn parse_aggregate_assignment_expression(
-    aggregate_assignment_expression_rule: Pair<Rule>,
+pub(crate) fn parse_aggregate_expression(
+    aggregate_expression_rule: Pair<Rule>,
     scope: &dyn ParserScope,
 ) -> Result<(Box<str>, AggregationExpression), ParserError> {
-    let mut aggregate_assignment_rules = aggregate_assignment_expression_rule.into_inner();
+    let mut aggregate_expression_rules = aggregate_expression_rule.into_inner();
+
+    let first_rule = aggregate_expression_rules.next().unwrap();
+    let first_rule_location = to_query_location(&first_rule);
 
-    let destination_rule = aggregate_assignment_rules.next().unwrap();
-    let destination_rule_location = to_query_location(&destination_rule);
+    match first_rule.as_rule() {
+        Rule::identifier_literal => {
+            let identifier = crate::tabular_expressions::validate_summary_identifier(
+                scope,
+                first_rule_location,
+                &[first_rule.as_str().trim().into()],
+            )?;
 
-    let identifier = destination_rule.as_str().trim();
+            let aggregation_expression =
+                parse_aggregate_expressions(aggregate_expression_rules.next().unwrap(), scope)?;
 
-    crate::tabular_expressions::validate_summary_identifier(
-        scope,
-        destination_rule_location,
-        identifier,
-    )?;
+            Ok((identifier.into(), aggregation_expression))
+        }
+        _ => {
+            let aggregation_expression = parse_aggregate_expressions(first_rule, scope)?;
 
-    let aggregation_expression =
-        parse_aggregate_expression(aggregate_assignment_rules.next().unwrap(), scope)?;
+            let identifier = resolve_identifier(&aggregation_expression, scope)?;
 
-    Ok((identifier.into(), aggregation_expression))
+            let full_identifier = crate::tabular_expressions::validate_summary_identifier(
+                scope,
+                first_rule_location,
+                &identifier,
+            )?;
+
+            Ok((full_identifier.into(), aggregation_expression))
+        }
+    }
 }
 
-fn parse_aggregate_expression(
-    aggregate_expression_rule: Pair<Rule>,
+fn parse_aggregate_expressions(
+    aggregate_expressions_rule: Pair<Rule>,
     scope: &dyn ParserScope,
 ) -> Result<AggregationExpression, ParserError> {
-    let query_location = to_query_location(&aggregate_expression_rule);
+    let query_location = to_query_location(&aggregate_expressions_rule);
 
-    let aggregate = match aggregate_expression_rule.as_rule() {
+    let aggregate = match aggregate_expressions_rule.as_rule() {
         Rule::average_aggregate_expression => AggregationExpression::new(
             query_location,
             AggregationFunction::Average,
-            Some(parse_scalar_expression(aggregate_expression_rule, scope)?),
+            Some(parse_scalar_expression(aggregate_expressions_rule, scope)?),
         ),
         Rule::count_aggregate_expression => {
             AggregationExpression::new(query_location, AggregationFunction::Count, None)
         }
         Rule::maximum_aggregate_expression => AggregationExpression::new(
             query_location,
             AggregationFunction::Maximum,
-            Some(parse_scalar_expression(aggregate_expression_rule, scope)?),
+            Some(parse_scalar_expression(aggregate_expressions_rule, scope)?),
         ),
         Rule::minimum_aggregate_expression => AggregationExpression::new(
             query_location,
             AggregationFunction::Minimum,
-            Some(parse_scalar_expression(aggregate_expression_rule, scope)?),
+            Some(parse_scalar_expression(aggregate_expressions_rule, scope)?),
         ),
         Rule::sum_aggregate_expression => AggregationExpression::new(
             query_location,
             AggregationFunction::Sum,
-            Some(parse_scalar_expression(aggregate_expression_rule, scope)?),
+            Some(parse_scalar_expression(aggregate_expressions_rule, scope)?),
         ),
-        _ => panic!("Unexpected rule in aggregate_expression: {aggregate_expression_rule}"),
+        _ => panic!("Unexpected rule in aggregate_expression: {aggregate_expressions_rule}"),
     };
 
     Ok(aggregate)
 }
 
+fn resolve_identifier(
+    aggregation_expression: &AggregationExpression,
+    scope: &dyn ParserScope,
+) -> Result<Vec<Box<str>>, ParserError> {
+    let f = match aggregation_expression.get_aggregation_function() {
+        AggregationFunction::Average => "avg",
+        AggregationFunction::Count => "count",
+        AggregationFunction::Maximum => "max",
+        AggregationFunction::Minimum => "min",
+        AggregationFunction::Sum => "sum",
+    };
+
+    if let Some(s) = &aggregation_expression.get_value_expression()
+        && let Some(mut i) = try_resolve_identifier(s, scope)?
+    {
+        i.insert(0, f.into());
+        return Ok(i);
+    }
+
+    Ok(vec![f.into()])
+}
+
 #[cfg(test)]
 mod tests {
     use pest::Parser;
@@ -81,11 +121,9 @@ mod tests {
 
             let state = ParserState::new(input);
 
-            let mut result =
-                KqlPestParser::parse(Rule::aggregate_assignment_expression, input).unwrap();
+            let mut result = KqlPestParser::parse(Rule::aggregate_expression, input).unwrap();
 
-            let expression =
-                parse_aggregate_assignment_expression(result.next().unwrap(), &state).unwrap();
+            let expression = parse_aggregate_expression(result.next().unwrap(), &state).unwrap();
 
             assert_eq!(expected, expression);
         };
@@ -114,7 +152,7 @@ mod tests {
 
             let expression = parse_aggregate_expression(result.next().unwrap(), &state).unwrap();
 
-            assert_eq!(expected, expression);
+            assert_eq!(expected, expression.1);
         };
 
         run_test_success(

rust/experimental/query_engine/kql-parser/src/kql.pest

@@ -284,26 +284,27 @@ sum_aggregate_expression = {
     "sum" ~ "(" ~ scalar_expression ~ ")"
 }
 
-aggregate_expression = _{
+aggregate_expressions = _{
     average_aggregate_expression
     | count_aggregate_expression
     | maximum_aggregate_expression
     | minimum_aggregate_expression
     | sum_aggregate_expression
 }
 
-aggregate_assignment_expression = {
-    identifier_literal ~ "=" ~ aggregate_expression
+aggregate_expression = {
+    identifier_literal ~ "=" ~ aggregate_expressions
+    | aggregate_expressions
 }
 
 group_by_expression = {
     identifier_literal ~ "=" ~ scalar_expression
-    | accessor_expression
+    | scalar_expression
 }
 
 summarize_expression = {
     "summarize"
-        ~ (aggregate_assignment_expression ~ ("," ~ aggregate_assignment_expression)*)?
+        ~ (aggregate_expression ~ ("," ~ aggregate_expression)*)?
         ~ ("by" ~ group_by_expression ~ ("," ~ group_by_expression)*)?
         ~ ("|" ~ tabular_expressions)*
 }

rust/experimental/query_engine/kql-parser/src/scalar_expression.rs

@@ -310,6 +310,143 @@ pub(crate) fn parse_scalar_unary_expression(
     })
 }
 
+pub(crate) fn try_resolve_identifier(
+    scalar_expression: &ScalarExpression,
+    scope: &dyn ParserScope,
+) -> Result<Option<Vec<Box<str>>>, ParserError> {
+    let r = match scalar_expression {
+        ScalarExpression::Attached(a) => {
+            parse_identifier_from_accessor(a.get_name().get_value(), a.get_value_accessor(), scope)
+        }
+        ScalarExpression::Constant(c) => {
+            let name = scope
+                .get_constant_name(c.get_constant_id())
+                .expect("Constant not found");
+
+            parse_identifier_from_accessor(name, c.get_value_accessor(), scope)
+        }
+        ScalarExpression::Source(s) => {
+            parse_identifier_from_accessor("source", s.get_value_accessor(), scope)
+        }
+        ScalarExpression::Variable(v) => {
+            parse_identifier_from_accessor(v.get_name().get_value(), v.get_value_accessor(), scope)
+        }
+        ScalarExpression::Math(m) => match m {
+            MathScalarExpression::Add(_) => Ok(None),
+            MathScalarExpression::Bin(b) => try_resolve_identifier(b.get_left_expression(), scope),
+            MathScalarExpression::Ceiling(u) => {
+                try_resolve_identifier(u.get_value_expression(), scope)
+            }
+            MathScalarExpression::Divide(_) => Ok(None),
+            MathScalarExpression::Floor(u) => {
+                try_resolve_identifier(u.get_value_expression(), scope)
+            }
+            MathScalarExpression::Modulus(_) => Ok(None),
+            MathScalarExpression::Multiply(_) => Ok(None),
+            MathScalarExpression::Negate(_) => Ok(None),
+            MathScalarExpression::Subtract(_) => Ok(None),
+        },
+        ScalarExpression::Convert(c) => match c {
+            ConvertScalarExpression::Boolean(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+            ConvertScalarExpression::DateTime(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+            ConvertScalarExpression::Double(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+            ConvertScalarExpression::Integer(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+            ConvertScalarExpression::String(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+            ConvertScalarExpression::TimeSpan(c) => {
+                try_resolve_identifier(c.get_inner_expression(), scope)
+            }
+        },
+        ScalarExpression::Case(_) => Ok(None),
+        ScalarExpression::Coalesce(_) => Ok(None),
+        ScalarExpression::Collection(_) => Ok(None),
+        ScalarExpression::Conditional(_) => Ok(None),
+        ScalarExpression::Temporal(_) => Ok(None),
+        ScalarExpression::Length(l) => {
+            if let Some(mut i) = try_resolve_identifier(l.get_inner_expression(), scope)? {
+                i.insert(0, "len".into());
+                return Ok(Some(i));
+            }
+
+            Ok(None)
+        }
+        ScalarExpression::Logical(_) => Ok(None),
+        ScalarExpression::Parse(_) => Ok(None),
+        ScalarExpression::Slice(_) => Ok(None),
+        ScalarExpression::Static(_) => Ok(None),
+        ScalarExpression::Text(_) => Ok(None),
+    };
+
+    if let Ok(Some(mut identifier)) = r {
+        // Note: The identifier path may contain source.[default_map_key]. We always
+        // remove "source" and then remove the default_map_key if the mode is
+        // enabled.
+        if let Some("source") = identifier.first().map(|v| v.as_ref()) {
+            identifier.remove(0);
+
+            if let Some(schema) = scope.get_source_schema()
+                && let Some((key, _)) = schema.get_default_map()
+                && let Some(first) = identifier.first().map(|v| v.as_ref())
+                && key == first
+            {
+                identifier.remove(0);
+            }
+        }
+        return Ok(Some(identifier));
+    }
+
+    r
+}
+
+fn parse_identifier_from_accessor(
+    root: &str,
+    value_accessor: &ValueAccessor,
+    scope: &dyn ParserScope,
+) -> Result<Option<Vec<Box<str>>>, ParserError> {
+    let mut identifier: Vec<Box<str>> = vec![root.into()];
+    for selector in value_accessor.get_selectors() {
+        match selector {
+            ScalarExpression::Static(s) => match s.to_value() {
+                Value::String(s) => {
+                    identifier.push(s.get_value().into());
+                }
+                Value::Integer(i) => {
+                    identifier.push(format!("{}", i.get_value()).into());
+                }
+                _ => {
+                    return Ok(None);
+                }
+            },
+            ScalarExpression::Constant(c) => {
+                let name = scope
+                    .get_constant_name(c.get_constant_id())
+                    .expect("Constant not found");
+
+                if let Some(mut i) =
+                    parse_identifier_from_accessor(name, c.get_value_accessor(), scope)?
+                {
+                    identifier.append(&mut i);
+                } else {
+                    return Ok(None);
+                }
+            }
+            _ => {
+                return Ok(None);
+            }
+        }
+    }
+    Ok(Some(identifier))
+}
+
 #[cfg(test)]
 mod tests {
     use pest::Parser;

rust/experimental/query_engine/kql-parser/src/scalar_primitive_expressions.rs

@@ -544,7 +544,7 @@ pub(crate) fn parse_accessor_expression(
         //   source['const_str'] = 1 so the const_str is not evaluated.
         if allow_root_scalar {
             if let Some((constant_id, value_type)) =
-                scope.get_constant(root_accessor_identity.get_value())
+                scope.get_constant_id(root_accessor_identity.get_value())
             {
                 return Ok(ScalarExpression::Constant(
                     ReferenceConstantScalarExpression::new(