From 183f35a82fca72f754362c5ded0293367e795719 Mon Sep 17 00:00:00 2001 From: YusefSyed <211442445+YusefSyed@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:57:13 -0400 Subject: [PATCH] GH-51145: [Python] Accept Arrow integer scalars in slice methods --- python/pyarrow/array.pxi | 3 ++ python/pyarrow/table.pxi | 6 ++++ python/pyarrow/tests/test_array.py | 57 ++++++++++++++++++++++++++++++ python/pyarrow/tests/test_table.py | 42 ++++++++++++++++++++++ 4 files changed, 108 insertions(+) diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 2b2130e992e..7adc77924da 100644 --- a/python/pyarrow/array.pxi +++ b/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,6 +1632,7 @@ cdef class Array(_PandasConvertible): """ cdef shared_ptr[CArray] result + offset = operator.index(offset) if offset < 0: raise IndexError('Offset must be non-negative') @@ -1638,6 +1640,7 @@ cdef class Array(_PandasConvertible): 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) diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi index cb6a2e0acb4..7cecbe5b866 100644 --- a/python/pyarrow/table.pxi +++ b/python/pyarrow/table.pxi @@ -911,6 +911,7 @@ cdef class ChunkedArray(_PandasConvertible): """ cdef shared_ptr[CChunkedArray] result + offset = operator.index(offset) if offset < 0: raise IndexError('Offset must be non-negative') @@ -918,6 +919,7 @@ cdef class ChunkedArray(_PandasConvertible): 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,6 +3174,7 @@ cdef class RecordBatch(_Tabular): """ cdef shared_ptr[CRecordBatch] result + offset = operator.index(offset) if offset < 0: raise IndexError('Offset must be non-negative') @@ -3179,6 +3182,7 @@ cdef class RecordBatch(_Tabular): 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,6 +4296,7 @@ cdef class Table(_Tabular): """ cdef shared_ptr[CTable] result + offset = operator.index(offset) if offset < 0: raise IndexError('Offset must be non-negative') @@ -4299,6 +4304,7 @@ cdef class Table(_Tabular): 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) diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index a1e3616c9ce..eb81a4409e8 100644 --- a/python/pyarrow/tests/test_array.py +++ b/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)) diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index cfe47e4ed05..eb02051693f 100644 --- a/python/pyarrow/tests/test_table.py +++ b/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)