feat: support series.at[row_label] = scaler and loc equivalent

bigframes/core/indexers.py

@@ -59,23 +59,26 @@ def __setitem__(self, key, value) -> None:
 
         # Assume the key is for the index label.
         block = self._series._block
-        value_column = self._series._value_column
+        original_column = self._series
         index_column = block.index_columns[0]
-
-        # if index == key return value else value_colum
         block, insert_cond = block.apply_unary_op(
-            index_column, ops.partial_right(ops.eq_op, key)
+            index_column,
+            ops.partial_right(ops.eq_op, key),
+            result_label=self._series.name,
         )
-        block, result_id = block.apply_binary_op(
-            insert_cond,
-            self._series._value_column,
-            ops.partial_arg1(ops.where_op, value),
+        insert_cond_bool_series = bigframes.series.Series(
+            block.select_column(insert_cond)
         )
-        block = block.copy_values(result_id, value_column).drop_columns(
-            [insert_cond, result_id]
+        new_column = insert_cond_bool_series.map(
+            {True: value, False: None}, verify_integrity=False
         )
-
-        self._series._set_block(block)
+        try:
+            new_column = new_column.fillna(self._series)
+        except ibis.common.exceptions.IbisTypeError:
+            raise TypeError(
+                f"Cannot assign scalar of type {type(value)} to column of type {original_column.dtype}."
+            )
+        self._series._set_block(new_column._block)
 
 
 class IlocSeriesIndexer:
@@ -117,6 +120,20 @@ def __getitem__(
     ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]:
         return self._series.loc[key]
 
+    def __setitem__(
+        self,
+        key: LocSingleKey,
+        value,
+    ):
+        if pd.api.types.is_list_like(value) or isinstance(
+            value, bigframes.series.Series
+        ):
+            raise NotImplementedError(
+                "series.at.__setitem__ only supports scalar right-hand values. "
+                f"{constants.FEEDBACK_LINK}"
+            )
+        self._series.loc[key] = value
+
 
 class LocDataFrameIndexer:
     def __init__(self, dataframe: bigframes.dataframe.DataFrame):

tests/system/small/test_series.py

@@ -1012,6 +1012,16 @@ def test_loc_setitem_cell(scalars_df_index, scalars_pandas_df_index):
     pd.testing.assert_series_equal(bf_original.to_pandas(), pd_original)
 
 
+def test_at_setitem_row_label_scalar(scalars_df_index, scalars_pandas_df_index):
+    bf_series = scalars_df_index["int64_col"]
+    pd_series = scalars_pandas_df_index["int64_col"]
+    bf_series.at[1] = 1000
+    pd_series.at[1] = 1000
+    bf_result = bf_series.to_pandas()
+    pd_result = pd_series.astype("Float64")  # type difference is due to NA treatment
+    pd.testing.assert_series_equal(bf_result, pd_result)
+
+
 def test_ne_obj_series(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     col_name = "string_col"