From 166b33b9f147cf73cdf302c1fdd807716149aa5c Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Fri, 5 Jan 2024 00:10:43 +0000 Subject: [PATCH] feat: support assigning to columns like a property --- bigframes/core/log_adapter.py | 3 +- bigframes/dataframe.py | 41 ++++++++++++++++++++++ bigframes/series.py | 4 +++ tests/system/small/test_dataframe.py | 52 ++++++++++++++++++++++++++-- 4 files changed, 97 insertions(+), 3 deletions(-) diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py index 556851fa1b..860d394cd2 100644 --- a/bigframes/core/log_adapter.py +++ b/bigframes/core/log_adapter.py @@ -19,12 +19,13 @@ _lock = threading.Lock() MAX_LABELS_COUNT = 64 _api_methods: List = [] +_excluded_methods = ["__setattr__", "__getattr__"] def class_logger(decorated_cls): """Decorator that adds logging functionality to each method of the class.""" for attr_name, attr_value in decorated_cls.__dict__.items(): - if callable(attr_value): + if callable(attr_value) and (attr_name not in _excluded_methods): setattr(decorated_cls, attr_name, method_logger(attr_value, decorated_cls)) return decorated_cls diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 595670b0b6..43568383d4 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -251,6 +251,15 @@ def index( ) -> indexes.Index: return indexes.Index(self) + @index.setter + def index(self, value): + # TODO: Handle assigning MultiIndex + result = self._assign_single_item("_new_bf_index", value).set_index( + "_new_bf_index" + ) + self._set_block(result._get_block()) + self.index.name = value.name if hasattr(value, "name") else None + @property def loc(self) -> indexers.LocDataFrameIndexer: return indexers.LocDataFrameIndexer(self) @@ -545,6 +554,29 @@ def __getattr__(self, key: str): else: raise AttributeError(key) + def __setattr__(self, key: str, value): + if key in ["_block", "_query_job"]: + object.__setattr__(self, key, value) + return + # Can this be removed??? + try: + # boring attributes go through boring old path + object.__getattribute__(self, key) + return object.__setattr__(self, key, value) + except AttributeError: + pass + + # if this fails, go on to more involved attribute setting + # (note that this matches __getattr__, above). + try: + if key in self.columns: + self[key] = value + else: + object.__setattr__(self, key, value) + # Can this be removed? + except (AttributeError, TypeError): + object.__setattr__(self, key, value) + def __repr__(self) -> str: """Converts a DataFrame to a string. Calls to_pandas. @@ -1246,6 +1278,15 @@ def _assign_single_item_listlike(self, k: str, v: Sequence) -> DataFrame: [get_column_left[col_id] for col_id in original_index_column_ids], index_labels=self._block.index_labels, ) + src_col = get_column_right[new_column_block.value_columns[0]] + # Check to see if key exists, and modify in place + col_ids = self._block.cols_matching_label(k) + for col_id in col_ids: + result_block = result_block.copy_values( + src_col, get_column_left[col_id] + ) + if len(col_ids) > 0: + result_block = result_block.drop_columns([src_col]) return DataFrame(result_block) def _assign_scalar(self, label: str, value: Union[int, float]) -> DataFrame: diff --git a/bigframes/series.py b/bigframes/series.py index eefd2b755d..1247883aa5 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -139,6 +139,10 @@ def struct(self) -> structs.StructAccessor: def T(self) -> Series: return self.transpose() + @property + def _info_axis(self) -> indexes.Index: + return self.index + def transpose(self) -> Series: return self diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index fa3d5148a8..9557475b46 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3140,17 +3140,65 @@ def test_df___array__(scalars_df_index, scalars_pandas_df_index): ) -def test_getattr_attribute_error_when_pandas_has(scalars_df_index): +def test_df_getattr_attribute_error_when_pandas_has(scalars_df_index): # swapaxes is implemented in pandas but not in bigframes with pytest.raises(AttributeError): scalars_df_index.swapaxes() -def test_getattr_attribute_error(scalars_df_index): +def test_df_getattr_attribute_error(scalars_df_index): with pytest.raises(AttributeError): scalars_df_index.not_a_method() +def test_df_getattr_axes(): + df = dataframe.DataFrame( + [[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"] + ) + assert isinstance(df.index, bigframes.core.indexes.Index) + assert isinstance(df.columns, pandas.Index) + assert isinstance(df.my_column, series.Series) + + +def test_df_setattr_index(): + pd_df = pandas.DataFrame( + [[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"] + ) + bf_df = dataframe.DataFrame(pd_df) + pd_df.index = [4, 5] + bf_df.index = [4, 5] + + assert_pandas_df_equal( + pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False + ) + + +def test_df_setattr_columns(): + pd_df = pandas.DataFrame( + [[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"] + ) + bf_df = dataframe.DataFrame(pd_df) + pd_df.columns = [4, 5, 6] + bf_df.columns = [4, 5, 6] + + assert_pandas_df_equal( + pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False + ) + + +def test_df_setattr_modify_column(): + pd_df = pandas.DataFrame( + [[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"] + ) + bf_df = dataframe.DataFrame(pd_df) + pd_df.my_column = [4, 5] + bf_df.my_column = [4, 5] + + assert_pandas_df_equal( + pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False + ) + + def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index): index_list = scalars_pandas_df_index.string_col.iloc[[0, 1, 1, 5]].values