Repository: modin-project/modin

Issue:
FEAT: [Dataframe and Series API] xs function implementation
- dataframe: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.xs.html
- series: https://pandas.pydata.org/docs/reference/api/pandas.Series.xs.html


Discussion:

Resolving patch:
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -25,6 +25,7 @@
     is_dtype_equal,
     is_object_dtype,
     pandas_dtype,
+    is_integer,
 )
 from pandas.core.indexes.api import ensure_index
 import pandas.core.window.rolling
@@ -3673,6 +3674,92 @@ def __getitem__(self, key):
         else:
             return self._getitem(key)
 
+    def xs(
+        self,
+        key,
+        axis=0,
+        level=None,
+        drop_level: bool = True,
+    ):  # noqa: PR01, RT01, D200
+        """
+        Return cross-section from the Series/DataFrame.
+        """
+        axis = self._get_axis_number(axis)
+        labels = self.columns if axis else self.index
+
+        if isinstance(key, list):
+            # deprecated in pandas, to be removed in 2.0
+            warnings.warn(
+                "Passing lists as key for xs is deprecated and will be removed in a "
+                + "future version. Pass key as a tuple instead.",
+                FutureWarning,
+            )
+
+        if level is not None:
+            if not isinstance(labels, pandas.MultiIndex):
+                raise TypeError("Index must be a MultiIndex")
+            loc, new_ax = labels.get_loc_level(key, level=level, drop_level=drop_level)
+
+            # create the tuple of the indexer
+            _indexer = [slice(None)] * self.ndim
+            _indexer[axis] = loc
+            indexer = tuple(_indexer)
+
+            result = self.iloc[indexer]
+            setattr(result, self._pandas_class._get_axis_name(axis), new_ax)
+            return result
+
+        if axis == 1:
+            if drop_level:
+                return self[key]
+            index = self.columns
+        else:
+            index = self.index
+
+        new_index = None
+        if isinstance(index, pandas.MultiIndex):
+            loc, new_index = index._get_loc_level(key, level=0)
+            if not drop_level:
+                if is_integer(loc):
+                    new_index = index[loc : loc + 1]
+                else:
+                    new_index = index[loc]
+        else:
+            loc = index.get_loc(key)
+
+            if isinstance(loc, np.ndarray):
+                if loc.dtype == np.bool_:
+                    (loc,) = loc.nonzero()
+                # Note: pandas uses self._take_with_is_copy here
+                return self.take(loc, axis=axis)
+
+            if not is_scalar(loc):
+                new_index = index[loc]
+
+        if is_scalar(loc) and axis == 0:
+            # In this case loc should be an integer
+            if self.ndim == 1:
+                # if we encounter an array-like and we only have 1 dim
+                # that means that their are list/ndarrays inside the Series!
+                # so just return them (pandas GH 6394)
+                return self.iloc[loc]
+
+            result = self.iloc[loc]
+        elif is_scalar(loc):
+            result = self.iloc[:, slice(loc, loc + 1)]
+        elif axis == 1:
+            result = self.iloc[:, loc]
+        else:
+            result = self.iloc[loc]
+            if new_index is None:
+                raise RuntimeError(
+                    "`new_index` variable shouldn't be equal to None here, something went wrong."
+                )
+            result.index = new_index
+
+        # Note: pandas does result._set_is_copy here
+        return result
+
     __hash__ = None
 
     def _setitem_slice(self, key: slice, value):
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -2567,14 +2567,6 @@ def where(
         )
         return self._create_or_update_from_compiler(query_compiler, inplace)
 
-    def xs(self, key, axis=0, level=None, drop_level=True):  # noqa: PR01, RT01, D200
-        """
-        Return cross-section from the ``DataFrame``.
-        """
-        return self._default_to_pandas(
-            pandas.DataFrame.xs, key, axis=axis, level=level, drop_level=drop_level
-        )
-
     def _getitem_column(self, key):
         """
         Get column specified by `key`.
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -2111,14 +2111,6 @@ def where(
             try_cast=try_cast,
         )
 
-    def xs(
-        self, key, axis=0, level=None, drop_level=True
-    ):  # pragma: no cover # noqa: PR01, D200
-        """
-        Return cross-section from the Series/DataFrame.
-        """
-        raise NotImplementedError("Not Yet implemented.")
-
     @property
     def attrs(self):  # noqa: RT01, D200
         """


Test patch:
diff --git a/modin/pandas/test/dataframe/test_indexing.py b/modin/pandas/test/dataframe/test_indexing.py
--- a/modin/pandas/test/dataframe/test_indexing.py
+++ b/modin/pandas/test/dataframe/test_indexing.py
@@ -1772,17 +1772,36 @@ def test_tail(data, n):
 
 
 def test_xs():
-    d = {
+    # example is based on the doctest in the upstream pandas docstring
+    data = {
         "num_legs": [4, 4, 2, 2],
         "num_wings": [0, 0, 2, 2],
         "class": ["mammal", "mammal", "mammal", "bird"],
         "animal": ["cat", "dog", "bat", "penguin"],
         "locomotion": ["walks", "walks", "flies", "walks"],
     }
-    df = pd.DataFrame(data=d)
-    df = df.set_index(["class", "animal", "locomotion"])
-    with warns_that_defaulting_to_pandas():
-        df.xs("mammal")
+    modin_df, pandas_df = create_test_dfs(data)
+
+    def prepare_dataframes(df):
+        # to make several partitions (only for Modin dataframe)
+        df = (pd if isinstance(df, pd.DataFrame) else pandas).concat([df, df], axis=0)
+        # looks like pandas is sorting the index whereas modin is not, performing a join operation.
+        df = df.reset_index(drop=True)
+        df = df.join(df, rsuffix="_y")
+        return df.set_index(["class", "animal", "locomotion"])
+
+    modin_df = prepare_dataframes(modin_df)
+    pandas_df = prepare_dataframes(pandas_df)
+    eval_general(modin_df, pandas_df, lambda df: df.xs("mammal"))
+    eval_general(modin_df, pandas_df, lambda df: df.xs("cat", level=1))
+    eval_general(modin_df, pandas_df, lambda df: df.xs("num_legs", axis=1))
+    eval_general(
+        modin_df, pandas_df, lambda df: df.xs("cat", level=1, drop_level=False)
+    )
+    eval_general(modin_df, pandas_df, lambda df: df.xs(("mammal", "cat")))
+    eval_general(
+        modin_df, pandas_df, lambda df: df.xs(("mammal", "cat"), drop_level=False)
+    )
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)