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()
  )