Optimizer: Add partition by support for derived TopN(filter on row_nu…

[ghazalfamilyusa]

Issue Number: Ref #39792

Problem Summary:

This is a continuation of #41209 which introduced deriving top N from "row_number as RN ... where RN <= N". This PR adds support for partition by in the window function "row_number as RN over (partition by X) ... where RN <= N" and restricted if the partition by is a prefix of the PK. Execution support for this extension is here tikv/tikv#14116

Design

This optimization is a logical rewrite that (1) finds the desired pattern and (2) rewrites the query by inserting a TopN node that is expected to be pushed down to storage later by the physical planner.

• Pattern
  The logical rewrite looks like for the pattern selection(filter)->window->data-source
  with the following restrictions:
  - Filter is simple condition of row_number < value or row_number <= value
  - The window function is a simple row number with default frame. If the row_number has partition by then the partition
  - by fields should be a prefix of clustered index. This is reflects the limitation in the executor that works only
  for this case.
  - Data source has no tiflash option.

• Rewrite
  Add topN above (down stream) of data source. The write change selection(filter)->window->data-source to selection(filter)->window->topN->data-source. TopN node is extended to store the partition by of row_number if any. The physical planner may change TopN to Limit if sort is met at the source. For this reason, we also extended PhysicalLimit to include partition by as well. Another change made to physical planner is to only apply TopN/Limit to TiKv since it does not help to apply it at the root. Applying TopN at the root with partition by also requires changes to executor which is not needed.

• Dependencies
  The change in this PR address only the planner enhancements. The end to end solution requires: (a) extending the interface between physical planner and execution to include the partition by fields for both Limit and TopN (b) enhancing the execution engine code for Limit and TopN to support partitioning when data is fully/partially ordered on partition key.

Scope and user interface
The change adds a session/global variable to enable/disable this optimization called "tidb_opt_derive_topn" and it is off by default to reduce the risk. The change also limits the optimization to TikV and rule not applied when TiFlash exists. This is the case since TiFllash execution does not support the extended TopN (with partition). Also, TiFlash can push down window functions and probably does not need this optimization.

Customer use case
The original user case that motivates this change is tested and work as expected in the planner. See planner/core/casetest/testdata/derive_topn_from_window_out.json in this PR.

        "SQL": "select * from (select *, row_number() over (partition by primary_key, secondary_key order by c_timestamp) as rownum from customer where primary_key = 0x002 and secondary_key >= 0x001 and c_timestamp >= 1661883508511000000) as nested where rownum <= 10 order by secondary_key desc;",
        "Plan": [
          "Sort 0.89 root  test.customer.secondary_key:desc",
          "└─Selection 0.89 root  le(Column#6, 10)",
          "  └─Window 1.11 root  row_number()->Column#6 over(partition by test.customer.primary_key, test.customer.secondary_key order by test.customer.c_timestamp rows between current row and current row)",
          "    └─Sort 1.11 root  test.customer.primary_key, test.customer.secondary_key, test.customer.c_timestamp",
          "      └─TableReader 1.11 root  data:TopN",
          "        └─TopN 1.11 cop[tikv]  partition by test.customer.primary_key, test.customer.secondary_key order by test.customer.c_timestamp, offset:0, count:10",
          "          └─Selection 1.11 cop[tikv]  ge(test.customer.c_timestamp, 1661883508511000000)",
          "            └─TableRangeScan 33.33 cop[tikv] table:customer range:[0x0002 0x0001,0x0002 +inf], keep order:false, stats:pseudo"

Tests

Unit tests

Side effects

None

Documentation

None

Release note

None

[ti-chi-bot]

[REVIEW NOTIFICATION]

This pull request has been approved by:

• chrysan
• qw4990

To complete the pull request process, please ask the reviewers in the list to review by filling /cc @reviewer in the comment.
After your PR has acquired the required number of LGTMs, you can assign this pull request to the committer in the list by filling /assign @committer in the comment to help you merge this pull request.

The full list of commands accepted by this bot can be found here.

Reviewer can indicate their review by submitting an approval review.
Reviewer can cancel approval by submitting a request changes review.

[LittleFall]

can we add partition_by_column to explain result?
like "TopN 1.11 cop[tikv] order_by:XXX desc, partition_by: XXX, offset: X, count: X"

[ghazalfamilyusa]

Good point and will do in next update.

[chrysan]

Suggested change

	// Table not clustered and window has partition by. Can not do the TopN piush down.
	// Table not clustered and window has partition by. Can not do the TopN push down.

A typo here.

[ghazalfamilyusa]

Got it and will be fixed in next update.

[chrysan]

should have some tests to cover result correctness? Rest LGTM.

[LittleFall]

Should we control this optimizer rule open or not by the blocklist, instead of a new variable?

ref: https://docs.pingcap.com/tidb/stable/blocklist-control-plan#important-optimization-rules

[ghazalfamilyusa]

I think the use case is different for system variables and blocklist. Blocklist to me seems a configuration which lasts longer. Session/global variables can be used to control a more granular level on query level as a workaround to force using or not using an optimization.

[qw4990]

Yes, actually the optimizer rule blocklist is rarely used now, so a system variable seems better here.

[qw4990]

/merge

[ti-chi-bot]

This pull request has been accepted and is ready to merge.

Commit hash: 45a7402

[qw4990]

/retest