From 9f54d238ba6e36c35092a98d7ed73406e4cbf270 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Wed, 17 Jun 2020 12:34:26 -0500
Subject: [PATCH] Refactoring and more split-by flexibility.

---
 crates/nu-cli/src/commands/split_by.rs |  8 ++--
 crates/nu-cli/src/utils/data/split.rs  | 53 ++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 crates/nu-cli/src/utils/data/split.rs

diff --git a/crates/nu-cli/src/commands/split_by.rs b/crates/nu-cli/src/commands/split_by.rs
index ba44d32cdc..d7c299c993 100644
--- a/crates/nu-cli/src/commands/split_by.rs
+++ b/crates/nu-cli/src/commands/split_by.rs
@@ -112,17 +112,17 @@ pub fn suggestions(tried: Tagged<&str>, for_value: &Value) -> ShellError {
     possible_matches.sort();
 
     if !possible_matches.is_empty() {
-        return ShellError::labeled_error(
+        ShellError::labeled_error(
             "Unknown column",
             format!("did you mean '{}'?", possible_matches[0].1),
             tried.tag(),
-        );
+        )
     } else {
-        return ShellError::labeled_error(
+        ShellError::labeled_error(
             "Unknown column",
             "row does not contain this column",
             tried.tag(),
-        );
+        )
     }
 }
 
diff --git a/crates/nu-cli/src/utils/data/split.rs b/crates/nu-cli/src/utils/data/split.rs
new file mode 100644
index 0000000000..86bfeeb107
--- /dev/null
+++ b/crates/nu-cli/src/utils/data/split.rs
@@ -0,0 +1,53 @@
+use nu_errors::ShellError;
+use nu_protocol::{SpannedTypeName, TaggedDictBuilder, UntaggedValue, Value};
+use nu_source::Tag;
+
+use crate::utils::data::group;
+
+#[allow(clippy::type_complexity)]
+pub fn split(
+    value: &Value,
+    splitter: &Option<Box<dyn Fn(&Value) -> Result<String, ShellError> + Send>>,
+    tag: impl Into<Tag>,
+) -> Result<Value, ShellError> {
+    let tag = tag.into();
+
+    let mut splits = indexmap::IndexMap::new();
+
+    for (column, value) in value.row_entries() {
+        if !&value.is_table() {
+            return Err(ShellError::type_error(
+                "a table value",
+                value.spanned_type_name(),
+            ));
+        }
+
+        match group(&value, splitter, &tag) {
+            Ok(grouped) => {
+                for (split_label, subset) in grouped.row_entries() {
+                    let s = splits
+                        .entry(split_label.clone())
+                        .or_insert(indexmap::IndexMap::new());
+
+                    if !&subset.is_table() {
+                        return Err(ShellError::type_error(
+                            "a table value",
+                            subset.spanned_type_name(),
+                        ));
+                    }
+
+                    s.insert(column.clone(), subset.clone());
+                }
+            }
+            Err(err) => return Err(err),
+        }
+    }
+
+    let mut out = TaggedDictBuilder::new(&tag);
+
+    for (k, v) in splits.into_iter() {
+        out.insert_untagged(k, UntaggedValue::row(v));
+    }
+
+    Ok(out.into_value())
+}