For a dataset comprised of both empty and non-empty blocks, let the non-empty blocks determine the schema

[jianoaix]

Why are these changes needed?

There is a bug in combining the results from map_batches: if we create two dataset out of the same data, but with different num of partitions, we may get different results when run the same map_batches() on them. That is, num of partitions is affecting the map_batches() results, which should not.

Related issue number

Closes #22673

Checks

• [ Y] 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
  • [ Y] Unit tests
  • Release tests
  • This PR is not tested :(

[jianoaix]

I realize I should have created a new branch for a different issue. The previous commit logs make it a bit messy.

[clarkzinzow]

Nice!

[jjyao]

Could you explain the root cause of the issue? Schema of empty block is wrong? Should we always set schema of empty block to None?

[jianoaix]

Could you explain the root cause of the issue? Schema of empty block is wrong? Should we always set schema of empty block to None?

The schema for empty block is default to pyarrow (#22673 (comment)).

[jjyao]

If we don't want to trust the schema of empty block even if it's set, should we just set to None in the first place? Otherwise we will see inconsistent schemas for blocks in the same dataset?

[jianoaix]

The block is of type Union[List[T], "pyarrow.Table", "pandas.DataFrame", bytes], so even if it's empty, it already has a type and hence an implied the schema type.
Ideally if we know the output block type, we might pass it around and then use BlockBuilder.for_block_type(block_type).build() to create an empty one, which will have the right type and schema, even if it's empty. I was initially suggesting this approach (in the comment linked above), but it looks we might not know the output type, so here we use non-empty block's schema.

[ericl]

Seems tests are failing.

[jianoaix]

Seems tests are failing.

Yeah. As I was working on fixing it, I realized that it is not feasible to rely on num_rows in block metadata to determine whether a block is empty. IIUC, we do not necessarily always know num_rows because we employ lazy execution/lazy data loading.