Fix Dataframe.melt function & Add doctest case for melt function

databricks/koalas/frame.py

@@ -6782,6 +6782,27 @@ def melt(self, id_vars=None, value_vars=None, var_name=None,
         0  a         B          1
         1  b         B          3
         2  c         B          5
+
+        >>> df.melt(id_vars='Z')  # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+          ...
+        KeyError: "The following 'id_vars' are not present in the DataFrame: [('Z',)]"
+
+        >>> df.melt(value_vars='Z')  # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+          ...
+        KeyError: "The following 'value_vars' are not present in the DataFrame: [('Z',)]"
+
+        >>> df.columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C')])
+        >>> df.melt(id_vars=('A',))  # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+          ...
+        ValueError: id_vars must be a list of tuples when columns are a MultiIndex
+
+        >>> df.melt(value_vars=('A',))  # doctest: +ELLIPSIS
+        Traceback (most recent call last):
+          ...
+        ValueError: value_vars must be a list of tuples when columns are a MultiIndex
         """
         if id_vars is None:
             id_vars = []
@@ -6797,14 +6818,29 @@ def melt(self, id_vars=None, value_vars=None, var_name=None,
 
         column_index = self._internal.column_index
 
+        missing = [True for id_var in id_vars if id_var not in column_index]
+        if True in missing:
+            raise KeyError("The following 'id_vars' are not present"
+                           " in the DataFrame: {}"
+                           .format(id_vars))
+
         if value_vars is None:
             value_vars = []
         elif isinstance(value_vars, str):
             value_vars = [(value_vars,)]
         elif isinstance(value_vars, tuple):
-            value_vars = [value_vars]
+            if self._internal.column_index_level == 1:
+                value_vars = [valv if isinstance(valv, tuple) else (valv,) for valv in value_vars]
+            else:
+                raise ValueError('value_vars must be a list of tuples'
+                                 ' when columns are a MultiIndex')
         else:
             value_vars = [valv if isinstance(valv, tuple) else (valv,) for valv in value_vars]
+        missing = [True for value_var in value_vars if value_var not in column_index]
+        if True in missing:
+            raise KeyError("The following 'value_vars' are not present"
+                           " in the DataFrame: {}"
+                           .format(value_vars))
         if len(value_vars) == 0:
             value_vars = column_index
 

databricks/koalas/tests/test_dataframe.py

@@ -1822,6 +1822,9 @@ def test_melt(self):
                        pdf.melt(id_vars=['A'], value_vars=['B'],
                                 var_name='myVarname', value_name='myValname')
                        .sort_values(['myVarname', 'myValname']))
+        self.assert_eq(kdf.melt(value_vars=('A', 'B')).sort_values(['variable', 'value'])
+                       .reset_index(drop=True),
+                       pdf.melt(value_vars=('A', 'B')).sort_values(['variable', 'value']))
 
         # multi-index columns
         columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C')])