Repository: pydata/xarray

Issue:
Fail to sel() when index comes from categorical pandas Series
Dear xarray team,

Thank you very much for your work on this useful package.
Here is a bug I just found in my code.

#### MCVE Code Sample

Creating a Dataset from pandas when the coordinate is a categorical series:
```python
import pandas as pd

ind = pd.Series(['foo', 'bar'], dtype='category')
df = pd.DataFrame({'ind': ind, 'values': [1, 2]})
df = df.set_index('ind')

ds = df.to_xarray()

print(ds.sel(ind='foo'))
```

#### Expected Output

When `ind` is not categorical, the code returns the expected output:
```
<xarray.Dataset>
Dimensions:  ()
Coordinates:
    ind      <U3 'foo'
Data variables:
    values   int64 1
```

#### Problem Description

When `ind` is categorical, it fails and gives the following traceback:
```
Traceback (most recent call last):
  File "/home/matthieu/foo.py", line 9, in <module>
    print(ds.sel(ind='foo'))
  File "/opt/anaconda3/envs/issue/lib/python3.8/site-packages/xarray/core/dataset.py", line 2013,
in sel
    pos_indexers, new_indexes = remap_label_indexers(
  File "/opt/anaconda3/envs/issue/lib/python3.8/site-packages/xarray/core/coordinates.py", line 39
1, in remap_label_indexers
    pos_indexers, new_indexes = indexing.remap_label_indexers(
  File "/opt/anaconda3/envs/issue/lib/python3.8/site-packages/xarray/core/indexing.py", line 260,
in remap_label_indexers
    idxr, new_idx = convert_label_indexer(index, label, dim, method, tolerance)
  File "/opt/anaconda3/envs/issue/lib/python3.8/site-packages/xarray/core/indexing.py", line 179,
in convert_label_indexer
    indexer = index.get_loc(
TypeError: get_loc() got an unexpected keyword argument 'tolerance'
```


#### Output of ``xr.show_versions()``
<details>
commit: None
python: 3.8.0 (default, Nov  6 2019, 21:49:08)
[GCC 7.3.0]
python-bits: 64
OS: Linux
OS-release: 4.19.91-1-MANJARO
machine: x86_64
processor:
byteorder: little
LC_ALL: None
LANG: fr_FR.UTF-8
LOCALE: fr_FR.UTF-8
libhdf5: None
libnetcdf: None

xarray: 0.14.1
pandas: 0.25.3
numpy: 1.17.4
scipy: None
netCDF4: None
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: None
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: None
distributed: None
matplotlib: None
cartopy: None
seaborn: None
numbagg: None
setuptools: 44.0.0.post20200106
pip: 19.3.1
conda: None
pytest: None
IPython: None
sphinx: None
</details>



Discussion:
fujiisoup: I don't think the check failure is related.
pep8speaks: Hello @fujiisoup! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:







* In the file [`xarray/tests/test_dataset.py`](https://github.com/pydata/xarray/blob/98d38b09f45ba2a5df6b6b19705eaf478b5a96ef/xarray/tests/test_dataset.py):

> [Line 1409:53](https://github.com/pydata/xarray/blob/98d38b09f45ba2a5df6b6b19705eaf478b5a96ef/xarray/tests/test_dataset.py#L1409): [W291](https://duckduckgo.com/?q=pep8%20W291) trailing whitespace


fujiisoup: I think this PR is ready for review.
fujiisoup: I'll merge this tomorrow if no more commens.
keewis: unless I'm missing something this should be really close (run black once and fix the merge conflict in `whats-new.rst`)
keewis: should be ready for merging now
fujiisoup: Thanks, @dcherian and @keewis , for keeping this updated. 
Merging. 

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 9fbeaa95..0e48dfcf 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -40,6 +40,8 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
+- :py:meth:`DataArray.sel` and :py:meth:`Dataset.sel` now support :py:class:`pandas.CategoricalIndex`. (:issue:`3669`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Support using an existing, opened h5netcdf ``File`` with
   :py:class:`~xarray.backends.H5NetCDFStore`. This permits creating an
   :py:class:`~xarray.Dataset` from a h5netcdf ``File`` that has been opened
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 5ac79999..6f1ea76b 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -64,6 +64,7 @@ from .indexes import (
     default_indexes,
     isel_variable_and_index,
     propagate_indexes,
+    remove_unused_levels_categories,
     roll_index,
 )
 from .indexing import is_fancy_indexer
@@ -3411,7 +3412,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
 
     def _unstack_once(self, dim: Hashable, fill_value, sparse) -> "Dataset":
         index = self.get_index(dim)
-        index = index.remove_unused_levels()
+        index = remove_unused_levels_categories(index)
         full_idx = pd.MultiIndex.from_product(index.levels, names=index.names)
 
         # take a shortcut in case the MultiIndex was not modified.
@@ -4460,7 +4461,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         return self._to_dataframe(self.dims)
 
     def _set_sparse_data_from_dataframe(
-        self, dataframe: pd.DataFrame, dims: tuple, shape: Tuple[int, ...]
+        self, dataframe: pd.DataFrame, dims: tuple
     ) -> None:
         from sparse import COO
 
@@ -4468,9 +4469,11 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         if isinstance(idx, pd.MultiIndex):
             coords = np.stack([np.asarray(code) for code in idx.codes], axis=0)
             is_sorted = idx.is_lexsorted
+            shape = tuple(lev.size for lev in idx.levels)
         else:
             coords = np.arange(idx.size).reshape(1, -1)
             is_sorted = True
+            shape = (idx.size,)
 
         for name, series in dataframe.items():
             # Cast to a NumPy array first, in case the Series is a pandas
@@ -4495,14 +4498,16 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
             self[name] = (dims, data)
 
     def _set_numpy_data_from_dataframe(
-        self, dataframe: pd.DataFrame, dims: tuple, shape: Tuple[int, ...]
+        self, dataframe: pd.DataFrame, dims: tuple
     ) -> None:
         idx = dataframe.index
         if isinstance(idx, pd.MultiIndex):
             # expand the DataFrame to include the product of all levels
             full_idx = pd.MultiIndex.from_product(idx.levels, names=idx.names)
             dataframe = dataframe.reindex(full_idx)
-
+            shape = tuple(lev.size for lev in idx.levels)
+        else:
+            shape = (idx.size,)
         for name, series in dataframe.items():
             data = np.asarray(series).reshape(shape)
             self[name] = (dims, data)
@@ -4543,7 +4548,8 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         if not dataframe.columns.is_unique:
             raise ValueError("cannot convert DataFrame with non-unique columns")
 
-        idx = dataframe.index
+        idx = remove_unused_levels_categories(dataframe.index)
+        dataframe = dataframe.set_index(idx)
         obj = cls()
 
         if isinstance(idx, pd.MultiIndex):
@@ -4553,17 +4559,15 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
             )
             for dim, lev in zip(dims, idx.levels):
                 obj[dim] = (dim, lev)
-            shape = tuple(lev.size for lev in idx.levels)
         else:
             index_name = idx.name if idx.name is not None else "index"
             dims = (index_name,)
             obj[index_name] = (dims, idx)
-            shape = (idx.size,)
 
         if sparse:
-            obj._set_sparse_data_from_dataframe(dataframe, dims, shape)
+            obj._set_sparse_data_from_dataframe(dataframe, dims)
         else:
-            obj._set_numpy_data_from_dataframe(dataframe, dims, shape)
+            obj._set_numpy_data_from_dataframe(dataframe, dims)
         return obj
 
     def to_dask_dataframe(self, dim_order=None, set_index=False):
diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py
index 8337a0f0..06bf08ce 100644
--- a/xarray/core/indexes.py
+++ b/xarray/core/indexes.py
@@ -9,6 +9,26 @@ from .utils import is_scalar
 from .variable import Variable
 
 
+def remove_unused_levels_categories(index):
+    """
+    Remove unused levels from MultiIndex and unused categories from CategoricalIndex
+    """
+    if isinstance(index, pd.MultiIndex):
+        index = index.remove_unused_levels()
+        # if it contains CategoricalIndex, we need to remove unused categories
+        # manually. See https://github.com/pandas-dev/pandas/issues/30846
+        if any(isinstance(lev, pd.CategoricalIndex) for lev in index.levels):
+            levels = []
+            for i, level in enumerate(index.levels):
+                if isinstance(level, pd.CategoricalIndex):
+                    level = level[index.codes[i]].remove_unused_categories()
+                levels.append(level)
+            index = pd.MultiIndex.from_arrays(levels, names=index.names)
+    elif isinstance(index, pd.CategoricalIndex):
+        index = index.remove_unused_categories()
+    return index
+
+
 class Indexes(collections.abc.Mapping):
     """Immutable proxy for Dataset or DataArrary indexes."""
 
diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py
index 8e851b39..4e58be1a 100644
--- a/xarray/core/indexing.py
+++ b/xarray/core/indexing.py
@@ -175,6 +175,16 @@ def convert_label_indexer(index, label, index_name="", method=None, tolerance=No
         if label.ndim == 0:
             if isinstance(index, pd.MultiIndex):
                 indexer, new_index = index.get_loc_level(label.item(), level=0)
+            elif isinstance(index, pd.CategoricalIndex):
+                if method is not None:
+                    raise ValueError(
+                        "'method' is not a valid kwarg when indexing using a CategoricalIndex."
+                    )
+                if tolerance is not None:
+                    raise ValueError(
+                        "'tolerance' is not a valid kwarg when indexing using a CategoricalIndex."
+                    )
+                indexer = index.get_loc(label.item())
             else:
                 indexer = index.get_loc(
                     label.item(), method=method, tolerance=tolerance


Test patch:
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index f9eb37db..4e51e229 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -1408,6 +1408,56 @@ class TestDataset:
                 )
             )
 
+    def test_sel_categorical(self):
+        ind = pd.Series(["foo", "bar"], dtype="category")
+        df = pd.DataFrame({"ind": ind, "values": [1, 2]})
+        ds = df.set_index("ind").to_xarray()
+        actual = ds.sel(ind="bar")
+        expected = ds.isel(ind=1)
+        assert_identical(expected, actual)
+
+    def test_sel_categorical_error(self):
+        ind = pd.Series(["foo", "bar"], dtype="category")
+        df = pd.DataFrame({"ind": ind, "values": [1, 2]})
+        ds = df.set_index("ind").to_xarray()
+        with pytest.raises(ValueError):
+            ds.sel(ind="bar", method="nearest")
+        with pytest.raises(ValueError):
+            ds.sel(ind="bar", tolerance="nearest")
+
+    def test_categorical_index(self):
+        cat = pd.CategoricalIndex(
+            ["foo", "bar", "foo"],
+            categories=["foo", "bar", "baz", "qux", "quux", "corge"],
+        )
+        ds = xr.Dataset(
+            {"var": ("cat", np.arange(3))},
+            coords={"cat": ("cat", cat), "c": ("cat", [0, 1, 1])},
+        )
+        # test slice
+        actual = ds.sel(cat="foo")
+        expected = ds.isel(cat=[0, 2])
+        assert_identical(expected, actual)
+        # make sure the conversion to the array works
+        actual = ds.sel(cat="foo")["cat"].values
+        assert (actual == np.array(["foo", "foo"])).all()
+
+        ds = ds.set_index(index=["cat", "c"])
+        actual = ds.unstack("index")
+        assert actual["var"].shape == (2, 2)
+
+    def test_categorical_reindex(self):
+        cat = pd.CategoricalIndex(
+            ["foo", "bar", "baz"],
+            categories=["foo", "bar", "baz", "qux", "quux", "corge"],
+        )
+        ds = xr.Dataset(
+            {"var": ("cat", np.arange(3))},
+            coords={"cat": ("cat", cat), "c": ("cat", [0, 1, 2])},
+        )
+        actual = ds.reindex(cat=["foo"])["cat"].values
+        assert (actual == np.array(["foo"])).all()
+
     def test_sel_drop(self):
         data = Dataset({"foo": ("x", [1, 2, 3])}, {"x": [0, 1, 2]})
         expected = Dataset({"foo": 1})
@@ -3865,6 +3915,21 @@ class TestDataset:
         expected = pd.DataFrame([[]], index=idx)
         assert expected.equals(actual), (expected, actual)
 
+    def test_from_dataframe_categorical(self):
+        cat = pd.CategoricalDtype(
+            categories=["foo", "bar", "baz", "qux", "quux", "corge"]
+        )
+        i1 = pd.Series(["foo", "bar", "foo"], dtype=cat)
+        i2 = pd.Series(["bar", "bar", "baz"], dtype=cat)
+
+        df = pd.DataFrame({"i1": i1, "i2": i2, "values": [1, 2, 3]})
+        ds = df.set_index("i1").to_xarray()
+        assert len(ds["i1"]) == 3
+
+        ds = df.set_index(["i1", "i2"]).to_xarray()
+        assert len(ds["i1"]) == 2
+        assert len(ds["i2"]) == 2
+
     @requires_sparse
     def test_from_dataframe_sparse(self):
         import sparse