Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[feature request] assign a tuple to one row of recarray #9476

Open
dlee992 opened this issue Mar 5, 2024 · 8 comments · May be fixed by #9478
Open

[feature request] assign a tuple to one row of recarray #9476

dlee992 opened this issue Mar 5, 2024 · 8 comments · May be fixed by #9478

Comments

@dlee992
Copy link
Contributor

dlee992 commented Mar 5, 2024

test code:

import numpy as np
from numba import njit


@njit
def foo(arr, name, age, height):
    items = (name, age, height)
    arr[1] = items
    return arr


dtype = [("Name", "U10"), ("Age", "i4"), ("Height", "f4")]
recarr = np.empty((2,), dtype=dtype)
recarr[0] = ("John", 25, 180.5)

print(recarr)
print(foo(recarr, "John", 25, 180.5))

np supports directly assign a tuple of given types to one row of recarray, but numba doesn't.

No implementation of function Function(<built-in function setitem>) found for signature:
 
 >>> setitem(unaligned array(Record(Name[type=[unichr x 10];offset=0],Age[type=int64;offset=40],Height[type=float64;offset=48];56;False), 1d, C), Literal[int](1), Tuple(unicode_type, int64, float64))

If could, the user code can be only one-line, rather than multiple lines to assign each item.

kind of relating with #8471

@sklam
Copy link
Member

sklam commented Mar 5, 2024

Yes, this is doable without too much difficulty.

We should expect to see as output from the OP's program:

[('John', 25, 180.5) ('',  0,   0. )]
[('John', 25, 180.5) ('John', 25, 180.5)]

@gmarkall
Copy link
Member

gmarkall commented Mar 5, 2024

Example in-progress I had:

import numpy as np
import operator
from numba import cuda
from numba.core import types
from numba.core.extending import lower_cast
from numba.core.typeconv import Conversion
from numba.core.typing.templates import (AbstractTemplate, infer_global,
                                         signature)


@lower_cast(types.UniTuple, types.Record)
@lower_cast(types.Tuple, types.Record)
def tuple_to_record(context, builder, fromty, toty, val):
    breakpoint()


possible_conversions = (
    Conversion.exact,
    Conversion.promote,
    Conversion.safe,
    Conversion.unsafe
)


@infer_global(operator.setitem)
class SetItemRecordArray(AbstractTemplate):
    def generic(self, args, kws):
        assert not kws
        ary, idx, val = args
        for a, b in zip(ary.dtype.fields.values(), val.types):
            conv = self.context.can_convert(b, a.type)

            if conv not in possible_conversions:
                return

        return signature(types.none, ary, idx, val)


def test():

    dt = np.dtype([('x', np.float32), ('y', np.float32)], align=True)

    @cuda.jit
    def f(input, output):
        # v = input[0]
        # x, y = v.x, v.y
        # output[0] = v # works
        #output[0] = np.record((3.14, 2.71), dtype=dt)  # nope
        # output[0].x = 3.14 # OK
        # output[0].y = 2.71 # OK
        output[0] = (3.14, 2.71)

    input = np.array([(3.14, 2.71)], dtype=dt)
    output = np.zeros_like(input)

    f[1, 1](input, output)
    cuda.synchronize()

    print(input, output)
    assert np.allclose(input.view(dtype=np.float32),
                       output.view(dtype=np.float32))


test()

@sklam
Copy link
Member

sklam commented Mar 5, 2024

Just talked about this in triage, maybe it is not that straightforward once nested structs are considered.

@dlee992
Copy link
Contributor Author

dlee992 commented Mar 5, 2024

Thanks for the impl draft! Will follow up soon. I think it's still good if the case without nested structs can work for now.

@gmarkall
Copy link
Member

gmarkall commented Mar 5, 2024

I think @stuartarchibald has a better draft than me - taking mine and going further with it might not be the best way forward.

@dlee992
Copy link
Contributor Author

dlee992 commented Mar 5, 2024

Can @stuartarchibald's draft be shared with me? Thanks! I can follow up to make it as a PR.

@dlee992
Copy link
Contributor Author

dlee992 commented Mar 5, 2024

maybe I can directly change some code in

@infer_global(operator.setitem)
class SetItemBuffer(AbstractTemplate):
for typing.

@dlee992 dlee992 linked a pull request Mar 6, 2024 that will close this issue
@stuartarchibald
Copy link
Contributor

I think another way of doing this would involve using literal_unroll to unroll the assignment expression across the record/tuple, am still working on the details as the nested case (assignment to a record containing records) is a little tricky and might not be possible via this approach.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging a pull request may close this issue.

4 participants