diff --git a/src/awkward/_connect/numba/arrayview.py b/src/awkward/_connect/numba/arrayview.py index c09b470536..ec19c1c7d1 100644 --- a/src/awkward/_connect/numba/arrayview.py +++ b/src/awkward/_connect/numba/arrayview.py @@ -1068,6 +1068,7 @@ def lower_asarray(context, builder, sig, args): whichpos = ak._connect.numba.layout.posat( context, builder, viewproxy.pos, viewtype.type.ARRAY ) + arrayptr = ak._connect.numba.layout.getat( context, builder, viewproxy.arrayptrs, whichpos ) @@ -1075,11 +1076,16 @@ def lower_asarray(context, builder, sig, args): bitwidth = ak._connect.numba.layout.type_bitwidth(rettype.dtype) itemsize = context.get_constant(numba.intp, bitwidth // 8) - data = numba.core.cgutils.pointer_add( - builder, + data_pointer_type = context.get_value_type(numba.types.CPointer(rettype.dtype)) + + typed_arrayptr = builder.inttoptr( arrayptr, - builder.mul(viewproxy.start, itemsize), - context.get_value_type(numba.types.CPointer(rettype.dtype)), + data_pointer_type, + ) + + data = builder.gep( + typed_arrayptr, + [viewproxy.start], ) shape = context.make_tuple( diff --git a/src/awkward/_connect/numba/layout.py b/src/awkward/_connect/numba/layout.py index ebf614a88f..2051c0f670 100644 --- a/src/awkward/_connect/numba/layout.py +++ b/src/awkward/_connect/numba/layout.py @@ -4,7 +4,6 @@ import json -import llvmlite.ir import numba from numba.core.errors import NumbaTypeError, NumbaValueError @@ -58,25 +57,30 @@ def string_numba_lower( baseptr = ak._connect.numba.layout.getat( context, builder, viewproxy.arrayptrs, whichnextpos ) - rawptr = builder.add( - baseptr, - ak._connect.numba.layout.castint( - context, builder, viewtype.type.indextype.dtype, numba.intp, start - ), - ) - rawptr_cast = builder.inttoptr( - rawptr, - llvmlite.ir.PointerType(llvmlite.ir.IntType(numba.intp.bitwidth // 8)), + start_cast = ak._connect.numba.layout.castint( + context, builder, viewtype.type.indextype.dtype, numba.intp, start ) strsize = builder.sub(stop, start) strsize_cast = ak._connect.numba.layout.castint( context, builder, viewtype.type.indextype.dtype, numba.intp, strsize ) + # ``baseptr`` is stored in the Awkward lookup table as an integer + # address. Reconstruct a typed pointer explicitly before applying + # pointer arithmetic with GEP. + uint8_pointer_type = context.get_value_type(numba.types.CPointer(numba.types.uint8)) + typed_baseptr = builder.inttoptr( + baseptr, + uint8_pointer_type, + ) + rawptr = builder.gep( + typed_baseptr, + [start_cast], + ) pyapi = context.get_python_api(builder) gil = pyapi.gil_ensure() - strptr = builder.bitcast(rawptr_cast, pyapi.cstring) + strptr = builder.bitcast(rawptr, pyapi.cstring) if viewtype.type.parameters["__array__"] == "string": kind = context.get_constant(numba.types.int32, pyapi.py_unicode_1byte_kind) @@ -353,24 +357,33 @@ def posat(context, builder, pos, offset): def getat(context, builder, baseptr, offset, rettype=None): - ptrtype = None - if rettype is not None: - ptrtype = context.get_value_type(numba.types.CPointer(rettype)) - bitwidth = type_bitwidth(rettype) - else: - bitwidth = numba.intp.bitwidth - byteoffset = builder.mul(offset, context.get_constant(numba.intp, bitwidth // 8)) - out = builder.load( - numba.core.cgutils.pointer_add(builder, baseptr, byteoffset, ptrtype) + """ + Load an element from an external buffer address. + + Awkward's lookup table stores buffer pointers as integer addresses. + Convert the address explicitly to a typed pointer before using GEP. + """ + element_type = numba.intp if rettype is None else rettype + pointer_type = context.get_value_type(numba.types.CPointer(element_type)) + typed_baseptr = builder.inttoptr( + baseptr, + pointer_type, ) + + element_ptr = builder.gep( + typed_baseptr, + [offset], + ) + + out = builder.load(element_ptr) + if rettype is not None and isinstance(rettype, numba.types.Boolean): return builder.icmp_signed( "!=", out, context.get_constant(numba.int8, 0), ) - else: - return out + return out def regularize_atval(context, builder, viewproxy, attype, atval, wrapneg, checkbounds): diff --git a/tests/test_4251_numba_intaddr_pointer_reconstruction.py b/tests/test_4251_numba_intaddr_pointer_reconstruction.py new file mode 100644 index 0000000000..885858d39c --- /dev/null +++ b/tests/test_4251_numba_intaddr_pointer_reconstruction.py @@ -0,0 +1,68 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE + +from __future__ import annotations + +import pytest + +import awkward as ak + +numba = pytest.importorskip("numba") + +ak.numba.register_and_check() + + +@pytest.mark.parametrize( + "array", + [ + ak.Array([True, False, True]), + ak.Array([1, 2, 3]), + ak.Array([1.5, 2.5, 3.5]), + ak.Array([[1, 2], [], [3]]), + ak.Array(["hello", "world"]), + ak.Array([b"hello", b"\x00\xff"]), + ], +) +def test_numba_getitem_after_pointer_conversion(array): + @numba.njit + def getitem(array, index): + return array[index] + + for index, expected in enumerate(array.to_list()): + assert ak.to_list(getitem(array, index)) == expected + + +def test_numba_ragged_string_membership(): + array = ak.Array( + [ + ["TRA", "TRB"], + [], + None, + ["TRG", "TRA"], + ] + ) + + @numba.njit + def isin(array, haystack, builder): + for row in array: + builder.begin_list() + + if row is not None: + for value in row: + builder.append(value in haystack) + + builder.end_list() + + return builder + + result = isin( + array, + ("TRA", "TRB"), + ak.ArrayBuilder(), + ).snapshot() + + assert result.to_list() == [ + [True, True], + [], + [], + [False, True], + ]