Implement GroupBy.__getitem__ (#3691)

* Implement GroupBy.__getitem__

* fix test name

* move whats new entry

Co-authored-by: Mathias Hauser <[email protected]>

Author: dcherian

doc/groupby.rst

@@ -63,6 +63,12 @@ You can also iterate over groups in ``(label, group)`` pairs:
 
     list(ds.groupby("letters"))
 
+You can index out a particular group:
+
+.. ipython:: python
+
+    ds.groupby("letters")["b"]
+
 Just like in pandas, creating a GroupBy object is cheap: it does not actually
 split the data until you access particular values.
 

doc/whats-new.rst

@@ -51,6 +51,10 @@ New Features
   grant from the `Chan Zuckerberg Initiative <https://chanzuckerberg.com>`_ and
   developed by `B-Open <https://www.bopen.eu>`_.
   By `Aureliana Barghini <https://github.com/aurghs>`_ and `Alessandro Amici <https://github.com/alexamici>`_.
+- Implement ``__getitem__`` for both :py:class:`~core.groupby.DatasetGroupBy` and
+  :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas'
+  :py:meth:`~pandas.core.groupby.GroupBy.get_group`.
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 
 
 Breaking changes

xarray/core/groupby.py

@@ -415,11 +415,20 @@ def dims(self):
 
     @property
     def groups(self):
+        """
+        Mapping from group labels to indices. The indices can be used to index the underlying object.
+        """
         # provided to mimic pandas.groupby
         if self._groups is None:
             self._groups = dict(zip(self._unique_coord.values, self._group_indices))
         return self._groups
 
+    def __getitem__(self, key):
+        """
+        Get DataArray or Dataset corresponding to a particular group label.
+        """
+        return self._obj.isel({self._group_dim: self.groups[key]})
+
     def __len__(self):
         return self._unique_coord.size
 

xarray/tests/test_groupby.py

@@ -549,4 +549,17 @@ def test_groupby_none_group_name():
     assert "group" in mean.dims
 
 
+def test_groupby_getitem(dataset):
+
+    assert_identical(dataset.sel(x="a"), dataset.groupby("x")["a"])
+    assert_identical(dataset.sel(z=1), dataset.groupby("z")[1])
+
+    assert_identical(dataset.foo.sel(x="a"), dataset.foo.groupby("x")["a"])
+    assert_identical(dataset.foo.sel(z=1), dataset.foo.groupby("z")[1])
+
+    actual = dataset.groupby("boo")["f"].unstack().transpose("x", "y", "z")
+    expected = dataset.sel(y=[1], z=[1, 2]).transpose("x", "y", "z")
+    assert_identical(expected, actual)
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here