resolve the comments

Author: jialuoo

bigframes/series.py

@@ -1904,6 +1904,15 @@ def apply(
         *,
         args: typing.Tuple = (),
     ) -> Series:
+        # Note: This signature differs from pandas.Series.apply. Specifically,
+        # `args` is keyword-only and `by_row` is a custom parameter here. Full
+        # alignment would involve breaking changes. However, given that by_row
+        # is not frequently used, we defer any such changes until there is a
+        # clear need based on user feedback.
+        #
+        # See pandas docs for reference:
+        # https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.apply.html
+
         # TODO(shobs, b/274645634): Support convert_dtype, **kwargs
         # is actually a ternary op
 
@@ -1952,6 +1961,9 @@ def apply(
             result_series = self._apply_nary_op(
                 ops.NaryRemoteFunctionOp(function_def=func.udf_def), args
             )
+            # TODO(jialuo): Investigate why `_apply_nary_op` drops the series
+            # `name`. Manually reassigning it here as a temporary fix.
+            result_series.name = self.name
         else:
             result_series = self._apply_unary_op(
                 ops.RemoteFunctionOp(function_def=func.udf_def, apply_on_null=True)

tests/system/large/functions/test_managed_function.py

@@ -1124,22 +1124,17 @@ def foo_list(x: int, y0: float, y1: bytes, y2: bool) -> list[str]:
 
         scalars_df, scalars_pandas_df = scalars_dfs
 
-        bf_result_col = scalars_df["int64_too"].apply(
-            foo_list, args=(12.34, b"hello world", False)
-        )
         bf_result = (
-            scalars_df["int64_too"].to_frame().assign(result=bf_result_col).to_pandas()
+            scalars_df["int64_too"]
+            .apply(foo_list, args=(12.34, b"hello world", False))
+            .to_pandas()
         )
-
-        pd_result_col = scalars_pandas_df["int64_too"].apply(
+        pd_result = scalars_pandas_df["int64_too"].apply(
             foo_list, args=(12.34, b"hello world", False)
         )
-        pd_result = (
-            scalars_pandas_df["int64_too"].to_frame().assign(result=pd_result_col)
-        )
 
         # Ignore any dtype difference.
-        pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False)
+        pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False)
 
     finally:
         # Clean up the gcp assets created for the managed function.

tests/system/large/functions/test_remote_function.py

@@ -2990,34 +2990,20 @@ def foo(x: int, y: bool, z: float) -> str:
         scalars_df, scalars_pandas_df = scalars_dfs
 
         args1 = (True, 10.0)
-        bf_result_col = scalars_df["int64_too"].apply(foo, args=args1)
-        bf_result = (
-            scalars_df["int64_too"].to_frame().assign(result=bf_result_col).to_pandas()
-        )
-
-        pd_result_col = scalars_pandas_df["int64_too"].apply(foo, args=args1)
-        pd_result = (
-            scalars_pandas_df["int64_too"].to_frame().assign(result=pd_result_col)
-        )
+        bf_result = scalars_df["int64_too"].apply(foo, args=args1).to_pandas()
+        pd_result = scalars_pandas_df["int64_too"].apply(foo, args=args1)
 
         # Ignore any dtype difference.
-        pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False)
+        pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False)
 
         args2 = (False, -10.0)
         foo_ref = session.read_gbq_function(foo.bigframes_bigquery_function)
 
-        bf_result_col = scalars_df["int64_too"].apply(foo_ref, args=args2)
-        bf_result = (
-            scalars_df["int64_too"].to_frame().assign(result=bf_result_col).to_pandas()
-        )
-
-        pd_result_col = scalars_pandas_df["int64_too"].apply(foo, args=args2)
-        pd_result = (
-            scalars_pandas_df["int64_too"].to_frame().assign(result=pd_result_col)
-        )
+        bf_result = scalars_df["int64_too"].apply(foo_ref, args=args2).to_pandas()
+        pd_result = scalars_pandas_df["int64_too"].apply(foo, args=args2)
 
         # Ignore any dtype difference.
-        pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False)
+        pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False)
 
     finally:
         # Clean up the gcp assets created for the remote function.