[Data] read_images benchmark sometimes fails with ArrowVariableShapedTensorArray error

[bveeramani]

What happened + What you expected to happen

See https://buildkite.com/ray-project/release/builds/30278#01946784-2b88-46ea-a571-c4a7afb0c16f.

ray.exceptions.RayTaskError(ValueError): ray::ReadImage() (pid=2078, ip=10.0.36.228)
--
  | for b_out in map_transformer.apply_transform(iter(blocks), ctx):
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/data/_internal/execution/operators/map_transformer.py", line 395, in __call__
  | yield output_buffer.next()
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/data/_internal/output_buffer.py", line 99, in next
  | block_to_yield = block.slice(0, target_num_rows, copy=True)
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/data/_internal/arrow_block.py", line 222, in slice
  | view = transform_pyarrow.combine_chunks(view)
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/data/_internal/arrow_ops/transform_pyarrow.py", line 382, in combine_chunks
  | new_column_values_arrays.append(combine_chunked_array(col))
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/data/_internal/arrow_ops/transform_pyarrow.py", line 416, in combine_chunked_array
  | return _concatenate_extension_column(array)
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/air/util/transform_pyarrow.py", line 35, in _concatenate_extension_column
  | return ArrowTensorArray._concat_same_type(ca.chunks)
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/air/util/tensor_extensions/arrow.py", line 798, in _concat_same_type
  | return ArrowVariableShapedTensorArray.from_numpy(
  | File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/air/util/tensor_extensions/arrow.py", line 1012, in from_numpy
  | raise ValueError(
  | ValueError: ArrowVariableShapedTensorArray only supports tensor elements that all have the same number of dimensions, but got tensor elements with dimensions: 3, 2

Versions / Dependencies

4bea94c567

Reproduction script

Wasn't able to narrow down the issue to a minimal repro.

Issue Severity

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

[raulchen]

IIRC, It's because the input images are in multiple different formats and thus producing rows with different number of dimensions. Adding read_image(..., mode="RGB") fixes this issue.