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

[terraflops1048576]

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?