fix: make ak.enforce_type work with typetracers

[ikrommyd]

Closes #3736

As the title says, this makes ak.enforce_type work with typetracers. This is needed in order to make dak.enforce_type work in dask-awkward.

[ikrommyd]

Opening it as draft as I need to add tests for many possible combinations for the enforce type arguments. This is what I only saw immediately to break the infinite recursion reported in the issue.

[codecov]

Codecov Report

❌ Patch coverage is 70.00000% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.75%. Comparing base (5a6e69f) to head (1a29afd).

Files with missing lines	Patch %	Lines
src/awkward/operations/ak_from_buffers.py	33.33%	2 Missing ⚠️
src/awkward/forms/form.py	80.00%	1 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/operations/ak_enforce_type.py	82.16% <100.00%> (ø)
src/awkward/forms/form.py	81.31% <80.00%> (-0.02%)	⬇️
src/awkward/operations/ak_from_buffers.py	93.36% <33.33%> (-0.81%)	⬇️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3764

[pfackeldey]

Thanks @ikrommyd!

This looks like the correct fix, I've tested and confirmed it locally as well. There may be certainly cases like this one where special typetracer handling has been overlooked in the past.

Please add the test as you suggested, then I think it's ready from my side.

[ikrommyd]

@pfackeldey I have added 1K lines of tests all taken from https://github.com/scikit-hep/awkward/blob/main/tests/test_2365_enforce_type.py.
I basically took all those tests and adapted them to check equality of typetracer-backed layouts with the only difference from the original tests regarding assertions being the fact that errors being raised at runtime due to the contents of the data, cannot be raised by typetracers as they do not hold any data. enforce_type basically has to assume that you can actually enforce the type in such cases and the error will be raised during evaluation of the graph when actual data are loaded. For example, using dask-awkward with this PR, this is typetracable but the final line errors upon computation

In [28]: array = dak.from_awkward(ak.Array([1, "hi"]), 1)

In [29]: dak.lib.structure.enforce_type(array, "string")
Out[29]: dask.awkward<enforce-type, type='2 * string', npartitions=1>

In [30]: dak.lib.structure.enforce_type(array, "string").compute()

ValueError: UnionArray(s) can only be converted to string if they are equivalent to their projections

This error occurred while calling

    ak.enforce_type(
        <Array [1, 'hi'] type='2 * union[int64, string]'>
        type = 'string'
        behavior = None
        attrs = None
    )

However, if you slice out the number and keep the string only, it is computable even

In [31]: dak.lib.structure.enforce_type(array[1:], "string").compute()
Out[31]: <Array ['hi'] type='1 * string'>

So during tracing you need to assume that your type enforcement actually works until you look at the data.

In general, a few minor changes were needed. Inspired from the fact that length_zero/one_array uses ak.from_buffers, I also added an actual error message in from_buffers when typetracer backend is requested because it internally has no handling of typetracers and actually never expects the buffers to be typetracers.
to_buffers guards this here:

awkward/src/awkward/_do.py

Lines 87 to 88 in 0ad8b1e

	if not backend.nplike.known_data:
	raise TypeError("cannot call 'to_buffers' on an array without concrete data")

so from_buffers should too.

This should be ready for review from my side.

[ikrommyd]

To support this PR more, I have exposed it to the public API of dask-awkward in this commit in my fork and have added tests for it: dask-contrib/dask-awkward@43fa7ef. They all pass with awkward from this PR branch so I feel pretty confident that this is good to go.