[#1213] feat(rust): Support block filter by taskId when getting memory data

[wenlongbrother]

What changes were proposed in this pull request?

Support block filter by taskId when getting memory data

Why are the changes needed?

In AQE, after executing the sub-QueryStages, collect the shuffle data size, So if we can filter block,
it will improve the performance of AQE.

Fix: #1213

Does this PR introduce any user-facing change?

No.

How was this patch tested?

1. UTs

[wenlongbrother]

Hi @zuston , could you please help reivew this pr , I may have lost some context.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (d5de299) 53.35% compared to head (8335f44) 54.41%.
Report is 6 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master    #1311      +/-   ##
============================================
+ Coverage     53.35%   54.41%   +1.05%     
- Complexity     2690     2695       +5     
============================================
  Files           410      391      -19     
  Lines         23590    21283    -2307     
  Branches       2003     2006       +3     
============================================
- Hits          12587    11581    -1006     
+ Misses        10229     8997    -1232     
+ Partials        774      705      -69     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[zuston]

Could you help add some UTs?

[wenlongbrother]

Could you help add some UTs?

sure, @zuston plz do review again, thank you. 😄

[zuston]

This part could be referred by the java side shuffle-server

[wenlongbrother]

Hi @zuston , could you please help review this pr again? Thanks! I have update ut and parameter

[zuston]

Overall lgtm. I believe this pr is in good shape

[wenlongbrother]

Hi @zuston thank you for your review and helpful comment.

[zuston]

This part looks not correct. I hope this could be bound to the read_partial_data_with_max_size_limit .

Without the bitmap filter, we will read it the (last_block_id, end] without the limited size. But for now, I think it will read from (last_block_id, end] with filter and then in a limited size. Right?

So it's better to refactor the read_partial_data_with_max_size_limit, maybe the name of read_partial_data_with_max_size_limit_and_filter will be better. Like that

pub(crate) fn read_partial_data_with_max_size_limit_and_task_filter<'a>(
        &'a self,
        blocks: Vec<&'a PartitionedDataBlock>,
        fetched_size_limit: i64,
        serialized_expected_task_ids_bitmap: Option<Treemap>,
    ) -> (Vec<&PartitionedDataBlock>, i64) {
        let mut fetched = vec![];
        let mut fetched_size = 0;

        for block in blocks {
            if &serialized_expected_task_ids_bitmap.is_some() {
                if !&serialized_expected_task_ids_bitmap.unwrap().contains(block.task_attempt_id as u64) {
                    continue;
                }
            }
            if fetched_size >= fetched_size_limit {
                break;
            }
            fetched_size += block.length as i64;
            fetched.push(block);
        }

        (fetched, fetched_size)
    }

[wenlongbrother]

Hi @zuston Your consideration is correct.
Integrating filter logic into read_partial_data_with_max_size_limit_and_task_filter functions is indeed a better design choice, so that the filter can be applied directly when the data is read, rather than filtering and then reading.
The advantage of this is that data can be processed more efficiently, especially when the amount of data is large or the filtering conditions are complex.
Thank you for your meticulous guidance and patience.

[zuston]

Why the bitmap will be ignored when the last_block_id = -1 ?

[wenlongbrother]

You mentioned

Otherwise, when the last_block_id = -1, the filter is not valid. This case also should be added into test.

I may misunderstood you, let me update it, sorry

[zuston]

LGTM. Thanks @wenlongbrother