feat: Allow iloc to support lists of negative indices

bigframes/core/compile/aggregate_compiler.py

@@ -350,6 +350,12 @@ def _(
     return _apply_window_if_present(column.count(), window)
 
 
+@compile_unary_agg.register
+def _(op: agg_ops.ReverseRowNumberOp, column: ibis_types.NumericColumn, window=None):
+    row_count = _apply_window_if_present(ibis_ops.count(1), window)
+    return column - row_count
+
+
 @compile_unary_agg.register
 def _(
     op: agg_ops.CutOp,

bigframes/core/indexers.py

@@ -27,6 +27,7 @@
 import bigframes.core.guid as guid
 import bigframes.core.indexes as indexes
 import bigframes.core.scalar
+import bigframes.core.window_spec as windows
 import bigframes.dataframe
 import bigframes.dtypes
 import bigframes.exceptions as bfe
@@ -477,6 +478,24 @@ def _iloc_getitem_series_or_dataframe(
                 Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
                 series_or_dataframe.iloc[0:0],
             )
+
+        # Check if both positive index and negative index are necessary
+        if isinstance(key, bigframes.series.Series):
+            # Avoid data download
+            is_key_unisigned = False
+        elif "shape" in series_or_dataframe._block.__dict__:
+            # If there is a cache, we convert all indices to positive.
+            row_count = series_or_dataframe._block.shape[0]
+            key = [k if k >= 0 else row_count + k for k in key]
+            is_key_unisigned = True
+        else:
+            first_sign = key[0] >= 0
+            is_key_unisigned = True
+            for k in key:
+                if (k >= 0) != first_sign:
+                    is_key_unisigned = False
+                    break
+
         if isinstance(series_or_dataframe, bigframes.series.Series):
             original_series_name = series_or_dataframe.name
             series_name = (
@@ -497,7 +516,23 @@ def _iloc_getitem_series_or_dataframe(
         block = df._block
         # explicitly set index to offsets, reset_index may not generate offsets in some modes
         block, offsets_id = block.promote_offsets("temp_iloc_offsets_")
-        block = block.set_index([offsets_id])
+        pos_block = block.set_index([offsets_id])
+
+        if not is_key_unisigned or key[0] < 0:
+            neg_block, _ = block.apply_window_op(
+                offsets_id,
+                ops.aggregations.ReverseRowNumberOp(),
+                window_spec=windows.rows(),
+                result_label="Reversed_Index",
+            )
+
+            neg_block = neg_block.set_index([neg_block.value_columns[-1]])
+
+        if is_key_unisigned:
+            block = pos_block if key[0] >= 0 else neg_block
+        else:
+            block = pos_block.concat([neg_block], how="inner")
+
         df = bigframes.dataframe.DataFrame(block)
 
         result = df.loc[key]

bigframes/operations/aggregations.py

@@ -335,6 +335,14 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         )
 
 
+@dataclasses.dataclass(frozen=True)
+class ReverseRowNumberOp(UnaryWindowOp):
+    name: ClassVar[str] = "size"
+
+    def output_type(self, *input_types):
+        return dtypes.INT_DTYPE
+
+
 @dataclasses.dataclass(frozen=True)
 class CutOp(UnaryWindowOp):
     # TODO: Unintuitive, refactor into multiple ops?

tests/system/small/test_dataframe.py

@@ -4409,7 +4409,7 @@ def test_loc_list_multiindex(scalars_dfs_maybe_ordered):
 
 
 def test_iloc_list(scalars_df_index, scalars_pandas_df_index):
-    index_list = [0, 0, 0, 5, 4, 7]
+    index_list = [0, 0, 0, 5, 4, 7, -2, -5, 3]
 
     bf_result = scalars_df_index.iloc[index_list]
     pd_result = scalars_pandas_df_index.iloc[index_list]
@@ -4423,7 +4423,7 @@ def test_iloc_list(scalars_df_index, scalars_pandas_df_index):
 def test_iloc_list_partial_ordering(
     scalars_df_partial_ordering, scalars_pandas_df_index
 ):
-    index_list = [0, 0, 0, 5, 4, 7]
+    index_list = [0, 0, 0, 5, 4, 7, -2, -5, 3]
 
     bf_result = scalars_df_partial_ordering.iloc[index_list]
     pd_result = scalars_pandas_df_index.iloc[index_list]