feat(bigframes): support offset-based column access via iloc

packages/bigframes/bigframes/core/indexers.py

@@ -14,20 +14,25 @@
 
 from __future__ import annotations
 
+import numbers
 import typing
 import warnings
-from typing import Tuple, Union
+from typing import Any, Sequence, Tuple, Union, cast
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.ibis.common.exceptions as ibis_exceptions
+import numpy as np
 import pandas as pd
+import pyarrow as pa
+import pyarrow.types
 
 import bigframes.core.blocks
 import bigframes.core.col
 import bigframes.core.expression as ex
 import bigframes.core.guid as guid
 import bigframes.core.indexes as indexes
 import bigframes.core.scalar
+import bigframes.core.validations as validations
 import bigframes.core.window_spec as windows
 import bigframes.dataframe
 import bigframes.dtypes
@@ -45,6 +50,9 @@
     ]
 
 
+_DATAFRAME_ILOC_ERROR = "Only DataFrame.iloc[:, col_indexer] = value is supported."
+
+
 class LocSeriesIndexer:
     def __init__(self, series: bigframes.series.Series):
         self._series = series
@@ -102,6 +110,9 @@ def __getitem__(
 
         Other key types are not yet supported.
         """
+        if not _is_noop_slice(key):
+            validations.enforce_ordered(self._series, "iloc")
+
         return _iloc_getitem_series_or_dataframe(self._series, key)
 
 
@@ -186,14 +197,7 @@ def __setitem__(
         key: Tuple[slice, str],
         value: bigframes.dataframe.SingleItemValue,
     ):
-        if (
-            isinstance(key, tuple)
-            and len(key) == 2
-            and isinstance(key[0], slice)
-            and (key[0].start is None or key[0].start == 0)
-            and (key[0].step is None or key[0].step == 1)
-            and key[0].stop is None
-        ):
+        if isinstance(key, tuple) and len(key) == 2 and _is_noop_slice(key[0]):
             # TODO(swast): Support setting multiple columns with key[1] as a list
             # of labels and value as a DataFrame.
             df = self._dataframe.assign(**{key[1]: value})
@@ -244,8 +248,41 @@ def __getitem__(self, key) -> Union[bigframes.dataframe.DataFrame, pd.Series]:
 
         Other key types are not yet supported.
         """
+        requires_ordering = True
+        if isinstance(key, tuple):
+            if len(key) > 0:
+                row_indexer = key[0]
+                if _is_noop_slice(row_indexer):
+                    requires_ordering = False
+        elif _is_noop_slice(key):
+            requires_ordering = False
+
+        if requires_ordering:
+            validations.enforce_ordered(self._dataframe, "iloc")
+
         return _iloc_getitem_series_or_dataframe(self._dataframe, key)
 
+    def __setitem__(
+        self,
+        key: Tuple[
+            slice, Union[int, typing.Sequence[int], slice, typing.Sequence[bool]]
+        ],
+        value: Union[
+            bigframes.dataframe.SingleItemValue, bigframes.dataframe.DataFrame
+        ],
+    ):
+        if not (isinstance(key, tuple) and len(key) == 2):
+            raise NotImplementedError(_DATAFRAME_ILOC_ERROR)
+
+        row_indexer, col_indexer = key
+
+        if not _is_noop_slice(row_indexer):
+            raise NotImplementedError(_DATAFRAME_ILOC_ERROR)
+
+        col_offsets = _iloc_col_indexer_to_offsets(self._dataframe, col_indexer)
+        df = self._dataframe._assign_multi_items_by_offsets(col_offsets, value)
+        self._dataframe._set_block(df._get_block())
+
 
 class IatDataFrameIndexer:
     def __init__(self, dataframe: bigframes.dataframe.DataFrame):
@@ -283,6 +320,16 @@ def __getitem__(
         return self._dataframe.loc[key]
 
 
+def _is_noop_slice(key: Any) -> bool:
+    """Return True if key is a slice selecting all elements in the original order."""
+    return (
+        isinstance(key, slice)
+        and (key.start is None or key.start == 0)
+        and (key.step is None or key.step == 1)
+        and key.stop is None
+    )
+
+
 @typing.overload
 def _loc_getitem_series_or_dataframe(
     series_or_dataframe: bigframes.series.Series, key
@@ -304,12 +351,14 @@ def _loc_getitem_series_or_dataframe(
     pd.Series,
     bigframes.core.scalar.Scalar,
 ]:
+    if _is_noop_slice(key):
+        return series_or_dataframe.copy()
+
     if isinstance(key, slice):
-        if (key.start is None) and (key.stop is None) and (key.step is None):
-            return series_or_dataframe.copy()
         raise NotImplementedError(
             f"loc does not yet support indexing with a slice. {constants.FEEDBACK_LINK}"
         )
+
     if isinstance(key, bigframes.core.col.Expression):
         label_to_col_ref = {
             label: ex.deref(id)
@@ -426,6 +475,74 @@ def _struct_accessor_check_and_warn(
         warnings.warn(msg, stacklevel=7, category=bfe.BadIndexerKeyWarning)
 
 
+def _iloc_clip_to_offset(index: int, length: int, name: str) -> int:
+    """Support negative values for offsets."""
+    offset = index
+    if offset < 0:
+        offset += length
+
+    if offset < 0 or offset >= length:
+        raise IndexError(f"{name} {index} is out-of-bounds")
+
+    return index
+
+
+def _is_integer_scalar(value: Any) -> bool:
+    return not (
+        isinstance(value, bool)
+        or isinstance(value, np.bool_)
+        or (isinstance(value, pa.Scalar) and pyarrow.types.is_boolean(value))
+    ) and (
+        isinstance(value, numbers.Integral)
+        or (isinstance(value, pa.Scalar) and pyarrow.types.is_integer(value))
+    )
+
+
+def _iloc_col_indexer_to_offsets(
+    df: bigframes.dataframe.DataFrame, col_indexer: Any
+) -> Sequence[int]:
+    """Convert col_indexer from one of the many pandas-compatible formats to a list of offsets."""
+    n_cols = len(df.columns)
+
+    if _is_integer_scalar(col_indexer):
+        col_offset = int(col_indexer)
+        return [
+            _iloc_clip_to_offset(
+                col_offset, n_cols, "single positional iloc column indexer"
+            )
+        ]
+
+    elif isinstance(col_indexer, slice):
+        return list(range(*col_indexer.indices(n_cols)))
+
+    elif pd.api.types.is_list_like(col_indexer):
+        col_indexer_list = list(col_indexer)
+
+        if len(col_indexer_list) > 0 and all(
+            isinstance(x, bool) for x in col_indexer_list
+        ):
+            if len(col_indexer_list) != n_cols:
+                raise ValueError(
+                    f"Boolean iloc column indexer has wrong length: {len(col_indexer_list)} instead of {n_cols}"
+                )
+            return [i for i, val in enumerate(col_indexer_list) if val]
+        else:
+            return [
+                _iloc_clip_to_offset(idx, n_cols, "iloc column indexer")
+                for idx in col_indexer_list
+            ]
+
+    raise TypeError(f"got unexpected {type(col_indexer)} for iloc column indexer")
+
+
+def _iloc_df_from_column_offsets(
+    df: bigframes.dataframe.DataFrame, key: Sequence[int]
+) -> bigframes.dataframe.DataFrame:
+    block = df._block
+    selected_ids = tuple(block.value_columns[offset] for offset in key)
+    return bigframes.dataframe.DataFrame(block.select_columns(selected_ids))
+
+
 @typing.overload
 def _iloc_getitem_series_or_dataframe(
     series_or_dataframe: bigframes.series.Series, key
@@ -470,14 +587,22 @@ def _iloc_getitem_series_or_dataframe(
 
         # len(key) == 2
         df = typing.cast(bigframes.dataframe.DataFrame, series_or_dataframe)
-        if isinstance(key[1], int):
+        if isinstance(key[0], int) and isinstance(key[1], int):
             return df.iat[key]
-        elif isinstance(key[1], list):
-            columns = df.columns[key[1]]
-            return _iloc_getitem_series_or_dataframe(df[columns], key[0])
-        raise NotImplementedError(
-            f"iloc does not yet support indexing with {key}. {constants.FEEDBACK_LINK}"
-        )
+
+        row_indexer, column_indexer = key
+        column_offsets = _iloc_col_indexer_to_offsets(df, column_indexer)
+        df_subset = _iloc_df_from_column_offsets(df, column_offsets)
+
+        if _is_integer_scalar(column_indexer):
+            selected_columns = cast(
+                Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
+                df_subset[df_subset.columns[0]],
+            )
+        else:
+            selected_columns = df_subset
+
+        return _iloc_getitem_series_or_dataframe(selected_columns, row_indexer)
     elif pd.api.types.is_list_like(key):
         if len(key) == 0:
             return typing.cast(

packages/bigframes/bigframes/core/local_data.py

@@ -22,7 +22,17 @@
 import itertools
 import json
 import uuid
-from typing import Any, Callable, Generator, Iterable, Literal, Optional, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    Iterable,
+    Literal,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+)
 
 import geopandas  # type: ignore
 import numpy
@@ -521,3 +531,77 @@ def _schema_durations_to_ints(schema: pa.Schema) -> pa.Schema:
     return pa.schema(
         pa.field(field.name, _durations_to_ints(field.type)) for field in schema
     )
+
+
+def _has_extension_type(pa_type: pa.DataType) -> bool:
+    """Check if the given pyarrow DataType contains an ExtensionType at any nesting level."""
+    if isinstance(pa_type, pa.ExtensionType):
+        return True
+    if pa.types.is_struct(pa_type):
+        struct_type = cast(pa.StructType, pa_type)
+        return any(
+            _has_extension_type(struct_type.field(i).type)
+            for i in range(struct_type.num_fields)
+        )
+    if pa.types.is_list(pa_type) or pa.types.is_large_list(pa_type):
+        list_type = cast(pa.ListType, pa_type)
+        return _has_extension_type(list_type.value_type)
+    return False
+
+
+def pyarrow_array_from_sequence(data: Any, pa_type: pa.DataType) -> pa.Array:
+    """Construct a pyarrow Array from a Sequence, DataFrame, or Series,
+    recursively handling ExtensionTypes inside structs and lists.
+
+    This is a workaround for issue b/456774399 where ArrowNotImplementedError is
+    thrown when constructing arrays of structs containing extension dtypes like
+    JSON.
+    """
+    if isinstance(data, (pa.Array, pa.ChunkedArray)):
+        return data.cast(pa_type) if data.type != pa_type else data
+    if isinstance(pa_type, pa.ExtensionType):
+        storage_arr = pyarrow_array_from_sequence(data, pa_type.storage_type)
+        return pa.ExtensionArray.from_storage(pa_type, storage_arr)
+    if pa.types.is_struct(pa_type):
+        struct_type = cast(pa.StructType, pa_type)
+        if isinstance(data, (pd.Series, pd.Index)):
+            df = pd.DataFrame(data.tolist())
+        elif isinstance(data, pd.DataFrame):
+            df = data
+        else:
+            df = pd.DataFrame(
+                list(data)
+                if not isinstance(data, (list, tuple, np.ndarray, pd.Series, pd.Index))
+                else data
+            )
+        arrays = []
+        fields = []
+        for i in range(struct_type.num_fields):
+            field = struct_type.field(i)
+            col_data = df[field.name] if field.name in df.columns else [None] * len(df)
+            arrays.append(pyarrow_array_from_sequence(col_data, field.type))
+            fields.append(field)
+        return pa.StructArray.from_arrays(arrays, fields=fields)
+    if pa.types.is_list(pa_type) or pa.types.is_large_list(pa_type):
+        list_type = cast(pa.ListType, pa_type)
+        if _has_extension_type(list_type.value_type):
+            flat_values: list[Any] = []
+            offsets = [0]
+            iterable = (
+                data.tolist() if isinstance(data, (pd.Series, pd.Index)) else data
+            )
+            for item in iterable:
+                if item is None or (
+                    not isinstance(
+                        item, (list, tuple, np.ndarray, pd.Series, pd.Index, pa.Array)
+                    )
+                    and pd.isna(item)
+                ):
+                    offsets.append(offsets[-1])
+                else:
+                    item_seq = cast(Sequence[Any], item)
+                    flat_values.extend(item_seq)
+                    offsets.append(offsets[-1] + len(item_seq))
+            values_arr = pyarrow_array_from_sequence(flat_values, list_type.value_type)
+            return pa.ListArray.from_arrays(pa.array(offsets), values_arr)
+    return pa.array(data, type=pa_type)