[query-engine] Add select expression (open-telemetry#1363)

Relates to open-telemetry#1362

## Changes

* Adds an expression for selecting data from a scalar result using a
path potentially not known until execution time

## Details

Working towards KQL `extract_json` expression support.

Author: CodeBlanch

rust/experimental/query_engine/engine-recordset-otlp-bridge/src/attached_records.rs

@@ -49,12 +49,16 @@ impl MapValue for Resource {
         key == "Attributes"
     }
 
-    fn get(&self, key: &str) -> Option<&(dyn AsStaticValue + 'static)> {
+    fn get(&self, key: &str) -> Option<&dyn AsValue> {
+        self.get_static(key).unwrap().map(|v| v as &dyn AsValue)
+    }
+
+    fn get_static(&self, key: &str) -> Result<Option<&(dyn AsStaticValue + 'static)>, String> {
         if key == "Attributes" {
-            return Some(&self.attributes);
+            return Ok(Some(&self.attributes));
         }
 
-        None
+        Ok(None)
     }
 
     fn get_items(&self, item_callback: &mut dyn KeyValueCallback) -> bool {
@@ -81,13 +85,17 @@ impl MapValue for InstrumentationScope {
         matches!(key, "Attributes" | "Name" | "Version")
     }
 
-    fn get(&self, key: &str) -> Option<&(dyn AsStaticValue + 'static)> {
-        match key {
+    fn get(&self, key: &str) -> Option<&dyn AsValue> {
+        self.get_static(key).unwrap().map(|v| v as &dyn AsValue)
+    }
+
+    fn get_static(&self, key: &str) -> Result<Option<&(dyn AsStaticValue + 'static)>, String> {
+        Ok(match key {
             "Attributes" => Some(&self.attributes),
             "Name" => self.name.as_ref().map(|v| v as &dyn AsStaticValue),
             "Version" => self.version.as_ref().map(|v| v as &dyn AsStaticValue),
             _ => None,
-        }
+        })
     }
 
     fn get_items(&self, item_callback: &mut dyn KeyValueCallback) -> bool {

rust/experimental/query_engine/engine-recordset-otlp-bridge/src/logs.rs

@@ -79,8 +79,12 @@ impl MapValue for LogRecord {
         }
     }
 
-    fn get(&self, key: &str) -> Option<&(dyn AsStaticValue + 'static)> {
-        match key {
+    fn get(&self, key: &str) -> Option<&dyn AsValue> {
+        self.get_static(key).unwrap().map(|v| v as &dyn AsValue)
+    }
+
+    fn get_static(&self, key: &str) -> Result<Option<&(dyn AsStaticValue + 'static)>, String> {
+        Ok(match key {
             "Attributes" => Some(&self.attributes),
             "Timestamp" => self.timestamp.as_ref().map(|v| v as &dyn AsStaticValue),
             "ObservedTimestamp" => self
@@ -98,7 +102,7 @@ impl MapValue for LogRecord {
             "TraceFlags" => self.flags.as_ref().map(|v| v as &dyn AsStaticValue),
             "EventName" => self.event_name.as_ref().map(|v| v as &dyn AsStaticValue),
             _ => None,
-        }
+        })
     }
 
     fn get_items(&self, item_callback: &mut dyn KeyValueCallback) -> bool {

rust/experimental/query_engine/engine-recordset/src/execution_context.rs

@@ -125,13 +125,20 @@ impl<'a> ExecutionContextVariables<'a> {
     ) -> Option<Ref<'_, dyn AsStaticValue + 'static>> {
         let vars = self.local_variables.borrow();
 
-        let var = Ref::filter_map(vars, |v| v.get(name));
+        let var = Ref::filter_map(vars, |v| {
+            v.get_static(name)
+                .expect("Static access not supported by underlying map")
+        });
 
         if let Ok(v) = var {
             return Some(v);
         }
 
-        Ref::filter_map(self.global_variables.borrow(), |v| v.get(name)).ok()
+        Ref::filter_map(self.global_variables.borrow(), |v| {
+            v.get_static(name)
+                .expect("Static access not supported by underlying map")
+        })
+        .ok()
     }
 
     #[cfg(test)]

rust/experimental/query_engine/engine-recordset/src/primitives/owned_value.rs

@@ -144,6 +144,19 @@ impl From<&dyn ArrayValue> for ArrayValueStorage<OwnedValue> {
     }
 }
 
+impl From<&dyn MapValue> for MapValueStorage<OwnedValue> {
+    fn from(value: &dyn MapValue) -> Self {
+        let mut values: HashMap<Box<str>, _> = HashMap::new();
+
+        value.get_items(&mut KeyValueClosureCallback::new(|k, v| {
+            values.insert(k.into(), v.into());
+            true
+        }));
+
+        MapValueStorage::new(values)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/experimental/query_engine/engine-recordset/src/primitives/resolved_value.rs

@@ -194,6 +194,48 @@ impl<'a> ResolvedValue<'a> {
             ResolvedValue::Sequence(s) => Ok(ResolvedArrayValue::Sequence(s)),
         }
     }
+
+    pub fn try_resolve_map(self) -> Result<ResolvedMapValue<'a>, Self> {
+        if self.get_value_type() != ValueType::Map {
+            return Err(self);
+        }
+
+        match self {
+            ResolvedValue::Value(v) => {
+                if let Value::Map(s) = v {
+                    Ok(ResolvedMapValue::Value(s))
+                } else {
+                    panic!()
+                }
+            }
+            ResolvedValue::Borrowed(s, b) => {
+                match Ref::filter_map(Ref::clone(&b), |v| {
+                    if let StaticValue::Map(s) = v.to_static_value() {
+                        Some(s)
+                    } else {
+                        None
+                    }
+                }) {
+                    Ok(v) => Ok(ResolvedMapValue::Borrowed(BorrowedMapValue {
+                        source: s,
+                        _orig: b,
+                        value: v,
+                    })),
+                    Err(_) => panic!(),
+                }
+            }
+            ResolvedValue::Computed(o) => {
+                if let OwnedValue::Map(s) = o {
+                    Ok(ResolvedMapValue::Computed(s))
+                } else {
+                    panic!()
+                }
+            }
+            ResolvedValue::Slice(_) => panic!(),
+            ResolvedValue::List(_) => panic!(),
+            ResolvedValue::Sequence(_) => panic!(),
+        }
+    }
 }
 
 impl From<ResolvedValue<'_>> for OwnedValue {
@@ -896,6 +938,159 @@ impl Display for ResolvedArrayValue<'_> {
     }
 }
 
+#[derive(Debug)]
+pub enum ResolvedMapValue<'a> {
+    /// A value resolved from the expression tree or an attached record
+    Value(&'a dyn MapValue),
+
+    /// A value borrowed from the record being modified by the engine
+    Borrowed(BorrowedMapValue<'a>),
+
+    /// A value computed by the engine as the result of a dynamic expression
+    Computed(MapValueStorage<OwnedValue>),
+}
+
+#[derive(Debug)]
+pub struct BorrowedMapValue<'a> {
+    source: BorrowSource,
+    // Note: orig is not currently used but left in the code to support future
+    // needing to return a resolved map as a static value. See usage on
+    // BorrowedArrayValue for example
+    _orig: Ref<'a, dyn AsStaticValue + 'static>,
+    value: Ref<'a, dyn MapValue + 'static>,
+}
+
+impl<'a> ResolvedMapValue<'a> {
+    pub fn copy_if_borrowed_from_target(&mut self, target: &MutableValueExpression) -> bool {
+        let v = match self {
+            ResolvedMapValue::Borrowed(b) => Some((&b.source, &b.value)),
+            _ => None,
+        };
+
+        if let Some((s, v)) = v {
+            let writing_while_holding_borrow = match target {
+                MutableValueExpression::Source(_) => {
+                    matches!(s, BorrowSource::Source)
+                }
+                MutableValueExpression::Variable(_) => {
+                    matches!(s, BorrowSource::Variable)
+                }
+            };
+
+            if writing_while_holding_borrow {
+                *self = ResolvedMapValue::Computed((&**v).into());
+                return true;
+            }
+        }
+
+        false
+    }
+
+    pub fn take<FConvert, FTake, R>(
+        self,
+        keys: &[&str],
+        convert: FConvert,
+        mut take: FTake,
+    ) -> Result<(), ExpressionError>
+    where
+        FConvert: Fn(&str, ResolvedValue<'a>) -> Result<R, ExpressionError>,
+        FTake: FnMut(R),
+    {
+        match self {
+            ResolvedMapValue::Value(m) => {
+                for key in keys {
+                    let v = (convert)(
+                        key,
+                        ResolvedValue::Value(m.get(key).expect("Map key was not found").to_value()),
+                    )?;
+                    (take)(v);
+                }
+                Ok(())
+            }
+            ResolvedMapValue::Borrowed(b) => {
+                let m = &b.value;
+                for key in keys {
+                    match Ref::filter_map(Ref::clone(m), |m| {
+                        m.get_static(key)
+                            .expect("Borrowed map does not implement get_static")
+                    }) {
+                        Ok(v) => {
+                            let v = (convert)(key, ResolvedValue::Borrowed(b.source.clone(), v))?;
+                            (take)(v);
+                        }
+                        Err(_) => panic!("Map key was not found"),
+                    }
+                }
+                Ok(())
+            }
+            ResolvedMapValue::Computed(mut m) => {
+                let values = m.get_values_mut();
+
+                for key in keys {
+                    match values.remove(*key) {
+                        Some(v) => {
+                            let v = (convert)(key, ResolvedValue::Computed(v))?;
+                            (take)(v);
+                        }
+                        None => panic!("Map key was not found"),
+                    }
+                }
+                Ok(())
+            }
+        }
+    }
+
+    fn get_map(&self) -> &dyn MapValue {
+        match self {
+            ResolvedMapValue::Value(v) => *v,
+            ResolvedMapValue::Borrowed(b) => &*b.value,
+            ResolvedMapValue::Computed(c) => c,
+        }
+    }
+}
+
+impl MapValue for ResolvedMapValue<'_> {
+    fn is_empty(&self) -> bool {
+        self.get_map().is_empty()
+    }
+
+    fn len(&self) -> usize {
+        self.get_map().len()
+    }
+
+    fn contains_key(&self, key: &str) -> bool {
+        self.get_map().contains_key(key)
+    }
+
+    fn get(&self, key: &str) -> Option<&dyn AsValue> {
+        self.get_map().get(key)
+    }
+
+    fn get_static(&self, key: &str) -> Result<Option<&(dyn AsStaticValue + 'static)>, String> {
+        self.get_map().get_static(key)
+    }
+
+    fn get_items(&self, item_callback: &mut dyn KeyValueCallback) -> bool {
+        self.get_map().get_items(item_callback)
+    }
+}
+
+impl AsValue for ResolvedMapValue<'_> {
+    fn get_value_type(&self) -> ValueType {
+        ValueType::Map
+    }
+
+    fn to_value(&self) -> Value<'_> {
+        Value::Map(self.get_map())
+    }
+}
+
+impl Display for ResolvedMapValue<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.to_value().fmt(f)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/experimental/query_engine/engine-recordset/src/primitives/value_storage.rs

@@ -450,8 +450,12 @@ impl<T: EnumerableValueSource<T>> MapValue for MapValueStorage<T> {
         self.values.contains_key(key)
     }
 
-    fn get(&self, key: &str) -> Option<&(dyn AsStaticValue + 'static)> {
-        self.values.get(key).map(|v| v as &dyn AsStaticValue)
+    fn get(&self, key: &str) -> Option<&(dyn AsValue + 'static)> {
+        self.values.get(key).map(|v| v as &dyn AsValue)
+    }
+
+    fn get_static(&self, key: &str) -> Result<Option<&(dyn AsStaticValue + 'static)>, String> {
+        Ok(self.values.get(key).map(|v| v as &dyn AsStaticValue))
     }
 
     fn get_items(&self, item_callback: &mut dyn KeyValueCallback) -> bool {