feat: add iat and iloc accessing by tuples of integers (#90)

* feat: add iat indexers

* select_columns correction

* add iloc for tuples

* add dataframe tests

* add series tests

* add third party comments

---------

Co-authored-by: Henry J Solberg <[email protected]>

Author: milkshakeiii

bigframes/core/indexers.py

@@ -97,6 +97,16 @@ def __getitem__(
         return _iloc_getitem_series_or_dataframe(self._series, key)
 
 
+class IatSeriesIndexer:
+    def __init__(self, series: bigframes.series.Series):
+        self._series = series
+
+    def __getitem__(self, key: int) -> bigframes.core.scalar.Scalar:
+        if not isinstance(key, int):
+            raise ValueError("Series iAt based indexing can only have integer indexers")
+        return self._series.iloc[key]
+
+
 class LocDataFrameIndexer:
     def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         self._dataframe = dataframe
@@ -188,6 +198,28 @@ def __getitem__(self, key) -> Union[bigframes.dataframe.DataFrame, pd.Series]:
         return _iloc_getitem_series_or_dataframe(self._dataframe, key)
 
 
+class IatDataFrameIndexer:
+    def __init__(self, dataframe: bigframes.dataframe.DataFrame):
+        self._dataframe = dataframe
+
+    def __getitem__(self, key: tuple) -> bigframes.core.scalar.Scalar:
+        error_message = "DataFrame.iat should be indexed by a tuple of exactly 2 ints"
+        # we raise TypeError or ValueError under the same conditions that pandas does
+        if isinstance(key, int):
+            raise TypeError(error_message)
+        if not isinstance(key, tuple):
+            raise ValueError(error_message)
+        key_values_are_ints = [isinstance(key_value, int) for key_value in key]
+        if not all(key_values_are_ints):
+            raise ValueError(error_message)
+        if len(key) != 2:
+            raise TypeError(error_message)
+        block = self._dataframe._block
+        column_block = block.select_columns([block.value_columns[key[1]]])
+        column = bigframes.series.Series(column_block)
+        return column.iloc[key[0]]
+
+
 @typing.overload
 def _loc_getitem_series_or_dataframe(
     series_or_dataframe: bigframes.series.Series, key
@@ -356,6 +388,18 @@ def _iloc_getitem_series_or_dataframe(
         return result_pd_df.iloc[0]
     elif isinstance(key, slice):
         return series_or_dataframe._slice(key.start, key.stop, key.step)
+    elif isinstance(key, tuple) and len(key) == 0:
+        return series_or_dataframe
+    elif isinstance(key, tuple) and len(key) == 1:
+        return _iloc_getitem_series_or_dataframe(series_or_dataframe, key[0])
+    elif (
+        isinstance(key, tuple)
+        and isinstance(series_or_dataframe, bigframes.dataframe.DataFrame)
+        and len(key) == 2
+    ):
+        return series_or_dataframe.iat[key]
+    elif isinstance(key, tuple):
+        raise pd.errors.IndexingError("Too many indexers")
     elif pd.api.types.is_list_like(key):
         if len(key) == 0:
             return typing.cast(

bigframes/dataframe.py

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

bigframes/series.py

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

tests/system/small/test_dataframe.py

@@ -2077,7 +2077,7 @@ def test_iloc_slice_nested(scalars_df_index, scalars_pandas_df_index):
 
 @pytest.mark.parametrize(
     "index",
-    [0, 5, -2],
+    [0, 5, -2, (2,)],
 )
 def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index):
     bf_result = scalars_df_index.iloc[index]
@@ -2089,6 +2089,59 @@ def test_iloc_single_integer(scalars_df_index, scalars_pandas_df_index, index):
     )
 
 
+@pytest.mark.parametrize(
+    "index",
+    [(2, 5), (5, 0), (0, 0)],
+)
+def test_iloc_tuple(scalars_df_index, scalars_pandas_df_index, index):
+    bf_result = scalars_df_index.iloc[index]
+    pd_result = scalars_pandas_df_index.iloc[index]
+
+    assert bf_result == pd_result
+
+
+@pytest.mark.parametrize(
+    ("index", "error"),
+    [
+        ((1, 1, 1), pd.errors.IndexingError),
+        (("asd", "asd", "asd"), pd.errors.IndexingError),
+        (("asd"), TypeError),
+    ],
+)
+def test_iloc_tuple_errors(scalars_df_index, scalars_pandas_df_index, index, error):
+    with pytest.raises(error):
+        scalars_df_index.iloc[index]
+    with pytest.raises(error):
+        scalars_pandas_df_index.iloc[index]
+
+
+@pytest.mark.parametrize(
+    "index",
+    [(2, 5), (5, 0), (0, 0)],
+)
+def test_iat(scalars_df_index, scalars_pandas_df_index, index):
+    bf_result = scalars_df_index.iat[index]
+    pd_result = scalars_pandas_df_index.iat[index]
+
+    assert bf_result == pd_result
+
+
+@pytest.mark.parametrize(
+    ("index", "error"),
+    [
+        (0, TypeError),
+        ("asd", ValueError),
+        ((1, 2, 3), TypeError),
+        (("asd", "asd"), ValueError),
+    ],
+)
+def test_iat_errors(scalars_df_index, scalars_pandas_df_index, index, error):
+    with pytest.raises(error):
+        scalars_pandas_df_index.iat[index]
+    with pytest.raises(error):
+        scalars_df_index.iat[index]
+
+
 def test_iloc_single_integer_out_of_bound_error(
     scalars_df_index, scalars_pandas_df_index
 ):

tests/system/small/test_series.py

@@ -1953,6 +1953,20 @@ def test_series_iloc(scalars_df_index, scalars_pandas_df_index, start, stop, ste
     )
 
 
+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]
+
+    assert bf_result == pd_result
+
+
+def test_iat_error(scalars_df_index, scalars_pandas_df_index):
+    with pytest.raises(ValueError):
+        scalars_pandas_df_index["int64_too"].iat["asd"]
+    with pytest.raises(ValueError):
+        scalars_df_index["int64_too"].iat["asd"]
+
+
 def test_series_add_prefix(scalars_df_index, scalars_pandas_df_index):
     bf_result = scalars_df_index["int64_too"].add_prefix("prefix_").to_pandas()
 

third_party/bigframes_vendored/pandas/core/frame.py

@@ -2097,3 +2097,13 @@ def fillna(self, value):
             DataFrame: Object with missing values filled
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    @property
+    def iloc(self):
+        """Purely integer-location based indexing for selection by position."""
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    @property
+    def iat(self):
+        """Access a single value for a row/column pair by integer position."""
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

third_party/bigframes_vendored/pandas/core/series.py

@@ -1823,3 +1823,13 @@ def map(
             Series: Same index as caller.
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    @property
+    def iloc(self):
+        """Purely integer-location based indexing for selection by position."""
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+    @property
+    def iat(self):
+        """Access a single value for a row/column pair by integer position."""
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)