From 4c8845db4d56610816390974b9e3e33d673776f3 Mon Sep 17 00:00:00 2001
From: Takuya UESHIN <ueshin@databricks.com>
Date: Thu, 9 Jul 2020 17:15:49 -0700
Subject: [PATCH] Enable to assign list.

---
 databricks/koalas/frame.py                | 16 ++++++++++++++++
 databricks/koalas/tests/test_dataframe.py | 15 +++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
index 36d3fcdee3..5b011f20bb 100644
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -10032,6 +10032,22 @@ def assign_columns(kdf, this_column_labels, that_column_labels):
                         yield (kdf._kser_for(this_label), this_label)
 
             kdf = align_diff_frames(assign_columns, self, value, fillna=False, how="left")
+        elif isinstance(value, list):
+            if len(self) != len(value):
+                raise ValueError("Length of values does not match length of index")
+
+            # TODO: avoid using default index?
+            with option_context(
+                "compute.default_index_type",
+                "distributed-sequence",
+                "compute.ops_on_diff_frames",
+                True,
+            ):
+                kdf = self.reset_index()
+                kdf[key] = ks.DataFrame(value)
+                kdf = kdf.set_index(kdf.columns[: len(self._internal.index_map)])
+                kdf.index.names = self.index.names
+
         elif isinstance(key, list):
             assert isinstance(value, DataFrame)
             # Same DataFrames.
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
index bb0b5ffc4e..8b4ae14acd 100644
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -92,6 +92,21 @@ def test_inplace(self):
         self.assert_eq(kdf, pdf)
         self.assert_eq(kser, pser)
 
+    def test_assign_list(self):
+        pdf, kdf = self.df_pair
+
+        pser = pdf.a
+        kser = kdf.a
+
+        pdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90]
+        kdf["x"] = [10, 20, 30, 40, 50, 60, 70, 80, 90]
+
+        self.assert_eq(kdf.sort_index(), pdf.sort_index())
+        self.assert_eq(kser, pser)
+
+        with self.assertRaisesRegex(ValueError, "Length of values does not match length of index"):
+            kdf["z"] = [10, 20, 30, 40, 50, 60, 70, 80]
+
     def test_dataframe_multiindex_columns(self):
         pdf = pd.DataFrame(
             {