Merge pull request rapidsai#21221 from rapidsai/main

Forward-merge main into pandas3

Author: AyodeAwe

python/cudf/cudf/core/dataframe.py

@@ -438,7 +438,19 @@ def _setitem_tuple_arg(self, key, value):
 
 
 class _DataFrameAtIndexer(_DataFrameLocIndexer):
-    pass
+    @_performance_tracking
+    def __getitem__(self, key):
+        indexing_utils.validate_scalar_key(
+            key, "Invalid call for scalar access (getting)!"
+        )
+        return super().__getitem__(key)
+
+    @_performance_tracking
+    def __setitem__(self, key, value):
+        indexing_utils.validate_scalar_key(
+            key, "Invalid call for scalar access (getting)!"
+        )
+        return super().__setitem__(key, value)
 
 
 class _DataFrameIlocIndexer(_DataFrameIndexer):
@@ -508,7 +520,19 @@ def _setitem_tuple_arg(self, key, value):
 
 
 class _DataFrameiAtIndexer(_DataFrameIlocIndexer):
-    pass
+    @_performance_tracking
+    def __getitem__(self, key):
+        indexing_utils.validate_scalar_key(
+            key, "iAt based indexing can only have integer indexers"
+        )
+        return super().__getitem__(key)
+
+    @_performance_tracking
+    def __setitem__(self, key, value):
+        indexing_utils.validate_scalar_key(
+            key, "iAt based indexing can only have integer indexers"
+        )
+        return super().__setitem__(key, value)
 
 
 @_performance_tracking

python/cudf/cudf/core/indexing_utils.py

@@ -14,6 +14,7 @@
 from cudf.api.types import (
     _is_scalar_or_zero_d_array,
     is_integer,
+    is_list_like,
 )
 from cudf.core.column.column import as_column
 from cudf.core.copy_types import BooleanMask, GatherMap
@@ -69,6 +70,30 @@ class ScalarIndexer:
 )
 
 
+def validate_scalar_key(key: Any, error_msg: str) -> None:
+    """Validate that key contains only scalar values for .at/.iat indexers.
+
+    Parameters
+    ----------
+    key : Any
+        The key to validate
+    error_msg : str
+        The error message to raise if validation fails
+
+    Raises
+    ------
+    ValueError
+        If the key contains list-like indexers
+    """
+    if not isinstance(key, tuple):
+        if is_list_like(key):
+            raise ValueError(error_msg)
+    else:
+        for k in key:
+            if is_list_like(k):
+                raise ValueError(error_msg)
+
+
 # Helpers for code-sharing between loc and iloc paths
 def expand_key(
     key: Any, frame: DataFrame | Series, method_type: Literal["iloc", "loc"]

python/cudf/cudf/core/series.py

@@ -255,7 +255,19 @@ def __setitem__(self, key, value):
 
 
 class _SeriesiAtIndexer(_SeriesIlocIndexer):
-    pass
+    @_performance_tracking
+    def __getitem__(self, key):
+        indexing_utils.validate_scalar_key(
+            key, "iAt based indexing can only have integer indexers"
+        )
+        return super().__getitem__(key)
+
+    @_performance_tracking
+    def __setitem__(self, key, value):
+        indexing_utils.validate_scalar_key(
+            key, "iAt based indexing can only have integer indexers"
+        )
+        return super().__setitem__(key, value)
 
 
 class _SeriesLocIndexer(_FrameIndexer):
@@ -379,7 +391,19 @@ def _loc_to_iloc(self, arg):
 
 
 class _SeriesAtIndexer(_SeriesLocIndexer):
-    pass
+    @_performance_tracking
+    def __getitem__(self, key):
+        indexing_utils.validate_scalar_key(
+            key, "Invalid call for scalar access (getting)!"
+        )
+        return super().__getitem__(key)
+
+    @_performance_tracking
+    def __setitem__(self, key, value):
+        indexing_utils.validate_scalar_key(
+            key, "Invalid call for scalar access (getting)!"
+        )
+        return super().__setitem__(key, value)
 
 
 class Series(SingleColumnFrame, IndexedFrame):

python/cudf/cudf/pandas/_wrappers/pandas.py

@@ -18,6 +18,7 @@
 # with this module https://github.com/rapidsai/cudf/issues/14521#issue-2015198786
 import pyarrow.dataset as ds  # noqa: F401
 from pandas._testing import at, getitem, iat, iloc, loc, setitem
+from pandas.compat._optional import import_optional_dependency
 from pandas.tseries.holiday import (
     AbstractHolidayCalendar as pd_AbstractHolidayCalendar,
     EasterMonday as pd_EasterMonday,
@@ -314,6 +315,20 @@ def _DataFrame_columns(self):
     return result
 
 
+def _to_xarray(self):
+    # Call xarray conversion functions directly with self (the proxy object).
+    # We must pass the proxy (self), not the slow pandas object, because xarray
+    # does isinstance checks against pd.MultiIndex and pd.api.extensions.ExtensionArray.
+    # After cudf.pandas.install(), these refer to proxy classes. The slow object
+    # contains real pandas types that don't pass isinstance checks against the proxy
+    # classes.
+    xr = import_optional_dependency("xarray")
+    if self.ndim == 1:
+        return xr.DataArray.from_series(self)
+    else:
+        return xr.Dataset.from_dataframe(self)
+
+
 DataFrame = make_final_proxy_type(
     "DataFrame",
     cudf.DataFrame,
@@ -346,6 +361,7 @@ def _DataFrame_columns(self):
         "flags": _FastSlowAttribute("flags", private=True),
         "memory_usage": _FastSlowAttribute("memory_usage"),
         "__sizeof__": _FastSlowAttribute("__sizeof__"),
+        "to_xarray": _to_xarray,
     },
 )
 
@@ -419,6 +435,7 @@ def _argsort(self, *args, **kwargs):
         "_accessors": set(),
         "dtype": property(_Series_dtype),
         "argsort": _argsort,
+        "to_xarray": _to_xarray,
         "attrs": _FastSlowAttribute("attrs"),
         "_mgr": _FastSlowAttribute("_mgr", private=True),
         "array": _FastSlowAttribute("array", private=True),
@@ -537,6 +554,7 @@ def Index__setattr__(self, name, value):
     fast_to_slow=lambda fast: fast,
     slow_to_fast=lambda slow: slow,
     additional_attributes={
+        "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "__from_arrow__": _FastSlowAttribute("__from_arrow__"),
         "__hash__": _FastSlowAttribute("__hash__"),
         "pyarrow_dtype": _FastSlowAttribute("pyarrow_dtype"),
@@ -575,6 +593,10 @@ def Index__setattr__(self, name, value):
     pd.Categorical,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
+    additional_attributes={
+        "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
+    },
 )
 
 CategoricalDtype = make_final_proxy_type(
@@ -614,7 +636,9 @@ def Index__setattr__(self, name, value):
     pd.arrays.DatetimeArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
+        "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
     },
@@ -687,7 +711,9 @@ def Index__setattr__(self, name, value):
     pd.arrays.TimedeltaArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
+        "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
     },
@@ -717,6 +743,7 @@ def Index__setattr__(self, name, value):
     pd.arrays.PeriodArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
@@ -799,11 +826,13 @@ def Index__setattr__(self, name, value):
     pd.arrays.StringArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
         "__array__": _FastSlowAttribute("__array__"),
         "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
+        "__arrow_array__": _FastSlowAttribute("__arrow_array__"),
     },
 )
 
@@ -835,6 +864,7 @@ def Index__setattr__(self, name, value):
     pd.core.arrays.string_arrow.ArrowStringArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "_pa_array": _FastSlowAttribute("_pa_array", private=True),
         "__array__": _FastSlowAttribute("__array__", private=True),
@@ -844,6 +874,7 @@ def Index__setattr__(self, name, value):
         "__abs__": _FastSlowAttribute("__abs__"),
         "__contains__": _FastSlowAttribute("__contains__"),
         "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
+        "__arrow_array__": _FastSlowAttribute("__arrow_array__"),
     },
 )
 
@@ -874,6 +905,7 @@ def Index__setattr__(self, name, value):
     pd.arrays.BooleanArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
@@ -898,6 +930,7 @@ def Index__setattr__(self, name, value):
     pd.arrays.IntegerArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "_data": _FastSlowAttribute("_data", private=True),
@@ -1018,7 +1051,9 @@ def Index__setattr__(self, name, value):
     pd.arrays.IntervalArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
+        "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
     },
@@ -1052,6 +1087,7 @@ def Index__setattr__(self, name, value):
     pd.arrays.FloatingArray,
     fast_to_slow=_Unusable(),
     slow_to_fast=_Unusable(),
+    bases=(pd.api.extensions.ExtensionArray,),
     additional_attributes={
         "__array_ufunc__": _FastSlowAttribute("__array_ufunc__"),
         "_data": _FastSlowAttribute("_data", private=True),
@@ -1081,6 +1117,7 @@ def Index__setattr__(self, name, value):
     },
 )
 
+
 SeriesGroupBy = make_intermediate_proxy_type(
     "SeriesGroupBy",
     cudf.core.groupby.groupby.SeriesGroupBy,