Improve errors raised by ds.groupby() of unsupported key type

python/ray/data/dataset.py

@@ -890,6 +890,7 @@ def aggregate(self, *aggs: AggregateFn) -> U:
         This is a blocking operation.
 
         Examples:
+            >>> from ray.data.aggregate import Max, Mean
             >>> ray.data.range(100).aggregate(Max())
             >>> ray.data.range_arrow(100).aggregate(
                 Max("value"), Mean("value"))

python/ray/data/grouped_dataset.py

@@ -40,9 +40,30 @@ def __init__(self, dataset: Dataset[T], key: GroupKeyT):
                 raise NotImplementedError(
                     "Multi-key groupby is not supported yet")
             else:
-                self._key = key[0]
-        else:
+                key = key[0]
+
+        try:
+            fmt = self._dataset._dataset_format()
+        except ValueError:
+            # Dataset is empty/cleared, let downstream ops handle this.
+            fmt = None
+
+        if key is None:
+            self._key = key
+        elif isinstance(key, str):
+            if fmt and fmt == "simple":
+                raise TypeError(
+                    "String key '{}' requires dataset format to be "
+                    "'arrow' or 'pandas', was '{}'.".format(key, fmt))
             self._key = key
+        elif callable(key):
+            if fmt and fmt != "simple":
+                raise NotImplementedError(
+                    "Callable key '{}' requires dataset format to be "
+                    "'simple', was '{}'.".format(key, fmt))
+            self._key = key
+        else:
+            raise TypeError("Invalid key type {} ({}).".format(key, type(key)))
 
     def aggregate(self, *aggs: AggregateFn) -> Dataset[U]:
         """Implements an accumulator-based aggregation.

python/ray/data/tests/test_dataset.py

@@ -3065,6 +3065,20 @@ def test_groupby_arrow(ray_start_regular_shared):
     assert agg_ds.count() == 0
 
 
+def test_groupby_errors(ray_start_regular_shared):
+    ds = ray.data.range(100)
+
+    ds.groupby(None).count().show()  # OK
+    ds.groupby(lambda x: x % 2).count().show()  # OK
+    with pytest.raises(TypeError):
+        ds.groupby("foo").count().show()
+
+    ds = ray.data.range_arrow(100)  # OK
+    ds.groupby(None).count().show()  # OK
+    with pytest.raises(NotImplementedError):
+        ds.groupby(lambda x: x % 2).count().show()
+
+
 @pytest.mark.parametrize("num_parts", [1, 15, 100])
 def test_groupby_agg_name_conflict(ray_start_regular_shared, num_parts):
     # Test aggregation name conflict.