[Datasets] Ragged Arrow datasets have incorrect schema

[bveeramani]

What happened + What you expected to happen

I created a dataset from ragged arrays. I expected the schema type to be ArrowVariableShapedTensorArray, but I got ArrowTensorArray instead.

Versions / Dependencies

Ray: 5aa1bb0
Python: 3.7.15
OS: Linux
Arrow: 6.0.1

Reproduction script

>>> import ray
>>> import numpy as np
>>> ds = ray.data.from_items([{"spam": np.zeros((32, 32, 3))}, {"spam": np.zeros((64, 64, 3))}])
>>> ds.schema().types
[ArrowTensorType(shape=(32, 32, 3), dtype=double)]

Issue Severity

Medium: It is a significant difficulty but I can work around it.

[clarkzinzow]

This is because the default parallelism of ray.data.from_items() will give a block per row, and we perform type inference via only the first block. Since the blocks are trivially homogeneous-shaped (they each contain a single tensor element), the type comes back as ArrowTensorType.

We could do some schema promotion/unification if/when we have metadata for all blocks on hand. We already have this schema unification logic in the table concatenation path, all that we'd need to do is to factor it out of table concatenation and use it when fetching schema.

[bveeramani]

Looks like this can an issue for read_images, too:

>>> ds = ray.data.read_images("python/ray/data/examples/data/image-datasets/different-sizes")
>>> ds
Dataset(num_blocks=3, num_rows=3, schema={image: ArrowTensorType(shape=(16, 16, 3), dtype=uint8)})

Since the blocks are trivially homogeneous-shaped (they each contain a single tensor element), the type comes back as ArrowTensorType.

To confirm my understanding, is it the case that whenever num_blocks = num_rows, ragged schemas won't be inferred?

[clarkzinzow]

To confirm my understanding, is it the case that whenever num_blocks = num_rows, ragged schemas won't be inferred?

Correct, that will always be the case with the current schema fetching.

[xwjiang2010]

@bveeramani Can you triage the issue as well? (by removing the triage label and adding a priority label) Thanks!

[bveeramani]

@xwjiang2010 done!

[bveeramani]

Copying relevant offline discussion:

[To implement this you] should just need to create a unify_schemas(schemas: List[pa.Schema]) -> pa.Schema utility that’s used by both table concatenation and when fetching schema, and then some minor refactoring of how the ArrowTensorType --> ArrowVariableShapedTensorType promotion is implemented. Right now, we chunk the columns for all tables and see if the type was promoted to the variable-shaped tensor type (link), but I think we’d want to change this to an explicit ArrowTensorType._need_variable_shaped_tensor_array(types) check, where we refactor the existing method to take a list of tensor types rather than a list of tensor arrays.

[bveeramani]

Bumping to P1 because this breaks prediction with ragged tensors