Merge branch 'main' into sycai_ai_gen_bool

Author: sycai

bigframes/core/log_adapter.py

@@ -149,49 +149,61 @@ def wrap(cls):
     return wrap(decorated_cls)
 
 
-def method_logger(method, /, *, custom_base_name: Optional[str] = None):
+def method_logger(method=None, /, *, custom_base_name: Optional[str] = None):
     """Decorator that adds logging functionality to a method."""
 
-    @functools.wraps(method)
-    def wrapper(*args, **kwargs):
-        api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
-        if custom_base_name is None:
-            qualname_parts = getattr(method, "__qualname__", method.__name__).split(".")
-            class_name = qualname_parts[-2] if len(qualname_parts) > 1 else ""
-            base_name = (
-                class_name if class_name else "_".join(method.__module__.split(".")[1:])
-            )
-        else:
-            base_name = custom_base_name
-
-        full_method_name = f"{base_name.lower()}-{api_method_name}"
-        # Track directly called methods
-        if len(_call_stack) == 0:
-            add_api_method(full_method_name)
-
-        _call_stack.append(full_method_name)
-
-        try:
-            return method(*args, **kwargs)
-        except (NotImplementedError, TypeError) as e:
-            # Log method parameters that are implemented in pandas but either missing (TypeError)
-            # or not fully supported (NotImplementedError) in BigFrames.
-            # Logging is currently supported only when we can access the bqclient through
-            # _block.session.bqclient.
-            if len(_call_stack) == 1:
-                submit_pandas_labels(
-                    _get_bq_client(*args, **kwargs),
-                    base_name,
-                    api_method_name,
-                    args,
-                    kwargs,
-                    task=PANDAS_PARAM_TRACKING_TASK,
+    def outer_wrapper(method):
+        @functools.wraps(method)
+        def wrapper(*args, **kwargs):
+            api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__)
+            if custom_base_name is None:
+                qualname_parts = getattr(method, "__qualname__", method.__name__).split(
+                    "."
+                )
+                class_name = qualname_parts[-2] if len(qualname_parts) > 1 else ""
+                base_name = (
+                    class_name
+                    if class_name
+                    else "_".join(method.__module__.split(".")[1:])
                 )
-            raise e
-        finally:
-            _call_stack.pop()
+            else:
+                base_name = custom_base_name
 
-    return wrapper
+            full_method_name = f"{base_name.lower()}-{api_method_name}"
+            # Track directly called methods
+            if len(_call_stack) == 0:
+                add_api_method(full_method_name)
+
+            _call_stack.append(full_method_name)
+
+            try:
+                return method(*args, **kwargs)
+            except (NotImplementedError, TypeError) as e:
+                # Log method parameters that are implemented in pandas but either missing (TypeError)
+                # or not fully supported (NotImplementedError) in BigFrames.
+                # Logging is currently supported only when we can access the bqclient through
+                # _block.session.bqclient.
+                if len(_call_stack) == 1:
+                    submit_pandas_labels(
+                        _get_bq_client(*args, **kwargs),
+                        base_name,
+                        api_method_name,
+                        args,
+                        kwargs,
+                        task=PANDAS_PARAM_TRACKING_TASK,
+                    )
+                raise e
+            finally:
+                _call_stack.pop()
+
+        return wrapper
+
+    if method is None:
+        # Called with parentheses
+        return outer_wrapper
+
+    # Called without parentheses
+    return outer_wrapper(method)
 
 
 def property_logger(prop):

bigframes/core/reshape/tile.py

@@ -20,6 +20,7 @@
 import bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile
 import pandas as pd
 
+import bigframes
 import bigframes.constants
 import bigframes.core.expression as ex
 import bigframes.core.ordering as order
@@ -32,7 +33,7 @@
 
 
 def cut(
-    x: bigframes.series.Series,
+    x,
     bins: typing.Union[
         int,
         pd.IntervalIndex,
@@ -60,9 +61,12 @@ def cut(
             f"but found {type(list(labels)[0])}. {constants.FEEDBACK_LINK}"
         )
 
-    if x.size == 0:
+    if len(x) == 0:
         raise ValueError("Cannot cut empty array.")
 
+    if not isinstance(x, bigframes.series.Series):
+        x = bigframes.series.Series(x)
+
     if isinstance(bins, int):
         if bins <= 0:
             raise ValueError("`bins` should be a positive integer.")

bigframes/dataframe.py

@@ -890,9 +890,11 @@ def __delitem__(self, key: str):
         self._set_block(df._get_block())
 
     def __setitem__(
-        self, key: str | list[str], value: SingleItemValue | MultiItemValue
+        self,
+        key: str | list[str] | pandas.Index,
+        value: SingleItemValue | MultiItemValue,
     ):
-        if isinstance(key, list):
+        if isinstance(key, (list, pandas.Index)):
             df = self._assign_multi_items(key, value)
         else:
             df = self._assign_single_item(key, value)
@@ -2246,7 +2248,7 @@ def _assign_single_item(
 
     def _assign_multi_items(
         self,
-        k: list[str],
+        k: list[str] | pandas.Index,
         v: SingleItemValue | MultiItemValue,
     ) -> DataFrame:
         value_sources: Sequence[Any] = []

tests/system/small/test_dataframe.py

@@ -1197,6 +1197,11 @@ def test_assign_new_column_w_setitem_list_error(scalars_dfs):
         pytest.param(
             ["new_col", "new_col_too"], [1, 2], id="sequence_to_full_new_column"
         ),
+        pytest.param(
+            pd.Index(("new_col", "new_col_too")),
+            [1, 2],
+            id="sequence_to_full_new_column_as_index",
+        ),
     ],
 )
 def test_setitem_multicolumn_with_literals(scalars_dfs, key, value):

tests/system/small/test_pandas.py

@@ -520,6 +520,18 @@ def _convert_pandas_category(pd_s: pd.Series):
     )
 
 
+def test_cut_for_array():
+    """Avoid regressions for internal issue 329866195"""
+    sc = [30, 80, 40, 90, 60, 45, 95, 75, 55, 100, 65, 85]
+    x = [20, 40, 60, 80, 100]
+
+    pd_result: pd.Series = pd.Series(pd.cut(sc, x))
+    bf_result = bpd.cut(sc, x)
+
+    pd_result = _convert_pandas_category(pd_result)
+    pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)
+
+
 @pytest.mark.parametrize(
     ("right", "labels"),
     [

tests/unit/core/test_log_adapter.py

@@ -101,6 +101,17 @@ def test_method_logging_with_custom_base_name(test_method_w_custom_base):
     assert "pandas-method1" in api_methods
 
 
+def test_method_logging_with_custom_base__logger_as_decorator():
+    @log_adapter.method_logger(custom_base_name="pandas")
+    def my_method():
+        pass
+
+    my_method()
+
+    api_methods = log_adapter.get_and_reset_api_methods()
+    assert "pandas-my_method" in api_methods
+
+
 def test_property_logging(test_instance):
     test_instance.my_field
 

tests/unit/test_pandas.py

@@ -122,6 +122,7 @@ def test_method_matches_session(method_name: str):
 )
 def test_cut_raises_with_invalid_labels(bins: int, labels, error_message: str):
     mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True)
+    mock_series.__len__.return_value = 5
     with pytest.raises(ValueError, match=error_message):
         bigframes.pandas.cut(mock_series, bins, labels=labels)
 
@@ -160,6 +161,8 @@ def test_cut_raises_with_unsupported_labels():
 )
 def test_cut_raises_with_invalid_bins(bins: int, error_message: str):
     mock_series = mock.create_autospec(bigframes.pandas.Series, instance=True)
+    mock_series.__len__.return_value = 5
+
     with pytest.raises(ValueError, match=error_message):
         bigframes.pandas.cut(mock_series, bins, labels=False)
 

third_party/bigframes_vendored/pandas/core/reshape/tile.py

@@ -8,11 +8,11 @@
 
 import pandas as pd
 
-from bigframes import constants, series
+from bigframes import constants
 
 
 def cut(
-    x: series.Series,
+    x,
     bins: typing.Union[
         int,
         pd.IntervalIndex,
@@ -113,7 +113,7 @@ def cut(
         dtype: struct<left_inclusive: int64, right_exclusive: int64>[pyarrow]
 
     Args:
-        x (bigframes.pandas.Series):
+        x (array-like):
             The input Series to be binned. Must be 1-dimensional.
         bins (int, pd.IntervalIndex, Iterable):
             The criteria to bin by.