fix: plot.scatter s parameter cannot accept float-like column (#563)

Fixes internal b/330574847 🦕

Author: chelsea-lin

bigframes/operations/_matplotlib/core.py

@@ -14,7 +14,6 @@
 
 import abc
 import typing
-import uuid
 
 import pandas as pd
 
@@ -115,6 +114,18 @@ def _compute_plot_data(self):
         if self._is_column_name(c, sample) and sample[c].dtype == dtypes.STRING_DTYPE:
             sample[c] = sample[c].astype("object")
 
+        # To avoid Matplotlib's automatic conversion of `Float64` or `Int64` columns
+        # to `object` types (which breaks float-like behavior), this code proactively
+        # converts the column to a compatible format.
+        s = self.kwargs.get("s", None)
+        if pd.core.dtypes.common.is_integer(s):
+            s = self.data.columns[s]
+        if self._is_column_name(s, sample):
+            if sample[s].dtype == dtypes.INT_DTYPE:
+                sample[s] = sample[s].astype("int64")
+            elif sample[s].dtype == dtypes.FLOAT_DTYPE:
+                sample[s] = sample[s].astype("float64")
+
         return sample
 
     def _is_sequence_arg(self, arg):
@@ -130,9 +141,3 @@ def _is_column_name(self, arg, data):
             and pd.core.dtypes.common.is_hashable(arg)
             and arg in data.columns
         )
-
-    def _generate_new_column_name(self, data):
-        col_name = None
-        while col_name is None or col_name in data.columns:
-            col_name = f"plot_temp_{str(uuid.uuid4())[:8]}"
-        return col_name

tests/system/small/operations/test_plotting.py

@@ -240,6 +240,32 @@ def test_scatter_args_c(c):
         )
 
 
+@pytest.mark.parametrize(
+    ("s"),
+    [
+        pytest.param([10, 34, 50], id="int"),
+        pytest.param([1.0, 3.4, 5.0], id="float"),
+        pytest.param(
+            [True, True, False], id="bool", marks=pytest.mark.xfail(raises=ValueError)
+        ),
+    ],
+)
+def test_scatter_args_s(s):
+    data = {
+        "a": [1, 2, 3],
+        "b": [1, 2, 3],
+    }
+    data["s"] = s
+    df = bpd.DataFrame(data)
+    pd_df = pd.DataFrame(data)
+
+    ax = df.plot.scatter(x="a", y="b", s="s")
+    pd_ax = pd_df.plot.scatter(x="a", y="b", s="s")
+    tm.assert_numpy_array_equal(
+        ax.collections[0].get_sizes(), pd_ax.collections[0].get_sizes()
+    )
+
+
 @pytest.mark.parametrize(
     ("arg_name"),
     [