Issue 1787 (#1827)

crates/nu-cli/src/commands/sum.rs

@@ -2,9 +2,11 @@ use crate::commands::WholeStreamCommand;
 use crate::prelude::*;
 use crate::utils::data_processing::{reducer_for, Reduce};
 use nu_errors::ShellError;
-use nu_protocol::{ReturnSuccess, ReturnValue, Signature, UntaggedValue, Value};
+use nu_protocol::{Dictionary, ReturnSuccess, ReturnValue, Signature, UntaggedValue, Value};
 use num_traits::identities::Zero;
 
+use indexmap::map::IndexMap;
+
 pub struct Sum;
 
 impl WholeStreamCommand for Sum {
@@ -54,13 +56,40 @@ impl WholeStreamCommand for Sum {
 
 fn sum(RunnableContext { mut input, .. }: RunnableContext) -> Result<OutputStream, ShellError> {
     let stream = async_stream! {
-        let mut values = input.drain_vec().await;
-
+        let mut values: Vec<Value> = input.drain_vec().await;
         let action = reducer_for(Reduce::Sum);
 
-        match action(Value::zero(), values) {
-            Ok(total) => yield ReturnSuccess::value(total),
-            Err(err) => yield Err(err),
+        if values.iter().all(|v| if let UntaggedValue::Primitive(_) = v.value {true} else {false}) {
+            let total = action(Value::zero(), values)?;
+            yield ReturnSuccess::value(total)
+        } else {
+            let mut column_values = IndexMap::new();
+            for value in values {
+                match value.value {
+                    UntaggedValue::Row(row_dict) => {
+                        for (key, value) in row_dict.entries.iter() {
+                            column_values
+                                .entry(key.clone())
+                                .and_modify(|v: &mut Vec<Value>| v.push(value.clone()))
+                                .or_insert(vec![value.clone()]);
+                        }
+                    },
+                    table => {},
+                };
+            }
+
+            let mut column_totals = IndexMap::new();
+            for (col_name, col_vals) in column_values {
+                let sum = action(Value::zero(), col_vals);
+                match sum {
+                    Ok(value) => {
+                        column_totals.insert(col_name, value);
+                    },
+                    Err(err) => yield Err(err),
+                };
+            }
+            yield ReturnSuccess::value(
+                UntaggedValue::Row(Dictionary {entries: column_totals}).into_untagged_value())
         }
     };
 

crates/nu-cli/src/utils/data_processing.rs

@@ -198,9 +198,20 @@ pub fn evaluate(
 }
 
 pub fn sum(data: Vec<Value>) -> Result<Value, ShellError> {
-    Ok(data
-        .into_iter()
-        .fold(Value::zero(), |acc: Value, value| acc + value))
+    let mut acc = Value::zero();
+    for value in data {
+        match value.value {
+            UntaggedValue::Primitive(_) => acc = acc + value,
+            _ => {
+                return Err(ShellError::labeled_error(
+                    "Attempted to compute the sum of a value that cannot be summed.",
+                    "value appears here",
+                    value.tag.span,
+                ))
+            }
+        }
+    }
+    Ok(acc)
 }
 
 fn formula(

crates/nu-cli/tests/commands/sum.rs

@@ -1,6 +1,7 @@
 use nu_test_support::fs::Stub::FileWithContentToBeTrimmed;
 use nu_test_support::playground::Playground;
 use nu_test_support::{nu, pipeline};
+use std::str::FromStr;
 
 #[test]
 fn all() {
@@ -60,3 +61,52 @@ fn outputs_zero_with_no_input() {
         assert_eq!(actual.out, "0");
     })
 }
+
+#[test]
+fn compute_sum_of_individual_row() -> Result<(), String> {
+    let answers_for_columns = [
+        ("cpu", 88.257434),
+        ("mem", 3032375296.),
+        ("virtual", 102579965952.),
+    ];
+    for (column_name, expected_value) in answers_for_columns.iter() {
+        let actual = nu!(
+            cwd: "tests/fixtures/formats/",
+            format!("open sample-ps-output.json | select {} | sum | get {}", column_name, column_name)
+        );
+        let result =
+            f64::from_str(&actual.out).map_err(|_| String::from("Failed to parse float."))?;
+        assert_eq!(result, *expected_value);
+    }
+    Ok(())
+}
+
+#[test]
+fn compute_sum_of_table() -> Result<(), String> {
+    let answers_for_columns = [
+        ("cpu", 88.257434),
+        ("mem", 3032375296.),
+        ("virtual", 102579965952.),
+    ];
+    for (column_name, expected_value) in answers_for_columns.iter() {
+        let actual = nu!(
+            cwd: "tests/fixtures/formats/",
+            format!("open sample-ps-output.json | select cpu mem virtual | sum | get {}", column_name)
+        );
+        let result =
+            f64::from_str(&actual.out).map_err(|_| String::from("Failed to parse float."))?;
+        assert_eq!(result, *expected_value);
+    }
+    Ok(())
+}
+
+#[test]
+fn sum_of_a_row_containing_a_table_is_an_error() {
+    let actual = nu!(
+        cwd: "tests/fixtures/formats/",
+        "open sample-sys-output.json | sum"
+    );
+    assert!(actual
+        .err
+        .contains("Attempted to compute the sum of a value that cannot be summed."));
+}