[Datasets] Correct schema unification for Datasets with ragged Arrow arrays

[scottjlee]

Signed-off-by: Scott Lee [email protected]

Why are these changes needed?

When creating Datasets with ragged arrays, the resulting Dataset incorrectly uses ArrowTensorArray instead of ArrowVariableShapedTensorArray as the underlying schema type. This PR refactors existing logic for schema unification into a separate function, which is now called during Arrow table concatenation and schema fetching to correct type promotion involving ragged arrays.

Related issue number

Closes #30082

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

[clarkzinzow]

LGTM!

[clarkzinzow]

This can be a future PR (I can do it as part of the type promotion PR), but we might want to try-except this pyarrow.unify_schemas() call, since this is the point at which we're validating that all of the schemas from different blocks are compatible. Propagating any exception raised from pyarrow.unify_schemas() seems fine for now, and in the future we can look at wrapping any raised exceptions with our own error indicating that the blocks have incompatible schemas and giving the user a path to rectifying this (e.g. manually specifying a schema at read time, so all blocks are consistent).