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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
