[data] [strict mode] Allow returning lists instead of arrays for numpy batches

[ericl]

Why are these changes needed?

Allow map_batches UDFs to return {"foo": [1, 2, 3]} in addition to {"foo": np.array([1, 2, 3])} by implicitly casting lists to arrays.

[amogkam]

awesome, this is great! are the internal numpy arrays converted back into python lists when iterating over the outputs via iter_batches? Otherwise, I can imagine this may be confusing

[pcmoritz]

Wow, this is genius, maybe it allows us to have the best of all worlds :D

[ericl]

awesome, this is great! are the internal numpy arrays converted back into python lists when iterating over the outputs via iter_batches? Otherwise, I can imagine this may be confusing

You still get np.arrays as output, but those are also Iterable (so if you try to use it as a list, it might be 95% the same?). Not sure if we can do better here without something like type extensions.

[pcmoritz]

One drastic possibility: Always convert np arrays of objects to lists. I.e. if map_batches returns a list, convert to np array of objects under the hood (this PR), if dataset has np arrays of objects, expose them to map_batches as a python list on the reading path. It is unconventional. We have to be careful. But it might be very ergonomic. If we do this, we would need to make it very clear in the docs, and tell people how to keep using numpy arrays of objects if they want to. If we do this, it needs careful user testing :)

[pcmoritz]

Alternatively, we document well how to use .tolist() on the reading path.

[ericl]

Yeah, I'm kind of skeptical of that because of the overhead. I think converting a dense numpy array into a python list may be zero copy if you use slice views, but if the user tries to concatenate that into a numpy array again that would almost certainly induce a copy of all the data.

If we have to choose one format, the numpy one is the one that allows the best performance in all cases.

[pcmoritz]

Sounds good to me. When I was playing around with the hugging face code, the return codepath was the most annoying thing to deal with, which is 100% addressed by this PR. Great job.

The .tolist() stuff for the reading path is very bearable I think -- if we document it well I think it is a good solution :)

[pcmoritz]

On the efficiency, one thought though: I'm pretty sure if you convert to a list, it will just do shallow copies of the objects and numpy arrays of objects are already not-so-performant (because the values are boxed, there is no vectorization and little cache locality). Just wanted to throw this out. So I would mostly focus on user experience here and less on efficiency (if we still want to experiment with this option).

[pcmoritz]

The main argument against having special handling for np arrays of objects for me would be it will be inconsistent between numpy arrays of int / float / etc and numpy arrays of objects, so might be confusing :)

[pcmoritz]

Ah I think I get your argument better now. If we do the conversion by default / standardize too much on this, all our users' data will be boxed and nobody will use numpy arrays of primitive types.

[pcmoritz]

I'm curious why you need the list( here -- shouldn't .items() always return an object that can be iterated on with the for key, value in batch.items() pattern?

[ericl]

This was to avoid mutating the dict we are iterating over. Though, maybe in this case it is ok, since it doesn't change the keys.

[pcmoritz]

if isinstance(value, list) would be a little clearer here (or is there a difference?)

[ericl]

Yup the issue precisely is how to "unbox" the array data without too much overheads.

If you default to arrays, then:

• you can call arr.tolist() if you need a list (cheap)
• otherwise keep it an array (cheap)

If you default to list, then:

• you need to call np.stack(arr) if you need an array (expensive)
• otherwise keep it a list (cheap)