Repository: huggingface/datasets
Pull request: Raise TypeError when indexing a dataset with bool
Fix #5858.

Problem statement:
Throw an error when dataset improperly indexed
### Describe the bug

Pandas-style subset indexing on dataset does not throw an error, when maybe it should. Instead returns the first instance of the dataset regardless of index condition.

### Steps to reproduce the bug

Steps to reproduce the behavior:

1. `squad = datasets.load_dataset("squad_v2", split="validation")`
2. `item = squad[squad['question'] ==  "Who was the Norse leader?"]`
or `it = squad[squad['id'] == '56ddde6b9a695914005b962b']`
3.  returns the first item in the dataset, which does not satisfy the above conditions:

`{'id': '56ddde6b9a695914005b9628', 'title': 'Normans', 'context': 'The Normans (Norman: Nourmands; French: Normands; Latin: Normanni) were the people who in the 10th and 11th centuries gave their name to Normandy, a region in France. They were descended from Norse ("Norman" comes from "Norseman") raiders and pirates from Denmark, Iceland and Norway who, under their leader Rollo, agreed to swear fealty to King Charles III of West Francia. Through generations of assimilation and mixing with the native Frankish and Roman-Gaulish populations, their descendants would gradually merge with the Carolingian-based cultures of West Francia. The distinct cultural and ethnic identity of the Normans emerged initially in the first half of the 10th century, and it continued to evolve over the succeeding centuries.', 'question': 'In what country is Normandy located?', 'answers': {'text': ['France', 'France', 'France', 'France'], 'answer_start': [159, 159, 159, 159]}}`

### Expected behavior

Should either throw an error message, or return the dataset item that satisfies the condition.

### Environment info

- `datasets` version: 2.9.0
- Platform: macOS-13.3.1-arm64-arm-64bit
- Python version: 3.10.8
- PyArrow version: 10.0.1
- Pandas version: 1.5.3

Hints from issue discussion:
Thanks for reporting, @sarahwie.

Please note that in `datasets` we do not have vectorized operation like `pandas`. Therefore, your equality comparisons above are `False`:
- For example: `squad['question']` returns a `list`, and this list is not equal to `"Who was the Norse leader?"`

The `False` value is equivalent to `0` when indexing a dataset, thus the reason why you get the first element (with index 0): 
- For example: `squad[False]` is equivalent to `squad[0]`

Maybe we should an exception instead of assuming that `False` is equivalent to `0` (and `True` is equivalent to `1`) in the context of indexing.

Commit message: Remove sequence of bool from signature

Patch to review:
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index 93637e2d839..9b1d86310cb 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -41,6 +41,7 @@
     Any,
     BinaryIO,
     Callable,
+    Collection,
     Dict,
     Iterable,
     Iterator,
@@ -2742,10 +2743,12 @@ def prepare_for_task(self, task: Union[str, TaskTemplate], id: int = 0) -> "Data
         dataset = dataset.cast(features=template.features)
         return dataset
 
-    def _getitem(self, key: Union[int, slice, str], **kwargs) -> Union[Dict, List]:
+    def _getitem(self, key: Union[int, slice, str, Collection[int]], **kwargs) -> Union[Dict, List]:
         """
-        Can be used to index columns (by string names) or rows (by integer index, slices, or iter of indices or bools)
+        Can be used to index columns (by string names) or rows (by integer index, slices, or collection of integer indices)
         """
+        if isinstance(key, bool):
+            raise TypeError("dataset index must be int, str, slice or collection of int, not bool")
         format_type = kwargs["format_type"] if "format_type" in kwargs else self._format_type
         format_columns = kwargs["format_columns"] if "format_columns" in kwargs else self._format_columns
         output_all_columns = (
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
index 06a75717b3d..68bbef8f745 100644
--- a/tests/test_arrow_dataset.py
+++ b/tests/test_arrow_dataset.py
@@ -4520,3 +4520,11 @@ def f(x):
     ds = ds.map(f)
     outputs = ds[:]
     assert outputs == {"a": [{"nested": [[i]]} for i in [-1, -1, 2, 3]]}
+
+
+def test_dataset_getitem_raises():
+    ds = Dataset.from_dict({"a": [0, 1, 2, 3]})
+    with pytest.raises(TypeError):
+        ds[False]
+    with pytest.raises(TypeError):
+        ds._getitem(True)


Code review comments:
diff_hunk: @@ -2742,10 +2743,12 @@ def prepare_for_task(self, task: Union[str, TaskTemplate], id: int = 0) -> "Data
         dataset = dataset.cast(features=template.features)
         return dataset
 
-    def _getitem(self, key: Union[int, slice, str], **kwargs) -> Union[Dict, List]:
+    def _getitem(self, key: Union[int, slice, str, Collection[int]], **kwargs) -> Union[Dict, List]:
path: src/datasets/arrow_dataset.py
text: @user1:
I think Sequence makes more sense, we shouldn't let users think they can should use a dictionary here
```suggestion
    def _getitem(self, key: Union[int, slice, str, Sequence_[int]], **kwargs) -> Union[Dict, List]:
```

@author:
@user1, I agree with you. Note that I added `Sequence` in a previous commit: https://github.com/huggingface/datasets/pull/5859/commits/a4c43e8b33e2d25e2ee702bd8b358fc6e568ad55

However, then I realized that `Sequence` is too restrictive: we also support `np.ndarray` (and also `pd.Series`) and these are not `Sequence` (but they are `Collection`). See: https://github.com/huggingface/datasets/blob/22d1d533e8ab831b1aa1aab3e7d3c72ba42a83e8/tests/test_arrow_dataset.py#L212

That is why I then changed it to `Collection`. 
But I agree that `Collection` is too permissive (we do not support `dict`).

What do you think could be a better solution? Maybe, instead of `Collection[int]`:
```python
ListLike[int]
```
See: https://github.com/pandas-dev/pandas/blob/2a270d8b04d9573b17b9d550a049d003a974a9d4/pandas/_typing.py#LL114C1-L114C9


@user1:
ListLike sounds good !