Skip to content

Commit f645c56

Browse files
feat: support assigning to columns like a property (#304)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent 1c63b45 commit f645c56

File tree

4 files changed

+97
-3
lines changed

4 files changed

+97
-3
lines changed

bigframes/core/log_adapter.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,13 @@
1919
_lock = threading.Lock()
2020
MAX_LABELS_COUNT = 64
2121
_api_methods: List = []
22+
_excluded_methods = ["__setattr__", "__getattr__"]
2223

2324

2425
def class_logger(decorated_cls):
2526
"""Decorator that adds logging functionality to each method of the class."""
2627
for attr_name, attr_value in decorated_cls.__dict__.items():
27-
if callable(attr_value):
28+
if callable(attr_value) and (attr_name not in _excluded_methods):
2829
setattr(decorated_cls, attr_name, method_logger(attr_value, decorated_cls))
2930
return decorated_cls
3031

bigframes/dataframe.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,15 @@ def index(
251251
) -> indexes.Index:
252252
return indexes.Index(self)
253253

254+
@index.setter
255+
def index(self, value):
256+
# TODO: Handle assigning MultiIndex
257+
result = self._assign_single_item("_new_bf_index", value).set_index(
258+
"_new_bf_index"
259+
)
260+
self._set_block(result._get_block())
261+
self.index.name = value.name if hasattr(value, "name") else None
262+
254263
@property
255264
def loc(self) -> indexers.LocDataFrameIndexer:
256265
return indexers.LocDataFrameIndexer(self)
@@ -545,6 +554,29 @@ def __getattr__(self, key: str):
545554
else:
546555
raise AttributeError(key)
547556

557+
def __setattr__(self, key: str, value):
558+
if key in ["_block", "_query_job"]:
559+
object.__setattr__(self, key, value)
560+
return
561+
# Can this be removed???
562+
try:
563+
# boring attributes go through boring old path
564+
object.__getattribute__(self, key)
565+
return object.__setattr__(self, key, value)
566+
except AttributeError:
567+
pass
568+
569+
# if this fails, go on to more involved attribute setting
570+
# (note that this matches __getattr__, above).
571+
try:
572+
if key in self.columns:
573+
self[key] = value
574+
else:
575+
object.__setattr__(self, key, value)
576+
# Can this be removed?
577+
except (AttributeError, TypeError):
578+
object.__setattr__(self, key, value)
579+
548580
def __repr__(self) -> str:
549581
"""Converts a DataFrame to a string. Calls to_pandas.
550582
@@ -1265,6 +1297,15 @@ def _assign_single_item_listlike(self, k: str, v: Sequence) -> DataFrame:
12651297
[get_column_left[col_id] for col_id in original_index_column_ids],
12661298
index_labels=self._block.index_labels,
12671299
)
1300+
src_col = get_column_right[new_column_block.value_columns[0]]
1301+
# Check to see if key exists, and modify in place
1302+
col_ids = self._block.cols_matching_label(k)
1303+
for col_id in col_ids:
1304+
result_block = result_block.copy_values(
1305+
src_col, get_column_left[col_id]
1306+
)
1307+
if len(col_ids) > 0:
1308+
result_block = result_block.drop_columns([src_col])
12681309
return DataFrame(result_block)
12691310

12701311
def _assign_scalar(self, label: str, value: Union[int, float]) -> DataFrame:

bigframes/series.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,10 @@ def struct(self) -> structs.StructAccessor:
139139
def T(self) -> Series:
140140
return self.transpose()
141141

142+
@property
143+
def _info_axis(self) -> indexes.Index:
144+
return self.index
145+
142146
def transpose(self) -> Series:
143147
return self
144148

tests/system/small/test_dataframe.py

Lines changed: 50 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3140,17 +3140,65 @@ def test_df___array__(scalars_df_index, scalars_pandas_df_index):
31403140
)
31413141

31423142

3143-
def test_getattr_attribute_error_when_pandas_has(scalars_df_index):
3143+
def test_df_getattr_attribute_error_when_pandas_has(scalars_df_index):
31443144
# swapaxes is implemented in pandas but not in bigframes
31453145
with pytest.raises(AttributeError):
31463146
scalars_df_index.swapaxes()
31473147

31483148

3149-
def test_getattr_attribute_error(scalars_df_index):
3149+
def test_df_getattr_attribute_error(scalars_df_index):
31503150
with pytest.raises(AttributeError):
31513151
scalars_df_index.not_a_method()
31523152

31533153

3154+
def test_df_getattr_axes():
3155+
df = dataframe.DataFrame(
3156+
[[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"]
3157+
)
3158+
assert isinstance(df.index, bigframes.core.indexes.Index)
3159+
assert isinstance(df.columns, pandas.Index)
3160+
assert isinstance(df.my_column, series.Series)
3161+
3162+
3163+
def test_df_setattr_index():
3164+
pd_df = pandas.DataFrame(
3165+
[[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"]
3166+
)
3167+
bf_df = dataframe.DataFrame(pd_df)
3168+
pd_df.index = [4, 5]
3169+
bf_df.index = [4, 5]
3170+
3171+
assert_pandas_df_equal(
3172+
pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False
3173+
)
3174+
3175+
3176+
def test_df_setattr_columns():
3177+
pd_df = pandas.DataFrame(
3178+
[[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"]
3179+
)
3180+
bf_df = dataframe.DataFrame(pd_df)
3181+
pd_df.columns = [4, 5, 6]
3182+
bf_df.columns = [4, 5, 6]
3183+
3184+
assert_pandas_df_equal(
3185+
pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False
3186+
)
3187+
3188+
3189+
def test_df_setattr_modify_column():
3190+
pd_df = pandas.DataFrame(
3191+
[[1, 1, 1], [1, 1, 1]], columns=["index", "columns", "my_column"]
3192+
)
3193+
bf_df = dataframe.DataFrame(pd_df)
3194+
pd_df.my_column = [4, 5]
3195+
bf_df.my_column = [4, 5]
3196+
3197+
assert_pandas_df_equal(
3198+
pd_df, bf_df.to_pandas(), check_index_type=False, check_dtype=False
3199+
)
3200+
3201+
31543202
def test_loc_list_string_index(scalars_df_index, scalars_pandas_df_index):
31553203
index_list = scalars_pandas_df_index.string_col.iloc[[0, 1, 1, 5]].values
31563204

0 commit comments

Comments
 (0)