feat: add AtIndexer getitems (#107)

* feat: add AtIndexer getitems * fix third party docstrings * use loc from at --------- Co-authored-by: Henry J Solberg <henryjsolberg@google.com>

Author: milkshakeiii

bigframes/core/indexers.py

@@ -108,6 +108,16 @@ def __getitem__(self, key: int) -> bigframes.core.scalar.Scalar:
  return self._series.iloc[key]
 
 
+class AtSeriesIndexer:
+ def __init__(self, series: bigframes.series.Series):
+ self._series = series
+
+ def __getitem__(
+ self, key: LocSingleKey
+ ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]:
+ return self._series.loc[key]
+
+
 class LocDataFrameIndexer:
  def __init__(self, dataframe: bigframes.dataframe.DataFrame):
  self._dataframe = dataframe
@@ -221,6 +231,20 @@ def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar:
  return column.iloc[key[0]]
 
 
+class AtDataFrameIndexer:
+ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
+ self._dataframe = dataframe
+
+ def __getitem__(
+ self, key: tuple
+ ) -> Union[bigframes.core.scalar.Scalar, bigframes.series.Series]:
+ if not isinstance(key, tuple):
+ raise TypeError(
+ "DataFrame.at should be indexed by a (row label, column name) tuple."
+ )
+ return self._dataframe.loc[key]
+
+
 @typing.overload
 def _loc_getitem_series_or_dataframe(
  series_or_dataframe: bigframes.series.Series, key

bigframes/dataframe.py

@@ -259,6 +259,10 @@ def iloc(self) -> indexers.ILocDataFrameIndexer:
  def iat(self) -> indexers.IatDataFrameIndexer:
  return indexers.IatDataFrameIndexer(self)
 
+ @property
+ def at(self) -> indexers.AtDataFrameIndexer:
+ return indexers.AtDataFrameIndexer(self)
+
  @property
  def dtypes(self) -> pandas.Series:
  return pandas.Series(data=self._block.dtypes, index=self._block.column_labels)

bigframes/series.py

@@ -91,6 +91,10 @@ def iloc(self) -> bigframes.core.indexers.IlocSeriesIndexer:
  def iat(self) -> bigframes.core.indexers.IatSeriesIndexer:
  return bigframes.core.indexers.IatSeriesIndexer(self)
 
+ @property
+ def at(self) -> bigframes.core.indexers.AtSeriesIndexer:
+ return bigframes.core.indexers.AtSeriesIndexer(self)
+
  @property
  def name(self) -> blocks.Label:
  return self._name

tests/system/small/test_dataframe.py

@@ -2196,6 +2196,29 @@ def test_loc_single_index_no_duplicate(scalars_df_index, scalars_pandas_df_index
  )
 
 
+def test_at_with_duplicate(scalars_df_index, scalars_pandas_df_index):
+ scalars_df_index = scalars_df_index.set_index("string_col", drop=False)
+ scalars_pandas_df_index = scalars_pandas_df_index.set_index(
+ "string_col", drop=False
+ )
+ index = "Hello, World!"
+ bf_result = scalars_df_index.at[index, "int64_too"]
+ pd_result = scalars_pandas_df_index.at[index, "int64_too"]
+ pd.testing.assert_series_equal(
+ bf_result.to_pandas(),
+ pd_result,
+ )
+
+
+def test_at_no_duplicate(scalars_df_index, scalars_pandas_df_index):
+ scalars_df_index = scalars_df_index.set_index("int64_too", drop=False)
+ scalars_pandas_df_index = scalars_pandas_df_index.set_index("int64_too", drop=False)
+ index = -2345
+ bf_result = scalars_df_index.at[index, "string_col"]
+ pd_result = scalars_pandas_df_index.at[index, "string_col"]
+ assert bf_result == pd_result
+
+
 def test_loc_setitem_bool_series_scalar_new_col(scalars_dfs):
  scalars_df, scalars_pandas_df = scalars_dfs
  bf_df = scalars_df.copy()

tests/system/small/test_series.py

@@ -1953,6 +1953,16 @@ def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, ste
  )
 
 
+def test_at(scalars_df_index, scalars_pandas_df_index):
+ scalars_df_index = scalars_df_index.set_index("int64_too", drop=False)
+ scalars_pandas_df_index = scalars_pandas_df_index.set_index("int64_too", drop=False)
+ index = -2345
+ bf_result = scalars_df_index["string_col"].at[index]
+ pd_result = scalars_pandas_df_index["string_col"].at[index]
+
+ assert bf_result == pd_result
+
+
 def test_iat(scalars_df_index, scalars_pandas_df_index):
  bf_result = scalars_df_index["int64_too"].iat[3]
  pd_result = scalars_pandas_df_index["int64_too"].iat[3]

third_party/bigframes_vendored/pandas/core/frame.py

@@ -2112,3 +2112,8 @@ def iloc(self):
  def iat(self):
  """Access a single value for a row/column pair by integer position."""
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+ @property
+ def at(self):
+ """Access a single value for a row/column label pair."""
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

third_party/bigframes_vendored/pandas/core/series.py

@@ -1833,3 +1833,8 @@ def iloc(self):
  def iat(self):
  """Access a single value for a row/column pair by integer position."""
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+ @property
+ def at(self):
+ """Access a single value for a row/column label pair."""
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)