Allow dataset row indexing with np.int types (#7423) (#7438)

Author: DavidRConnell

src/datasets/formatting/formatting.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numbers
 import operator
 from collections.abc import Iterable, Mapping, MutableMapping
 from functools import partial
@@ -565,7 +566,7 @@ def _check_valid_index_key(key: Union[int, slice, range, Iterable], size: int) -
 
 
 def key_to_query_type(key: Union[int, slice, range, str, Iterable]) -> str:
-    if isinstance(key, int):
+    if isinstance(key, numbers.Integral):
         return "row"
     elif isinstance(key, str):
         return "column"

tests/test_arrow_dataset.py

@@ -4605,12 +4605,22 @@ async def f(batch):
     assert len(out) == 1
 
 
+def test_dataset_getitem_int_np_equivalence():
+    ds = Dataset.from_dict({"a": [0, 1, 2, 3]})
+
+    assert ds[1] == ds[np.int64(1)]
+
+
 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)
+    with pytest.raises(TypeError):
+        ds[np.bool_(True)]
+    with pytest.raises(TypeError):
+        ds[1.0]
 
 
 def test_categorical_dataset(tmpdir):