GH-51145: [Python] Accept Arrow integer scalars in slice methods

python/pyarrow/array.pxi

@@ -18,6 +18,7 @@
 from cpython.pycapsule cimport PyCapsule_CheckExact, PyCapsule_GetPointer, PyCapsule_New
 
 from collections.abc import Sequence
+import operator
 import os
 import warnings
 from cython import sizeof
@@ -1631,13 +1632,15 @@ cdef class Array(_PandasConvertible):
         """
         cdef shared_ptr[CArray] result
 
+        offset = operator.index(offset)
         if offset < 0:
             raise IndexError('Offset must be non-negative')
 
         offset = min(len(self), offset)
         if length is None:
             result = self.ap.Slice(offset)
         else:
+            length = operator.index(length)
             if length < 0:
                 raise ValueError('Length must be non-negative')
             result = self.ap.Slice(offset, length)

python/pyarrow/table.pxi

@@ -911,13 +911,15 @@ cdef class ChunkedArray(_PandasConvertible):
         """
         cdef shared_ptr[CChunkedArray] result
 
+        offset = operator.index(offset)
         if offset < 0:
             raise IndexError('Offset must be non-negative')
 
         offset = min(len(self), offset)
         if length is None:
             result = self.chunked_array.Slice(offset)
         else:
+            length = operator.index(length)
             result = self.chunked_array.Slice(offset, length)
 
         return pyarrow_wrap_chunked_array(result)
@@ -3172,13 +3174,15 @@ cdef class RecordBatch(_Tabular):
         """
         cdef shared_ptr[CRecordBatch] result
 
+        offset = operator.index(offset)
         if offset < 0:
             raise IndexError('Offset must be non-negative')
 
         offset = min(len(self), offset)
         if length is None:
             result = self.batch.Slice(offset)
         else:
+            length = operator.index(length)
             result = self.batch.Slice(offset, length)
 
         return pyarrow_wrap_batch(result)
@@ -4292,13 +4296,15 @@ cdef class Table(_Tabular):
         """
         cdef shared_ptr[CTable] result
 
+        offset = operator.index(offset)
         if offset < 0:
             raise IndexError('Offset must be non-negative')
 
         offset = min(len(self), offset)
         if length is None:
             result = self.table.Slice(offset)
         else:
+            length = operator.index(length)
             result = self.table.Slice(offset, length)
 
         return pyarrow_wrap_table(result)

python/pyarrow/tests/test_array.py

@@ -563,6 +563,63 @@ def test_array_slice():
                 assert res.to_numpy().tolist() == expected
 
 
+@pytest.mark.parametrize("scalar_type", [
+    pa.int8(), pa.int16(), pa.int32(), pa.int64(),
+    pa.uint8(), pa.uint16(), pa.uint32(), pa.uint64(),
+])
+def test_array_slice_integer_scalars(scalar_type):
+    arr = pa.array(range(10))
+    offsets = pa.array([2, 4], type=scalar_type)
+
+    result = arr.slice(offsets[0], offsets[1])
+
+    assert result.equals(arr.slice(2, 4))
+
+
+@pytest.mark.numpy
+def test_array_slice_numpy_integer_scalars():
+    arr = pa.array(range(10))
+
+    result = arr.slice(np.int64(2), np.int64(4))
+
+    assert result.equals(arr.slice(2, 4))
+
+
+@pytest.mark.parametrize("scalar", [
+    pa.scalar(2.0),
+    pa.scalar(True),
+    pa.scalar(None, type=pa.int64()),
+])
+@pytest.mark.parametrize("args", [
+    lambda scalar: (scalar,),
+    lambda scalar: (0, scalar),
+])
+def test_array_slice_invalid_scalars(scalar, args):
+    arr = pa.array(range(10))
+
+    with pytest.raises(TypeError):
+        arr.slice(*args(scalar))
+
+
+def test_array_slice_negative_integer_scalars():
+    arr = pa.array(range(10))
+    negative = pa.scalar(-1, type=pa.int64())
+
+    with pytest.raises(IndexError):
+        arr.slice(negative)
+    with pytest.raises(ValueError):
+        arr.slice(0, negative)
+
+
+def test_array_slice_uint64_scalar_overflow():
+    arr = pa.array(range(10))
+    overflow = pa.scalar(2 ** 63, type=pa.uint64())
+
+    assert arr.slice(overflow).equals(arr.slice(len(arr)))
+    with pytest.raises(OverflowError):
+        arr.slice(0, overflow)
+
+
 def test_array_slice_negative_step():
     # ARROW-2714
     values = list(range(20))

python/pyarrow/tests/test_table.py

@@ -1372,6 +1372,48 @@ def _table_like_slice_tests(factory):
     assert obj.slice(len(obj) - 4, 2).equals(obj[-4:-2])
 
 
+@pytest.mark.parametrize("factory", [
+    pa.chunked_array,
+    pa.RecordBatch.from_arrays,
+    pa.table,
+])
+@pytest.mark.parametrize("scalar_type", [
+    pa.int8(), pa.int16(), pa.int32(), pa.int64(),
+    pa.uint8(), pa.uint16(), pa.uint32(), pa.uint64(),
+])
+def test_table_like_slice_integer_scalars(factory, scalar_type):
+    data = [pa.array(range(10))]
+    names = ["c0"]
+    if factory is pa.chunked_array:
+        obj = factory(data)
+    else:
+        obj = factory(data, names=names)
+    offsets = pa.array([2, 4], type=scalar_type)
+
+    result = obj.slice(offsets[0], offsets[1])
+
+    assert result.equals(obj.slice(2, 4))
+
+
+@pytest.mark.parametrize("factory", [
+    pa.chunked_array,
+    pa.RecordBatch.from_arrays,
+    pa.table,
+])
+def test_table_like_slice_uint64_scalar_overflow(factory):
+    data = [pa.array(range(10))]
+    names = ["c0"]
+    if factory is pa.chunked_array:
+        obj = factory(data)
+    else:
+        obj = factory(data, names=names)
+    overflow = pa.scalar(2 ** 63, type=pa.uint64())
+
+    assert obj.slice(overflow).equals(obj.slice(len(obj)))
+    with pytest.raises(OverflowError):
+        obj.slice(0, overflow)
+
+
 def test_recordbatch_slice_getitem():
     return _table_like_slice_tests(pa.RecordBatch.from_arrays)