From 5725e55abb5c1de401457b205adeaae09b1dc3ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Tue, 20 Oct 2020 05:37:40 -0500
Subject: [PATCH] Flatten rows containing same sub-table columns with distinct
 column names. (#2684)

---
 crates/nu-cli/src/commands/flatten.rs   | 12 ++++-
 crates/nu-cli/tests/commands/flatten.rs | 64 +++++++++++++++----------
 crates/nu-protocol/src/value/dict.rs    |  5 ++
 3 files changed, 55 insertions(+), 26 deletions(-)

diff --git a/crates/nu-cli/src/commands/flatten.rs b/crates/nu-cli/src/commands/flatten.rs
index ac0963d32..b9cc39bcd 100644
--- a/crates/nu-cli/src/commands/flatten.rs
+++ b/crates/nu-cli/src/commands/flatten.rs
@@ -100,12 +100,20 @@ fn flat_value(
                 } = value
                 {
                     if column_requested.is_none() && !columns.is_empty() {
-                        out.insert_value(column, value.clone());
+                        if out.contains_key(&column) {
+                            out.insert_value(format!("{}_{}", column, column), value.clone());
+                        } else {
+                            out.insert_value(column, value.clone());
+                        }
                         continue;
                     }
 
                     for (k, v) in mapa.into_iter() {
-                        out.insert_value(k, v.clone());
+                        if out.contains_key(k) {
+                            out.insert_value(format!("{}_{}", column, k), v.clone());
+                        } else {
+                            out.insert_value(k, v.clone());
+                        }
                     }
                 } else if value.is_table() {
                     if tables_explicitly_flattened >= 1 && column_requested.is_some() {
diff --git a/crates/nu-cli/tests/commands/flatten.rs b/crates/nu-cli/tests/commands/flatten.rs
index c65553d70..c97f69b29 100644
--- a/crates/nu-cli/tests/commands/flatten.rs
+++ b/crates/nu-cli/tests/commands/flatten.rs
@@ -54,24 +54,16 @@ fn flatten_row_column_explictly() {
             r#"
                 [
                     {
-                        "origin": "Ecuador",
                         "people": {
                             "name": "Andres",
                             "meal": "arepa"
-                        },
-                        "code": { "id": 1, "references": 2},
-                        "tags": ["carbohydrate", "corn", "maiz"],
-                        "city": ["Guayaquil", "Samborondón"]
+                        }
                     },
                     {
-                        "origin": "USA",
                         "people": {
                             "name": "Katz",
                             "meal": "nurepa"
-                        },
-                        "code": { "id": 2, "references": 1},
-                        "tags": ["carbohydrate", "shell food", "amigos flavor"],
-                        "city": ["Oregon", "Brooklin"]
+                        }
                     }
                 ]
             "#,
@@ -87,30 +79,58 @@ fn flatten_row_column_explictly() {
 }
 
 #[test]
-fn flatten_table_columns_explictly() {
+fn flatten_row_columns_having_same_column_names_flats_separately() {
     Playground::setup("flatten_test_2", |dirs, sandbox| {
         sandbox.with_files(vec![FileWithContentToBeTrimmed(
             "katz.json",
             r#"
                 [
                     {
-                        "origin": "Ecuador",
                         "people": {
                             "name": "Andres",
                             "meal": "arepa"
                         },
-                        "code": { "id": 1, "references": 2},
-                        "tags": ["carbohydrate", "corn", "maiz"],
-                        "city": ["Guayaquil", "Samborondón"]
+                        "city": [{"name": "Guayaquil"}, {"name": "Samborondón"}]
+                    },
+                    {
+                        "people": {
+                            "name": "Katz",
+                            "meal": "nurepa"
+                        },
+                        "city": [{"name": "Oregon"}, {"name": "Brooklin"}]
+                    }
+                ]
+            "#,
+        )]);
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            "open katz.json | flatten | flatten people city | get city_name | count"
+        );
+
+        assert_eq!(actual.out, "4");
+    })
+}
+
+#[test]
+fn flatten_table_columns_explictly() {
+    Playground::setup("flatten_test_3", |dirs, sandbox| {
+        sandbox.with_files(vec![FileWithContentToBeTrimmed(
+            "katz.json",
+            r#"
+                [
+                    {
+                        "people": {
+                            "name": "Andres",
+                            "meal": "arepa"
+                        },
+                        "city": ["Guayaquil", "Samborondón"]
                     },
                     {
-                        "origin": "USA",
                         "people": {
                             "name": "Katz",
                             "meal": "nurepa"
                         },
-                        "code": { "id": 2, "references": 1},
-                        "tags": ["carbohydrate", "shell food", "amigos flavor"],
                         "city": ["Oregon", "Brooklin"]
                     }
                 ]
@@ -128,28 +148,24 @@ fn flatten_table_columns_explictly() {
 
 #[test]
 fn flatten_more_than_one_column_that_are_subtables_not_supported() {
-    Playground::setup("flatten_test_3", |dirs, sandbox| {
+    Playground::setup("flatten_test_4", |dirs, sandbox| {
         sandbox.with_files(vec![FileWithContentToBeTrimmed(
             "katz.json",
             r#"
                 [
                     {
-                        "origin": "Ecuador",
                         "people": {
                             "name": "Andres",
                             "meal": "arepa"
-                        },
-                        "code": { "id": 1, "references": 2},
+                        }
                         "tags": ["carbohydrate", "corn", "maiz"],
                         "city": ["Guayaquil", "Samborondón"]
                     },
                     {
-                        "origin": "USA",
                         "people": {
                             "name": "Katz",
                             "meal": "nurepa"
                         },
-                        "code": { "id": 2, "references": 1},
                         "tags": ["carbohydrate", "shell food", "amigos flavor"],
                         "city": ["Oregon", "Brooklin"]
                     }
diff --git a/crates/nu-protocol/src/value/dict.rs b/crates/nu-protocol/src/value/dict.rs
index 98201b9da..43350885e 100644
--- a/crates/nu-protocol/src/value/dict.rs
+++ b/crates/nu-protocol/src/value/dict.rs
@@ -252,6 +252,11 @@ impl TaggedDictBuilder {
     pub fn is_empty(&self) -> bool {
         self.dict.is_empty()
     }
+
+    /// Checks if given key exists
+    pub fn contains_key(&self, key: &str) -> bool {
+        self.dict.contains_key(key)
+    }
 }
 
 impl From<TaggedDictBuilder> for Value {