Preserve dtype for numpy/torch/tf/jax arrays

[bhavitvyamalik]

Fixes #625. This lets the user preserve the dtype of numpy array to pyarrow array which was getting lost due to conversion of numpy array -> list -> pyarrow array.

[bhavitvyamalik]

Hi @lhoestq,
It turns out that pyarrow ListArray are not recognized as list-like when we get output from numpy_to_pyarrow_listarray. This might cause tests to fail. If possible can we convert that ListArray output to list inorder for tests to pass? Under the hood it'll maintain the dtype as that of numpy array passed during input only

[lhoestq]

Thanks ! To fix this I added a comment to actually keep the numpy array unchanged until it is passed to a TypedSequence. This way we don't have to deal with the ListArray issue

[bhavitvyamalik]

Brought down the failing tests from 7 to 4. Let me know if that part looks good. Failing tests are looking quite similar. In test_map_torch

datasets/tests/test_arrow_dataset.py

Line 1039 in 3d46bc3

Features({"filename": Value("string"), "tensor": Sequence(Value("float64"))}),

and test_map_tf

datasets/tests/test_arrow_dataset.py

Line 1056 in 3d46bc3

Features({"filename": Value("string"), "tensor": Sequence(Value("float64"))}),

they're expecting float64. Shouldn't that be float32 now?

[lhoestq]

It's normal: pytorch and tensorflow use float32 by default, unlike numpy which uses float64.

I think that we should always keep the precision of the original tensor (torch/tf/numpy).
It means that as it is in this PR it's fine (the precision is conserved when doing the torch/tf -> numpy conversion).

This is a breaking change but in my opinion the fact that we had Value("float64") for torch.float32 tensors was an issue already.

Let me know what you think. Cc @albertvillanova if you have an opinion on this

If we agree on doing this breaking change, we can just change the test.

[bhavitvyamalik]

Hi @lhoestq,
Merged master into this branch. Only changing the test is left for now (mentioned below) after which all tests should pass.

Brought down the failing tests from 7 to 4. Let me know if that part looks good. Failing tests are looking quite similar. In test_map_torch

datasets/tests/test_arrow_dataset.py

Line 1039 in 3d46bc3

Features({"filename": Value("string"), "tensor": Sequence(Value("float64"))}),

and test_map_tf

datasets/tests/test_arrow_dataset.py

Line 1056 in 3d46bc3

Features({"filename": Value("string"), "tensor": Sequence(Value("float64"))}),

they're expecting float64. Shouldn't that be float32 now?

[lhoestq]

they're expecting float64. Shouldn't that be float32 now?

Yes feel free to update those tests :)

It would be nice to have the same test for JAX as well

[bhavitvyamalik]

Added same test for for JAX too. Also, I saw that I missed changing test_cast_to_python_objects_jax like I did for TF and PyTorch. Finished that as well

[lhoestq]

It looks all good !
Thanks a lot :)