[Data] Limit operator push down

[goutamvenkat-anyscale]

Why are these changes needed?

Previously, the limit operator push down optimization was disabled because there was no guarantee that Map, MapBatches etc. were not modifying the number of rows. In this change, we pushdown the limit operator for Project, MapRows, MapBatches and for Union each union has its own limit operator pushed down on it.

Related issue number

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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[iamjustinhsu]

correct me if i'm wrong but don't u still need to

1. link limit_output_op.output_dependencies.input_dependencies to last_op.output_dependencies
2. ops_between_new_input_and_limit[-1] (most upstream op) input dependencies need to link with limit_output_op.output_dependencies

[goutamvenkat-anyscale]

Let me draw it out and double check this

[srinathk10]

We can pushdown limit into these Ops but not past it. For instance, limit pushdown into Sorts involve sorting on keys for the entire Dataset, but only yield limit rows on the Sorted Dataset.

Same for Joins, Groupbys.

[goutamvenkat-anyscale]

You're right. https://docs.rs/datafusion-optimizer/48.0.1/src/datafusion_optimizer/push_down_limit.rs.html#38 Datafusion does this for sort whereby limit is embedded into the sort operator.