[Datasets] Fix ndarray representation of single-element ragged tensor slices.

[clarkzinzow]

Single-element ragged tensor slices (e.g. a[5:6]) currently have the wrong NumPy representation; namely, although they are object-dtyped, they have a multi-dimensional shape, and its single tensor element isn't well-typed (in other words, it doesn't use the pointer-to-subdarrays representation). This is due to np.array([subndarray], dtype=object) trying to create a more consolidated representation than np.array([subndarray1, subndarray2], dtype=object). This causes single-element batches of ragged tensor slices failing to eventually be put back into the tensor extension representation.

This PR fixes this by doing a very explicit ragged tensor construction via the create-and-fill method: we allocate an empty, object-dtyped 1D array and fill it with the tensor elements. This prevents NumPy from trying to optimize the ragged tensor representation.

Example

Have subndarray = np.array([[1, 2], [3, 4]], dtype=np.int64) and arr be an (N, 2, 2) ndarray.

N > 1 (Single-Element) Case

With arr = np.array([subndarray, subndarray], dtype=object]), you get

ndarray(
    [
        ndarray([[1, 2], [3, 4]], dtype=np.int64),
        ndarray([[1, 2], [3, 4]], dtype=np.int64),
    ],
    dtype=object
)

I.e., you get a 1D array of pointers to the subndarrays, so the subndarrays are well-typed, e.g. arr[0].dtype == np.int64.

N == 1 (Multi-Element) Case

But the single-element case, arr = np.array([subndarray], dtype=object), tries to consolidate the inner subndarray into the outer ndarray representation:

ndarray([[[1, 2], [3, 4]]], dtype=object)

The biggest impact of this is arr[0].dtype == object, which breaks a big tensor extension array assumption: that individual tensor elements are well-typed.

Related issue number

Closes #30513, closes #30406, closes #30059

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 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 :(

[bveeramani]

LGTM w/ comments

[bveeramani]

Are we converting to NumPy here so we can get the dtype?

[clarkzinzow]

In addition to using NumPy as the "source of truth" for slicing semantics for multi-dimensional arrays (which we typically do), we're primarily interested in whether the slicing results in the expected semantics for the NumPy views of the data rather than the Arrow data, e.g. the fact that slicing preserves the Arrow-level typing doesn't need to be tested, that's guaranteed by Arrow's extension type slicing, but we need to make sure that the NumPy-level dtype is what we'd expect.

[bveeramani]

Why doesn't ArrowTensorArrow expose dtype and shape methods like TensorArray?

[clarkzinzow]

Both the data type and element shape are exposed on the underlying tensor array type, i.e. ata.type.storage.value_type for the Arrow type, ata.type.to_pandas_dtype().numpy_dtype for the NumPy dtype, and ata.type.shape for the element shape.

We haven't exposed them directly on ArrowTensorArray under the assumption that the user will typically be working with the Pandas-side TensorArray extension type (if they're interacting with tensor extension types at all; they're hidden by default), so the Arrow side isn't optimized for end-user interaction.

Agreed that making the Arrow-side extension type more user-friendly is worth doing though, we should open a ticket for that.