fix: ak.num should always return a useful (non-unknown length) type

[agoose77]

Replaces #2774 by ensuring the return type of ak.num is always a valid "index type" (something that kernels produce).

It might be that this is an integer rather than an np.ndarray (or cp.ndarray, etc); this is up to the nplike. Right now, nplikes don't promote Python integers to 0D arrays.

[codecov]

Codecov Report

Merging #2785 (b8ace29) into main (23f5322) will increase coverage by 0.00%.
The diff coverage is 66.66%.

Additional details and impacted files

Files	Coverage Δ
src/awkward/operations/ak_num.py	93.75% <66.66%> (+0.20%)	⬆️

[jpivarski]

Just one fix, and then it's ready to merge (i.e. no need for a second review).

[agoose77]

@jpivarski to confirm — forcibly returning ArrayLike values from ak.num means that it's not possible to use ak.num to return concrete lengths for typetracers with known shapes. I assume you're OK with that, but I want to surface before we merge.

[jpivarski]

@jpivarski to confirm — forcibly returning ArrayLike values from ak.num means that it's not possible to use ak.num to return concrete lengths for typetracers with known shapes. I assume you're OK with that, but I want to surface before we merge.

I don't see anything wrong with that—internally in the codebase, we should be using .length and users have len if they want to make a distinction between getting the length eagerly at DAG-building time (if the TypeTracer has a known shape) and getting the length as delayed data until it's taken up by a Dask worker.

In fact, if I had to make that distinction, I would expect the Awkward function to give me delayed data and the Python function to give me eager data, because Awkward knows about Dask and Python doesn't.

I guess we're losing a case of "give me the eager, numerical length if you have it and a delayed token if you don't." But if I wanted that case, I wouldn't have guessed that ak.num would provide it (and it would be a tricky thing to explain in the documentation—likely undiscoverable, anyway). If I needed this case, I'd write a try-except around len.