fix: handle unhashable behaviour type

[ianna]

fixes #2762

[codecov]

Codecov Report

Merging #2770 (2716f8f) into main (c310ec1) will decrease coverage by 0.02%.
Report is 1 commits behind head on main.
The diff coverage is 88.88%.

Additional details and impacted files

Files	Coverage Δ
src/awkward/_connect/numba/builder.py	85.02% <100.00%> (+0.06%)	⬆️
src/awkward/_connect/numba/layoutbuilder.py	87.24% <83.33%> (-0.03%)	⬇️

... and 73 files with indirect coverage changes

[jpivarski]

This was already dealt with in

awkward/src/awkward/_connect/numba/arrayview.py

Lines 343 to 354 in 17aabde

	def dict2serializable(obj):
	if obj is None:
	return None
	else:
	return tuple(obj.items())
	def serializable2dict(obj):
	if obj is None:
	return None
	else:
	return dict(obj)

Perhaps ArrayBuilder could use the same mechanism as ArrayView?

Also, @agoose77 was going to send an example of a behavior that can't be serialized due to a lambda function. I hope we don't have to pull in cloudpickle just for that. But maybe Numba's serialization already does something clever: we should try it and find out.

[agoose77]

Here's a non-pickleable array:

import pickle

import numpy as np
import numba as nb
import pytest

import awkward as ak


def make_add_xyr():
    def add_xyr(left, right):
        x = left.x + right.x
        y = left.y + right.y
        return ak.zip(
            {
                "x": x,
                "y": y,
                "r": np.sqrt(x**2 + y**2),
            },
            with_name="xyr",
        )

    return add_xyr


behavior = {(np.add, "xyr", "xyr"): make_add_xyr()}


a = ak.Array([{"x": 3, "y": 4, "r": 5}], with_name="xyr", behavior=behavior)
b = ak.Array([{"x": 3, "y": 4, "r": 5}], with_name="xyr", behavior=behavior)

builder = ak.ArrayBuilder(behavior=behavior)


with pytest.raises(AttributeError):
    behavior_pkl = pickle.dumps(behavior)


@nb.njit
def make_ab(builder):
    builder.begin_record("xyz")
    builder.field("x").integer(3)
    builder.field("y").integer(4)
    builder.field("z").integer(3)
    builder.end_record()

    builder.begin_record("xyz")
    builder.field("x").integer(3)
    builder.field("y").integer(4)
    builder.field("z").integer(3)
    builder.end_record()
    return builder


result = make_ab(builder).snapshot()

print(result)

We discussed in the meeting perhaps stashing a reference in the Python layer, and retrieving it later on? Just in case we forgot this (though I doubt we did).

[ianna]

@agoose77 - I'm done with the PR. Please, check and modify if needed. For example, I put both tests in one - it could be split. Thanks!

[ianna]

@agoose77 - I'm done with the PR. Please, check and modify if needed. For example, I put both tests in one - it could be split. Thanks!

Nope, not quite finished with it :-(

[ianna]

@jpivarski - I see that the lambda functions have different identity:

{(<ufunc 'add'>, 'xyr', 'xyr'): <function test_ArrayBuilder_non_picklable_behavior.<locals>.make_add_xyr.<locals>.add_xyr at 0x12395de40>}
{(<ufunc 'add'>, 'xyr', 'xyr'): <function test_ArrayBuilder_non_picklable_behavior.<locals>.make_add_xyr.<locals>.add_xyr at 0x1242d25c0>}

[agoose77]

I've loosened the test to just ensure that the .keys() agree. The other test ensures that the behavior with non-serialised lambdas is equal, which I think is sufficient.