fixing test errors

Author: chelsea-lin

bigframes/operations/_matplotlib/core.py

@@ -16,14 +16,14 @@
 import typing
 import uuid
 
-import matplotlib.pyplot as plt
 import pandas as pd
 
 import bigframes.dtypes as dtypes
 
 DEFAULT_SAMPLING_N = 1000
 DEFAULT_SAMPLING_STATE = 0
 
+
 class MPLPlot(abc.ABC):
     @abc.abstractmethod
     def generate(self):
@@ -105,7 +105,6 @@ def _compute_plot_data(self):
         c_id = None
         if self._is_sequence_arg(c):
             c_id = self._generate_new_column_name(data)
-            print(c_id)
             data[c_id] = c
 
         sample = self._compute_sample_data(data)

tests/system/small/operations/test_plotting.py

@@ -259,7 +259,11 @@ def test_scatter_args_c_sampling():
     pd_df = pd.DataFrame(data)
 
     ax = df.plot.scatter(x="plot_temp_0", y="plot_temp_1", c=c, sampling_n=3)
-    pd_ax = pd_df.plot.scatter(x="plot_temp_0", y="plot_temp_1", c=c)
+
+    sampling_index = [0, 1, 2]
+    pd_ax = pd_df.iloc[sampling_index].plot.scatter(
+        x="plot_temp_0", y="plot_temp_1", c=[c[i] for i in sampling_index]
+    )
     assert len(ax.collections[0].get_facecolor()) == len(
         pd_ax.collections[0].get_facecolor()
     )