diff --git a/python/pyarrow/pandas_compat.py b/python/pyarrow/pandas_compat.py index 5e2ee49437e..bb54c3b22c3 100644 --- a/python/pyarrow/pandas_compat.py +++ b/python/pyarrow/pandas_compat.py @@ -755,8 +755,8 @@ def _reconstruct_block(item, columns=None, extension_columns=None, return_block= # create ExtensionBlock arr = item['py_array'] assert len(placement) == 1 - name = columns[placement[0]] - pandas_dtype = extension_columns[name] + name = columns.get(placement[0], None) + pandas_dtype = extension_columns.get(name, None) if not hasattr(pandas_dtype, '__from_arrow__'): raise ValueError("This column does not support to be converted " "to a pandas ExtensionArray") diff --git a/python/pyarrow/tests/test_cuda_numba_interop.py b/python/pyarrow/tests/test_cuda_numba_interop.py index 876f3c7f761..3bd81d755f5 100644 --- a/python/pyarrow/tests/test_cuda_numba_interop.py +++ b/python/pyarrow/tests/test_cuda_numba_interop.py @@ -49,7 +49,7 @@ def teardown_module(module): @pytest.mark.parametrize("c", range(len(context_choice_ids)), ids=context_choice_ids) def test_context(c): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) assert ctx.handle == nb_ctx.handle.value assert ctx.handle == ctx.to_numba().handle.value ctx2 = cuda.Context.from_numba(nb_ctx) @@ -83,7 +83,7 @@ def make_random_buffer(size, target='host', dtype='uint8', ctx=None): @pytest.mark.parametrize("dtype", dtypes, ids=dtypes) @pytest.mark.parametrize("size", [0, 1, 8, 1000]) def test_from_object(c, dtype, size): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) arr, cbuf = make_random_buffer(size, target='device', dtype=dtype, ctx=ctx) # Creating device buffer from numba DeviceNDArray: @@ -161,7 +161,7 @@ def __cuda_array_interface__(self): ids=context_choice_ids) @pytest.mark.parametrize("dtype", dtypes, ids=dtypes) def test_numba_memalloc(c, dtype): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) dtype = np.dtype(dtype) # Allocate memory using numba context # Warning: this will not be reflected in pyarrow context manager @@ -184,7 +184,7 @@ def test_numba_memalloc(c, dtype): ids=context_choice_ids) @pytest.mark.parametrize("dtype", dtypes, ids=dtypes) def test_pyarrow_memalloc(c, dtype): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) size = 10 arr, cbuf = make_random_buffer(size, target='device', dtype=dtype, ctx=ctx) @@ -198,7 +198,7 @@ def test_pyarrow_memalloc(c, dtype): ids=context_choice_ids) @pytest.mark.parametrize("dtype", dtypes, ids=dtypes) def test_numba_context(c, dtype): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) size = 10 with nb_cuda.gpus[0]: arr, cbuf = make_random_buffer(size, target='device', @@ -217,7 +217,7 @@ def test_numba_context(c, dtype): ids=context_choice_ids) @pytest.mark.parametrize("dtype", dtypes, ids=dtypes) def test_pyarrow_jit(c, dtype): - ctx, nb_ctx = context_choices[c] + ctx, nb_ctx = context_choices.get(c, (None, None)) @nb_cuda.jit def increment_by_one(an_array): diff --git a/python/pyarrow/tests/test_gdb.py b/python/pyarrow/tests/test_gdb.py index 912953ae60d..58aabb7368e 100644 --- a/python/pyarrow/tests/test_gdb.py +++ b/python/pyarrow/tests/test_gdb.py @@ -159,7 +159,7 @@ def select_frame(self, func_name): if m is None: pytest.fail(f"Could not select frame for function {func_name}") - frame_num = int(m[1]) + frame_num = int(m.get(1, None)) out = self.run_command(f"frame {frame_num}") assert f"in {func_name}" in out diff --git a/python/stubs/__lib_pxi/array.pyi b/python/stubs/__lib_pxi/array.pyi index 17eb4c6d888..ffdb8a9c075 100644 --- a/python/stubs/__lib_pxi/array.pyi +++ b/python/stubs/__lib_pxi/array.pyi @@ -1992,6 +1992,7 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def __getitem__(self, key: int) -> _Scalar_co: ... @overload def __getitem__(self, key: builtins.slice) -> Self: ... + def __getitem__(self, key): """ Slice or return value at given index