Enable DataFrame setting value as list of labels. (Resolves #894) (#905)

Resolves #894

for below DataFrame,

```python
>>> df = ks.DataFrame([[1, 2], [4, 5], [7, 8]],
...                   index=['cobra', 'viper', 'sidewinder'],
...                   columns=['max_speed', 'shield'])
>>> df
            max_speed  shield
cobra               1       2
viper               4       5
sidewinder          7       8
```

now we can set value for all items matching the list of labels like Pandas.

```python
>>> df.loc[['viper', 'sidewinder'], ['shield']] = 50
>>> df
            max_speed  shield
cobra               1       2
viper               4      50
sidewinder          7      50
```

Author: itholic

databricks/koalas/indexing.py

@@ -283,12 +283,14 @@ class LocIndexer(object):
 
     **Setting values**
 
-    Setting value for all items matching the list of labels is not allowed
+    Setting value for all items matching the list of labels.
 
     >>> df.loc[['viper', 'sidewinder'], ['shield']] = 50
-    Traceback (most recent call last):
-     ...
-    databricks.koalas.exceptions.SparkPandasNotImplementedError: ...
+    >>> df
+                max_speed  shield
+    cobra               1       2
+    viper               4      50
+    sidewinder          7      50
 
     Setting value for an entire row is not allowed
 
@@ -303,17 +305,26 @@ class LocIndexer(object):
     >>> df
                 max_speed  shield
     cobra              30       2
-    viper              30       5
-    sidewinder         30       8
+    viper              30      50
+    sidewinder         30      50
+
+    Set value for an entire list of columns
+
+    >>> df.loc[:, ['max_speed', 'shield']] = 100
+    >>> df
+                max_speed  shield
+    cobra             100     100
+    viper             100     100
+    sidewinder        100     100
 
     Set value with Series
 
     >>> df.loc[:, 'shield'] = df['shield'] * 2
     >>> df
                 max_speed  shield
-    cobra              30       4
-    viper              30      10
-    sidewinder         30      16
+    cobra             100     200
+    viper             100     200
+    sidewinder        100     200
 
     **Getting values on a DataFrame with an index that has integer labels**
 
@@ -492,22 +503,44 @@ def __setitem__(self, key, value):
         rows_sel, cols_sel = key
 
         if (not isinstance(rows_sel, slice)) or (rows_sel != slice(None)):
-            raise SparkPandasNotImplementedError(
-                description="""Can only assign value to the whole dataframe, the row index
-                has to be `slice(None)` or `:`""",
-                pandas_function=".loc[..., ...] = ...",
-                spark_target_function="withColumn, select")
+            if isinstance(rows_sel, list):
+                if isinstance(cols_sel, str):
+                    cols_sel = [cols_sel]
+                kdf = self._kdf.copy()
+                for col_sel in cols_sel:
+                    # Uses `kdf` to allow operations on different DataFrames.
+                    # TODO: avoid temp column name or declare `__` prefix is
+                    #  reserved for Koalas' internal columns.
+                    kdf["__indexing_temp_col__"] = value
+                    new_col = kdf["__indexing_temp_col__"]._scol
+                    kdf[col_sel] = Series(kdf[col_sel]._internal.copy(
+                        scol=F.when(
+                            kdf._internal.index_scols[0].isin(rows_sel), new_col
+                        ).otherwise(kdf[col_sel]._scol)), anchor=kdf)
+                    kdf = kdf.drop(labels=['__indexing_temp_col__'])
+
+                self._kdf._internal = kdf._internal.copy()
+            else:
+                raise SparkPandasNotImplementedError(
+                    description="""Can only assign value to the whole dataframe, the row index
+                    has to be `slice(None)` or `:`""",
+                    pandas_function=".loc[..., ...] = ...",
+                    spark_target_function="withColumn, select")
 
-        if not isinstance(cols_sel, str):
-            raise ValueError("""only column names can be assigned""")
+        if not isinstance(cols_sel, (str, list)):
+            raise ValueError("""only column names or list of column names can be assigned""")
 
         if isinstance(value, DataFrame):
             if len(value.columns) == 1:
                 self._kdf[cols_sel] = _col(value)
             else:
                 raise ValueError("Only a dataframe with one column can be assigned")
         else:
-            self._kdf[cols_sel] = value
+            if isinstance(cols_sel, str):
+                cols_sel = [cols_sel]
+            if (not isinstance(rows_sel, list)) and (isinstance(cols_sel, list)):
+                for col_sel in cols_sel:
+                    self._kdf[col_sel] = value
 
 
 class ILocIndexer(object):

databricks/koalas/tests/test_indexing.py

@@ -507,6 +507,41 @@ def test_iloc_series(self):
         self.assert_eq(kseries.iloc[:1], pseries.iloc[:1])
         self.assert_eq(kseries.iloc[:-1], pseries.iloc[:-1])
 
+    def test_setitem(self):
+        pdf = pd.DataFrame([[1, 2], [4, 5], [7, 8]],
+                           index=['cobra', 'viper', 'sidewinder'],
+                           columns=['max_speed', 'shield'])
+        kdf = ks.from_pandas(pdf)
+
+        pdf.loc[['viper', 'sidewinder'], ['shield', 'max_speed']] = 10
+        kdf.loc[['viper', 'sidewinder'], ['shield', 'max_speed']] = 10
+        self.assert_eq(kdf, pdf)
+
+        pdf.loc[['viper', 'sidewinder'], 'shield'] = 50
+        kdf.loc[['viper', 'sidewinder'], 'shield'] = 50
+        self.assert_eq(kdf, pdf)
+
+        with self.assertRaisesRegex(ValueError,
+                                    'Only a dataframe with one column can be assigned'):
+            kdf.loc[:, 'max_speed'] = kdf
+        with self.assertRaisesRegex(ValueError,
+                                    'only column names or list of column names can be assigned'):
+            kdf.loc[['viper'], ('max_speed', 'shield')] = 10
+        msg = """Can only assign value to the whole dataframe, the row index
+        has to be `slice(None)` or `:`"""
+        msg = ("Can only assign value to the whole dataframe, the row index")
+        with self.assertRaisesRegex(SparkPandasNotImplementedError, msg):
+            kdf.loc['viper', 'max_speed'] = 10
+
+        pdf = pd.DataFrame([[1], [4], [7]],
+                           index=['cobra', 'viper', 'sidewinder'],
+                           columns=['max_speed'])
+        kdf = ks.from_pandas(pdf)
+
+        pdf.loc[:, 'max_speed'] = pdf
+        kdf.loc[:, 'max_speed'] = kdf
+        self.assert_eq(kdf, pdf)
+
     def test_iloc_raises(self):
         pdf = pd.DataFrame({"A": [1, 2], "B": [3, 4], "C": [5, 6]})
         kdf = ks.from_pandas(pdf)

databricks/koalas/tests/test_ops_on_diff_frames.py

@@ -395,6 +395,19 @@ def test_multi_index_assignment_frame(self):
 
         self.assert_eq(kdf.sort_index(), pdf.sort_index())
 
+    def test_loc_setitem(self):
+        pdf = pd.DataFrame(
+            [[1, 2], [4, 5], [7, 8]],
+            index=['cobra', 'viper', 'sidewinder'],
+            columns=['max_speed', 'shield'])
+        kdf = ks.DataFrame(pdf)
+        another_kdf = ks.DataFrame(pdf)
+
+        kdf.loc[['viper', 'sidewinder'], ['shield']] = another_kdf.max_speed
+        pdf.loc[['viper', 'sidewinder'], ['shield']] = pdf.max_speed
+
+        self.assert_eq(kdf.sort_index(), pdf.sort_index())
+
 
 class OpsOnDiffFramesDisabledTest(ReusedSQLTestCase, SQLTestUtils):
 
@@ -447,3 +460,14 @@ def test_assignment(self):
         with self.assertRaisesRegex(ValueError, "Cannot combine column argument"):
             kdf = ks.from_pandas(self.pdf1)
             kdf['c'] = self.kdf1.a
+
+    def test_loc_setitem(self):
+        pdf = pd.DataFrame(
+            [[1, 2], [4, 5], [7, 8]],
+            index=['cobra', 'viper', 'sidewinder'],
+            columns=['max_speed', 'shield'])
+        kdf = ks.DataFrame(pdf)
+        another_kdf = ks.DataFrame(pdf)
+
+        with self.assertRaisesRegex(ValueError, "Cannot combine column argument"):
+            kdf.loc[['viper', 'sidewinder'], ['shield']] = another_kdf.max_speed