[Data] Change offsets to int64 and change to LargeList for ArrowTensorArray

[terraflops1048576]

Closes #46434

Why are these changes needed?

Currently, the ArrowTensorArray and ArrowVariableShapedTensorArray types use int32 to encode the list offsets. This means that within a given PyArrow chunk in a column, the sum of the sizes of all of the tensors in that chunk must be less than 2^31; otherwise, depending on the overflow conditions, an error is thrown or the data is truncated. This usually doesn't manifest itself in Ray Data with the default settings, because it splits the blocks up to meet the target max block size (though this can be turned off!). However, it unavoidably shows up when one needs a large local shuffle buffer to feed into Ray Train.

This PR changes the offsets to be stored in 64-bit integers and updates the corresponding storage types of the TensorArrays.

As an example:

import ray.train
import ray.train.torch
import ray.data
import numpy as np

def f(batch):
   block1 = {"x": [np.ones((1000, 550), dtype=np.float16)] * 1000}
   return block1

dataset = ray.data.from_items([1, 2, 3, 4, 5]).map_batches(f, batch_size=None)
def train():
  data = ray.train.get_dataset_shard("train")
  for batch in data.iter_torch_batches(batch_size=100, local_shuffle_buffer_size=4000):
    pass

trainer = ray.train.torch.TorchTrainer(train_loop_per_worker=train, datasets={"train": dataset})
trainer.fit()

fails with

pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[c21]

Hi @terraflops1048576, thanks for fix! Can you also add a unit test?

[terraflops1048576]

I wasn't sure what unit test I should exactly add -- this PR is meant to change the list so its behavior still aligns with the existing unit tests, but works for large numbers of tensor elements. I added a unit test that checks that it actually works with the large numbers of tensors.

[scottjlee]

thanks for the fix! merging in master since it's been a few weeks, to ensure tests still pass

[raulchen]

Nice fix, thank you!

[terraflops1048576]

I think some of the test failures are unrelated, something to do with rllib. I see that one of the tests checks the size of the array; by my calculation, the array should use exactly 768 bytes to represent the tensors themselves, with 6 * 15 * 8 = 720 bytes for the values themselves, along with one 64-bit offset value per tensor. The rest is in the shape serialization, I believe. I updated the test to more accurately calculate the size, and hopefully it's within tolerance.

[can-anyscale]

Breaking many release tests: #46499, #46496, #46495. I'm reverting to unblock the release.