include Index too

rapidsai · Jan 14, 2025 · a99b7ed · a99b7ed
1 parent bf42b63
commit a99b7ed
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/cudf/cudf/pandas/_wrappers/pandas.py b/python/cudf/cudf/pandas/_wrappers/pandas.py
@@ -35,6 +35,7 @@
     _fast_slow_function_call,
     _FastSlowAttribute,
     _FunctionProxy,
+    _maybe_wrap_result,
     _Unusable,
     make_final_proxy_type as _make_final_proxy_type,
     make_intermediate_proxy_type as _make_intermediate_proxy_type,
@@ -262,7 +263,7 @@ def custom_repr_html(obj):
 
 def custom_getitem(self, arg):
     if cudf.api.types.is_scalar(arg):
-        return self._fsproxy_slow[arg]
+        return _maybe_wrap_result(self._fsproxy_slow[arg], None)
     else:
         return _fast_slow_function_call(
             lambda self, arg: self[arg],
@@ -352,6 +353,7 @@ def Index__setattr__(self, name, value):
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
         "name": _FastSlowAttribute("name"),
+        "__getitem__": custom_getitem,
     },
 )
 
@@ -430,6 +432,7 @@ def Index__setattr__(self, name, value):
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
         "name": _FastSlowAttribute("name"),
+        "__getitem__": custom_getitem,
     },
 )
 
@@ -467,6 +470,7 @@ def Index__setattr__(self, name, value):
         "_data": _FastSlowAttribute("_data", private=True),
         "_mask": _FastSlowAttribute("_mask", private=True),
         "name": _FastSlowAttribute("name"),
+        "__getitem__": custom_getitem,
     },
 )