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> 🦕

Author: TrevorBergeron

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
 

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.
 
@@ -1265,6 +1297,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:

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
 

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