Skip to content

Commit

Permalink
do not support sequence, and support column position
Browse files Browse the repository at this point in the history
  • Loading branch information
chelsea-lin committed Mar 22, 2024
1 parent 3340ccc commit 4af7b09
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 56 deletions.
28 changes: 11 additions & 17 deletions bigframes/operations/_matplotlib/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import pandas as pd

import bigframes.constants as constants
import bigframes.dtypes as dtypes

DEFAULT_SAMPLING_N = 1000
Expand Down Expand Up @@ -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):
Expand Down
36 changes: 1 addition & 35 deletions tests/system/small/operations/test_plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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()
Expand Down
4 changes: 0 additions & 4 deletions third_party/bigframes_vendored/pandas/plotting/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 4af7b09

Please sign in to comment.