feat: add index get_loc API

Author: shuoweil

bigframes/core/indexes/base.py

@@ -27,16 +27,21 @@
 import pandas
 
 from bigframes import dtypes
+from bigframes.core.array_value import ArrayValue
 import bigframes.core.block_transforms as block_ops
 import bigframes.core.blocks as blocks
 import bigframes.core.expression as ex
+import bigframes.core.identifiers as ids
+import bigframes.core.nodes as nodes
 import bigframes.core.ordering as order
 import bigframes.core.utils as utils
 import bigframes.core.validations as validations
+import bigframes.core.window_spec as window_spec
 import bigframes.dtypes
 import bigframes.formatting_helpers as formatter
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
+import bigframes.series
 
 if typing.TYPE_CHECKING:
     import bigframes.dataframe
@@ -247,6 +252,95 @@ def query_job(self) -> bigquery.QueryJob:
             self._query_job = query_job
         return self._query_job
 
+    def get_loc(self, key):
+        """Get integer location, slice or boolean mask for requested label.
+
+        Args:
+            key: The label to search for in the index.
+
+        Returns:
+            An integer, slice, or boolean mask representing the location(s) of the key.
+
+        Raises:
+            NotImplementedError: If the index has more than one level.
+            KeyError: If the key is not found in the index.
+        """
+
+        if self.nlevels != 1:
+            raise NotImplementedError("get_loc only supports single-level indexes")
+
+        # Get the index column from the block
+        index_column = self._block.index_columns[0]
+
+        # Apply row numbering to the original data
+        win_spec = window_spec.unbound()
+        row_num_agg = ex.NullaryAggregation(agg_ops.RowNumberOp())
+        row_num_col_id = ids.ColumnId.unique()
+
+        window_node = nodes.WindowOpNode(
+            child=self._block._expr.node,
+            expression=row_num_agg,
+            window_spec=win_spec,
+            output_name=row_num_col_id,
+            never_skip_nulls=True,
+        )
+
+        windowed_array = ArrayValue(window_node)
+        windowed_block = self._block.__class__(
+            windowed_array,
+            index_columns=self._block.index_columns,
+            column_labels=self._block.column_labels.insert(
+                len(self._block.column_labels), None
+            ),
+            index_labels=self._block._index_labels,
+        )
+
+        # Create expression to find matching positions
+        match_expr = ops.eq_op.as_expr(ex.deref(index_column), ex.const(key))
+        windowed_block, match_col_id = windowed_block.project_expr(match_expr)
+
+        # Filter to only rows where the key matches
+        filtered_block = windowed_block.filter_by_id(match_col_id)
+
+        # Check if key exists at all by counting on the filtered block
+        count_agg = ex.UnaryAggregation(agg_ops.count_op, ex.deref(row_num_col_id.name))
+        count_result = filtered_block._expr.aggregate([(count_agg, "count")])
+        count_scalar = self._block.session._executor.execute(
+            count_result
+        ).to_py_scalar()
+
+        if count_scalar == 0:
+            raise KeyError(f"'{key}' is not in index")
+
+        # If only one match, return integer position
+        if count_scalar == 1:
+            min_agg = ex.UnaryAggregation(agg_ops.min_op, ex.deref(row_num_col_id.name))
+            position_result = filtered_block._expr.aggregate([(min_agg, "position")])
+            position_scalar = self._block.session._executor.execute(
+                position_result
+            ).to_py_scalar()
+            return int(position_scalar)
+
+        # Multiple matches - need to determine if monotonic or not
+        is_monotonic = self.is_monotonic_increasing or self.is_monotonic_decreasing
+        if is_monotonic:
+            # Get min and max positions for slice
+            min_agg = ex.UnaryAggregation(agg_ops.min_op, ex.deref(row_num_col_id.name))
+            max_agg = ex.UnaryAggregation(agg_ops.max_op, ex.deref(row_num_col_id.name))
+            min_result = filtered_block._expr.aggregate([(min_agg, "min_pos")])
+            max_result = filtered_block._expr.aggregate([(max_agg, "max_pos")])
+            min_pos = self._block.session._executor.execute(min_result).to_py_scalar()
+            max_pos = self._block.session._executor.execute(max_result).to_py_scalar()
+
+            # create slice
+            start = int(min_pos)
+            stop = int(max_pos) + 1  # exclusive
+            return slice(start, stop, None)
+        else:
+            # Return boolean mask for non-monotonic duplicates
+            mask_block = windowed_block.select_columns([match_col_id])
+            return bigframes.series.Series(mask_block)
+
     def __repr__(self) -> str:
         # Protect against errors with uninitialized Series. See:
         # https://github.com/googleapis/python-bigquery-dataframes/issues/728

tests/system/small/test_index.py

@@ -32,6 +32,106 @@ def test_index_construct_from_list():
     pd.testing.assert_index_equal(bf_result, pd_result)
 
 
+@pytest.mark.parametrize("key, expected_loc", [("a", 0), ("b", 1), ("c", 2)])
+def test_get_loc_should_return_int_for_unique_index(key, expected_loc):
+    """Behavior: get_loc on a unique index returns an integer position."""
+    # The pandas result is used as the known-correct value.
+    # We assert our implementation matches it and the expected type.
+    bf_index = bpd.Index(["a", "b", "c"])
+
+    result = bf_index.get_loc(key)
+
+    assert result == expected_loc
+    assert isinstance(result, int)
+
+
+def test_get_loc_should_return_slice_for_monotonic_duplicates():
+    """Behavior: get_loc on a monotonic string index with duplicates returns a slice."""
+    bf_index = bpd.Index(["a", "b", "b", "c"])
+    pd_index = pd.Index(["a", "b", "b", "c"])
+
+    bf_result = bf_index.get_loc("b")
+    pd_result = pd_index.get_loc("b")
+
+    assert isinstance(bf_result, slice)
+    assert bf_result == pd_result  # Should be slice(1, 3, None)
+
+
+def test_get_loc_should_return_slice_for_monotonic_numeric_duplicates():
+    """Behavior: get_loc on a monotonic numeric index with duplicates returns a slice."""
+    bf_index = bpd.Index([1, 2, 2, 3])
+    pd_index = pd.Index([1, 2, 2, 3])
+
+    bf_result = bf_index.get_loc(2)
+    pd_result = pd_index.get_loc(2)
+
+    assert isinstance(bf_result, slice)
+    assert bf_result == pd_result  # Should be slice(1, 3, None)
+
+
+def test_get_loc_should_return_mask_for_non_monotonic_duplicates():
+    """Behavior: get_loc on a non-monotonic string index returns a boolean array."""
+    bf_index = bpd.Index(["a", "b", "c", "b"])
+    pd_index = pd.Index(["a", "b", "c", "b"])
+
+    bf_result = bf_index.get_loc("b")
+    if hasattr(bf_result, "to_numpy"):
+        bf_array = bf_result.to_numpy()
+    else:
+        bf_array = bf_result.to_pandas().to_numpy()
+    pd_result = pd_index.get_loc("b")
+
+    numpy.testing.assert_array_equal(bf_array, pd_result)
+
+
+def test_get_loc_should_return_mask_for_non_monotonic_numeric_duplicates():
+    """Behavior: get_loc on a non-monotonic numeric index returns a boolean array."""
+    bf_index = bpd.Index([1, 2, 3, 2])
+    pd_index = pd.Index([1, 2, 3, 2])
+
+    bf_result = bf_index.get_loc(2)
+    if hasattr(bf_result, "to_numpy"):
+        bf_array = bf_result.to_numpy()
+    else:
+        bf_array = bf_result.to_pandas().to_numpy()
+    pd_result = pd_index.get_loc(2)
+
+    numpy.testing.assert_array_equal(bf_array, pd_result)
+
+
+def test_get_loc_should_raise_error_for_missing_key():
+    """Behavior: get_loc raises KeyError when a string key is not found."""
+    bf_index = bpd.Index(["a", "b", "c"])
+
+    with pytest.raises(KeyError):
+        bf_index.get_loc("d")
+
+
+def test_get_loc_should_raise_error_for_missing_numeric_key():
+    """Behavior: get_loc raises KeyError when a numeric key is not found."""
+    bf_index = bpd.Index([1, 2, 3])
+
+    with pytest.raises(KeyError):
+        bf_index.get_loc(4)
+
+
+def test_get_loc_should_work_for_single_element_index():
+    """Behavior: get_loc on a single-element index returns 0."""
+    assert bpd.Index(["a"]).get_loc("a") == pd.Index(["a"]).get_loc("a")
+
+
+def test_get_loc_should_return_slice_when_all_elements_are_duplicates():
+    """Behavior: get_loc returns a full slice if all elements match the key."""
+    bf_index = bpd.Index(["a", "a", "a"])
+    pd_index = pd.Index(["a", "a", "a"])
+
+    bf_result = bf_index.get_loc("a")
+    pd_result = pd_index.get_loc("a")
+
+    assert isinstance(bf_result, slice)
+    assert bf_result == pd_result  # Should be slice(0, 3, None)
+
+
 def test_index_construct_from_series():
     bf_result = bpd.Index(
         bpd.Series([3, 14, 159], dtype=pd.Float64Dtype(), name="series_name"),

third_party/bigframes_vendored/pandas/core/indexes/base.py

@@ -741,6 +741,41 @@ def argmin(self) -> int:
         """
         raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+    def get_loc(self, key):
+        """
+        Get integer location, slice or boolean mask for requested label.
+
+        **Examples:**
+
+            >>> import bigframes.pandas as bpd
+            >>> bpd.options.display.progress_bar = None
+
+            >>> unique_index = bpd.Index(list('abc'))
+            >>> unique_index.get_loc('b')
+            1
+
+            >>> monotonic_index = bpd.Index(list('abbc'))
+            >>> monotonic_index.get_loc('b')
+            slice(1, 3, None)
+
+            >>> non_monotonic_index = bpd.Index(list('abcb'))
+            >>> non_monotonic_index.get_loc('b')
+            array([False,  True, False,  True])
+
+        Args:
+            key: Label to get the location for.
+
+        Returns:
+            int if unique index, slice if monotonic index with duplicates, else boolean array:
+                Integer position of the label for unique indexes.
+                Slice object for monotonic indexes with duplicates.
+                Boolean array mask for non-monotonic indexes with duplicates.
+
+        Raises:
+            KeyError: If the key is not found in the index.
+        """
+        raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
     def argmax(self) -> int:
         """
         Return int position of the largest value in the Series.