Support x and y properly in plots (both matplotlib and plotly)

Author: HyukjinKwon

databricks/koalas/plot/core.py

@@ -117,7 +117,7 @@ def prepare_hist_data(data, bins):
 
         if is_integer(bins):
             # computes boundaries for the column
-            bins = HistogramPlotBase.get_bins(data.to_spark(), bins)
+            bins = HistogramPlotBase.get_bins(numeric_data.to_spark(), bins)
 
         return numeric_data, bins
 

databricks/koalas/plot/matplotlib.py

@@ -23,6 +23,7 @@
 from pandas.core.dtypes.inference import is_list_like
 from pandas.io.formats.printing import pprint_thing
 
+from databricks.koalas import Series
 from databricks.koalas.plot import (
     TopNPlotBase,
     SampledPlotBase,
@@ -373,6 +374,12 @@ class KoalasHistPlot(PandasHistPlot, HistogramPlotBase):
     def _args_adjust(self):
         if is_list_like(self.bottom):
             self.bottom = np.array(self.bottom)
+        y = self.kwds.get("y")
+        if y:
+            if isinstance(self.data, Series):
+                self.data = self.data.to_frame()
+            # When y is explicitly specified, we can drop other columns.
+            self.data = self.data[[y]]
 
     def _compute_plot_data(self):
         self.data, self.bins = HistogramPlotBase.prepare_hist_data(self.data, self.bins)

databricks/koalas/plot/plotly.py

@@ -74,6 +74,12 @@ def plot_histogram(data: Union["ks.DataFrame", "ks.Series"], **kwargs):
     import plotly.graph_objs as go
 
     bins = kwargs.get("bins", 10)
+    x = kwargs.get("x", None)
+    if x:
+        if isinstance(data, ks.Series):
+            data = data.to_frame()
+        # Note that it's opposite with matplotlib which takes y as the bars.
+        data = data[[x]]
     kdf, bins = HistogramPlotBase.prepare_hist_data(data, bins)
     assert len(bins) > 2, "the number of buckets must be higher than 2."
     output_series = HistogramPlotBase.compute_hist(kdf, bins)

databricks/koalas/tests/plot/test_frame_plot_matplotlib.py

@@ -402,6 +402,19 @@ def check_hist_plot(pdf, kdf):
             bin2 = self.plot_to_base64(ax2)
             self.assertEqual(bin1, bin2)
 
+            non_numeric_pdf = self.pdf1.copy()
+            non_numeric_pdf.c = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"]
+            non_numeric_kdf = ks.from_pandas(non_numeric_pdf)
+            ax1 = non_numeric_pdf.plot.hist(
+                x=non_numeric_pdf.columns[0], y=non_numeric_pdf.columns[1], bins=3
+            )
+            bin1 = self.plot_to_base64(ax1)
+            ax2 = non_numeric_kdf.plot.hist(
+                x=non_numeric_pdf.columns[0], y=non_numeric_pdf.columns[1], bins=3
+            )
+            bin2 = self.plot_to_base64(ax2)
+            self.assertEqual(bin1, bin2)
+
         pdf1 = self.pdf1
         kdf1 = self.kdf1
         check_hist_plot(pdf1, kdf1)