[SPARK-41513][SQL] Implement an accumulator to collect per mapper row count metrics

[amaliujia]

What changes were proposed in this pull request?

In current Spark optimizer, a single partition shuffle might be created for a limit if this limit is not the last non-action operation (e.g. a filter following the limit and the data size exceeds a threshold). There is a possibility that the previous output partitions before go into this limit are sorted. The single partition shuffle approach has a correctness bug in this case: shuffle read partitions could be out of partition order and the limit exec just take the first limit rows which could lose the order thus result into wrong result. This is a shuffle so it is relatively costly. Meanwhile, to correct this bug, a native solution is to sort all the data fed into limit again, which is another overhead.

So we propose a row count based AQE algorithm that optimizes this problem by two folds:

1. Avoid the extra sort on the shuffle read side (or with the limit exec) to achieve the correct result.
2. Avoid reading all shuffle data from mappers for this single partition shuffle to reduce shuffle cost.

Note that 1. is only applied for the sorted partition case where 2. is applied for general single partition shuffle + limit case

The algorithm works as the following:

1. Each mapper will record a row count when writing shuffle data.
2. Since this is single shuffle partition case, there is only one partition but N mappers.
3. A accumulatorV2 is implemented to collect a list of tuple which records the mapping between mapper id and the number of row written by the mapper (row count metrics)
4. AQE framework detects a plan shape of shuffle plus a global limit.
5. AQE framework reads necessary data from mappers based on the limit. For example, if mapper 1 writes 200 rows and mapper 2 writes 300 rows, and the limit is 500, AQE creates shuffle read node to write from both mapper 1 and 2, thus skip the left mappers.
6. This is both correct for limit with the sorted or non-sorted partitions.

This is the first step to implement the idea in https://issues.apache.org/jira/browse/SPARK-41512, which is to implement a row count accumulator that will be used to collect row count metrics.

Why are the changes needed?

Optimization algorithm for global limit with single partition shuffle

Does this PR introduce any user-facing change?

NO

How was this patch tested?

UT

[amaliujia]

@cloud-fan

[cloud-fan]

We can put it in the sql module, or probably the same file with shuffle node.

[amaliujia]

Moved to sql module.

[AmplabJenkins]

Can one of the admins verify this patch?

[cloud-fan]

what does this do?

[amaliujia]

hmm I can remove this assignment and use id directly.

[cloud-fan]

TaskContext.partitionId is int, do we really need long here?

[amaliujia]

The firs long here (aka IN) defines

  /**
   * Takes the inputs and accumulates.
   */
  def add(v: IN): Unit

So this should be a long for the row count?

[cloud-fan]

row count should be long, map index should be int.

[amaliujia]

I think I wrongly understand the template type.

Changed to Integer for the map index.

[cloud-fan]

assert(getOrCreate.size == 1)?

[amaliujia]

done

[cloud-fan]

when can we hit this branch?

[amaliujia]

This is because I don't know the invocation sequence of the accumulator APIs on the executor side. So I added this branch for safe.

If the setPartitionId is always called before any add, then we can does an assert on isZero and remove this branch.

What do you think?

[cloud-fan]

let's add an assert

[amaliujia]

done

[cloud-fan]

thanks, merging to master!