do not support sequence, and support column position

Author: chelsea-lin

bigframes/operations/_matplotlib/core.py

@@ -18,6 +18,7 @@
 
 import pandas as pd
 
+import bigframes.constants as constants
 import bigframes.dtypes as dtypes
 
 DEFAULT_SAMPLING_N = 1000
@@ -92,32 +93,25 @@ def __init__(self, data, **kwargs) -> None:
         super().__init__(data, **kwargs)
 
         c = self.kwargs.get("c", None)
-        if self._is_sequence_arg(c) and len(c) != self.data.shape[0]:
-            raise ValueError(
-                f"'c' argument has {len(c)} elements, which is "
-                + f"inconsistent with 'x' and 'y' with size {self.data.shape[0]}"
+        if self._is_sequence_arg(c):
+            raise NotImplementedError(
+                f"Only support a single color string or a column name/posision. {constants.FEEDBACK_LINK}"
             )
 
     def _compute_plot_data(self):
-        data = self.data.copy()
-
-        c = self.kwargs.get("c", None)
-        c_id = None
-        if self._is_sequence_arg(c):
-            c_id = self._generate_new_column_name(data)
-            data[c_id] = c
-
-        sample = self._compute_sample_data(data)
+        sample = self._compute_sample_data(self.data)
 
         # Works around a pandas bug:
         # https://github.com/pandas-dev/pandas/commit/45b937d64f6b7b6971856a47e379c7c87af7e00a
+        c = self.kwargs.get("c", None)
+        if (
+            pd.core.dtypes.common.is_integer(c)
+            and not self.data.columns._holds_integer()
+        ):
+            c = self.data.columns[c]
         if self._is_column_name(c, sample) and sample[c].dtype == dtypes.STRING_DTYPE:
             sample[c] = sample[c].astype("object")
 
-        if c_id is not None:
-            self.kwargs["c"] = sample[c_id]
-            sample = sample.drop(columns=[c_id])
-
         return sample
 
     def _is_sequence_arg(self, arg):

tests/system/small/operations/test_plotting.py

@@ -215,15 +215,7 @@ def test_scatter(scalars_dfs):
         pytest.param("red", id="red"),
         pytest.param("c", id="int_column"),
         pytest.param("species", id="color_column"),
-        pytest.param(["red", "green", "blue"], id="color_sequence"),
-        pytest.param([3.4, 5.3, 2.0], id="number_sequence"),
-        pytest.param(
-            [3.4, 5.3],
-            id="length_mismatches_sequence",
-            marks=pytest.mark.xfail(
-                raises=ValueError,
-            ),
-        ),
+        pytest.param(3, id="column_index"),
     ],
 )
 def test_scatter_args_c(c):
@@ -248,32 +240,6 @@ def test_scatter_args_c(c):
         )
 
 
-def test_scatter_args_c_sampling():
-    data = {
-        "plot_temp_0": [1, 2, 3, 4, 5],
-        "plot_temp_1": [5, 4, 3, 2, 1],
-    }
-    c = ["red", "green", "blue", "orange", "black"]
-
-    df = bpd.DataFrame(data)
-    pd_df = pd.DataFrame(data)
-
-    ax = df.plot.scatter(x="plot_temp_0", y="plot_temp_1", c=c, sampling_n=3)
-
-    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()
-    )
-    for idx in range(len(ax.collections[0].get_facecolor())):
-        tm.assert_numpy_array_equal(
-            ax.collections[0].get_facecolor()[idx],
-            pd_ax.collections[0].get_facecolor()[idx],
-        )
-
-
 def test_sampling_plot_args_n():
     df = bpd.DataFrame(np.arange(bf_mpl.DEFAULT_SAMPLING_N * 10), columns=["one"])
     ax = df.plot.line()

third_party/bigframes_vendored/pandas/plotting/_core.py

@@ -266,10 +266,6 @@ def scatter(
 
                 - A single color string referred to by name, RGB or RGBA code,
                   for instance 'red' or '#a98d19'.
-                - A sequence of color strings referred to by name, RGB or RGBA
-                  code, which will be used for each point's color recursively. For
-                  instance ['green','yellow'] all points will be filled in green or
-                  yellow, alternatively.
                 - A column name or position whose values will be used to color the
                   marker points according to a colormap.