[GitHub] [beam] AnandInguva commented on a diff in pull request #26157: [python] [Dataframe] Support index in df set_index

["AnandInguva (via GitHub)" <gi...@apache.org>]

AnandInguva commented on code in PR #26157:
URL: https://github.com/apache/beam/pull/26157#discussion_r1160732479


##########
sdks/python/apache_beam/dataframe/frames.py:
##########
@@ -2636,18 +2636,13 @@ def get(self, key, default_value=None):
   @frame_base.populate_defaults(pd.DataFrame)
   @frame_base.maybe_inplace
   def set_index(self, keys, **kwargs):
-    """``keys`` must be a ``str`` or ``List[str]``. Passing an Index or Series
-    is not yet supported (`Issue 20759
-    <https://github.com/apache/beam/issues/20759>`_)."""
     if isinstance(keys, str):
       keys = [keys]
 
-    if any(isinstance(k, (_DeferredIndex, frame_base.DeferredFrame))
-           for k in keys):
-      raise NotImplementedError("set_index with Index or Series instances is "
-                                "not yet supported "
-                                "(https://github.com/apache/beam/issues/20759)"
-                                ".")
+    keys = [
+      k._frame._expr.proxy().index if isinstance(k, (_DeferredIndex, frame_base.DeferredFrame)) else k
+      for k in keys

Review Comment:
   Can you add some tests cases please?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org