This is an automated cherry-pick of pingcap#54609

Signed-off-by: ti-chi-bot <[email protected]>

pkg/planner/core/logical_limit.go

@@ -0,0 +1,179 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package core
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+
+	"github.com/pingcap/tidb/pkg/expression"
+	"github.com/pingcap/tidb/pkg/planner/core/base"
+	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	"github.com/pingcap/tidb/pkg/planner/property"
+	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
+	"github.com/pingcap/tidb/pkg/util/plancodec"
+)
+
+// LogicalLimit represents offset and limit plan.
+type LogicalLimit struct {
+	logicalop.LogicalSchemaProducer
+
+	PartitionBy      []property.SortItem // This is used for enhanced topN optimization
+	Offset           uint64
+	Count            uint64
+	PreferLimitToCop bool
+	IsPartial        bool
+}
+
+// Init initializes LogicalLimit.
+func (p LogicalLimit) Init(ctx base.PlanContext, offset int) *LogicalLimit {
+	p.BaseLogicalPlan = logicalop.NewBaseLogicalPlan(ctx, plancodec.TypeLimit, &p, offset)
+	return &p
+}
+
+// *************************** start implementation of Plan interface ***************************
+
+// ExplainInfo implements Plan interface.
+func (p *LogicalLimit) ExplainInfo() string {
+	ectx := p.SCtx().GetExprCtx().GetEvalCtx()
+	buffer := bytes.NewBufferString("")
+	if len(p.GetPartitionBy()) > 0 {
+		buffer = explainPartitionBy(ectx, buffer, p.GetPartitionBy(), false)
+		fmt.Fprintf(buffer, ", offset:%v, count:%v", p.Offset, p.Count)
+	} else {
+		fmt.Fprintf(buffer, "offset:%v, count:%v", p.Offset, p.Count)
+	}
+	return buffer.String()
+}
+
+// *************************** end implementation of Plan interface ***************************
+
+// *************************** start implementation of logicalPlan interface ***************************
+
+// HashCode implements LogicalPlan.<0th> interface.
+func (p *LogicalLimit) HashCode() []byte {
+	// PlanType + SelectOffset + Offset + Count
+	result := make([]byte, 24)
+	binary.BigEndian.PutUint32(result, uint32(plancodec.TypeStringToPhysicalID(p.TP())))
+	binary.BigEndian.PutUint32(result[4:], uint32(p.QueryBlockOffset()))
+	binary.BigEndian.PutUint64(result[8:], p.Offset)
+	binary.BigEndian.PutUint64(result[16:], p.Count)
+	return result
+}
+
+// PredicatePushDown implements base.LogicalPlan.<1st> interface.
+func (p *LogicalLimit) PredicatePushDown(predicates []expression.Expression, opt *optimizetrace.LogicalOptimizeOp) ([]expression.Expression, base.LogicalPlan) {
+	// Limit forbids any condition to push down.
+	p.BaseLogicalPlan.PredicatePushDown(nil, opt)
+	return predicates, p
+}
+
+// PruneColumns implements base.LogicalPlan.<2nd> interface.
+func (p *LogicalLimit) PruneColumns(parentUsedCols []*expression.Column, opt *optimizetrace.LogicalOptimizeOp) (base.LogicalPlan, error) {
+	savedUsedCols := make([]*expression.Column, len(parentUsedCols))
+	copy(savedUsedCols, parentUsedCols)
+	var err error
+	if p.Children()[0], err = p.Children()[0].PruneColumns(parentUsedCols, opt); err != nil {
+		return nil, err
+	}
+	p.SetSchema(nil)
+	p.InlineProjection(savedUsedCols, opt)
+	return p, nil
+}
+
+// FindBestTask inherits BaseLogicalPlan.LogicalPlan.<3rd> implementation.
+
+// BuildKeyInfo implements base.LogicalPlan.<4th> interface.
+func (p *LogicalLimit) BuildKeyInfo(selfSchema *expression.Schema, childSchema []*expression.Schema) {
+	p.LogicalSchemaProducer.BuildKeyInfo(selfSchema, childSchema)
+	if p.Count == 1 {
+		p.SetMaxOneRow(true)
+	}
+}
+
+// PushDownTopN implements the base.LogicalPlan.<5th> interface.
+func (p *LogicalLimit) PushDownTopN(topNLogicalPlan base.LogicalPlan, opt *optimizetrace.LogicalOptimizeOp) base.LogicalPlan {
+	var topN *LogicalTopN
+	if topNLogicalPlan != nil {
+		topN = topNLogicalPlan.(*LogicalTopN)
+	}
+	child := p.Children()[0].PushDownTopN(p.convertToTopN(opt), opt)
+	if topN != nil {
+		return topN.AttachChild(child, opt)
+	}
+	return child
+}
+
+// DeriveTopN inherits BaseLogicalPlan.LogicalPlan.<6th> implementation.
+
+// PredicateSimplification inherits BaseLogicalPlan.LogicalPlan.<7th> implementation.
+
+// ConstantPropagation inherits BaseLogicalPlan.LogicalPlan.<8th> implementation.
+
+// PullUpConstantPredicates inherits BaseLogicalPlan.LogicalPlan.<9th> implementation.
+
+// RecursiveDeriveStats inherits BaseLogicalPlan.LogicalPlan.<10th> implementation.
+
+// DeriveStats implement base.LogicalPlan.<11th> interface.
+func (p *LogicalLimit) DeriveStats(childStats []*property.StatsInfo, _ *expression.Schema, _ []*expression.Schema, _ [][]*expression.Column) (*property.StatsInfo, error) {
+	if p.StatsInfo() != nil {
+		return p.StatsInfo(), nil
+	}
+	p.SetStats(deriveLimitStats(childStats[0], float64(p.Count)))
+	return p.StatsInfo(), nil
+}
+
+// ExtractColGroups inherits BaseLogicalPlan.LogicalPlan.<12th> implementation.
+
+// PreparePossibleProperties inherits BaseLogicalPlan.LogicalPlan.<13th> implementation.
+
+// ExhaustPhysicalPlans implements base.LogicalPlan.<14th> interface.
+func (p *LogicalLimit) ExhaustPhysicalPlans(prop *property.PhysicalProperty) ([]base.PhysicalPlan, bool, error) {
+	return getLimitPhysicalPlans(p, prop)
+}
+
+// ExtractCorrelatedCols inherits BaseLogicalPlan.LogicalPlan.<15th> implementation.
+
+// MaxOneRow inherits BaseLogicalPlan.LogicalPlan.<16th> implementation.
+
+// Children inherits BaseLogicalPlan.LogicalPlan.<17th> implementation.
+
+// SetChildren inherits BaseLogicalPlan.LogicalPlan.<18th> implementation.
+
+// SetChild inherits BaseLogicalPlan.LogicalPlan.<19th> implementation.
+
+// RollBackTaskMap inherits BaseLogicalPlan.LogicalPlan.<20th> implementation.
+
+// CanPushToCop inherits BaseLogicalPlan.LogicalPlan.<21st> implementation.
+
+// ExtractFD inherits BaseLogicalPlan.LogicalPlan.<22nd> implementation.
+
+// GetBaseLogicalPlan inherits BaseLogicalPlan.LogicalPlan.<23rd> implementation.
+
+// ConvertOuterToInnerJoin inherits BaseLogicalPlan.LogicalPlan.<24th> implementation.
+
+// *************************** end implementation of logicalPlan interface ***************************
+
+// GetPartitionBy returns partition by fields
+func (p *LogicalLimit) GetPartitionBy() []property.SortItem {
+	return p.PartitionBy
+}
+
+func (p *LogicalLimit) convertToTopN(opt *optimizetrace.LogicalOptimizeOp) *LogicalTopN {
+	topn := LogicalTopN{Offset: p.Offset, Count: p.Count, PreferLimitToCop: p.PreferLimitToCop}.Init(p.SCtx(), p.QueryBlockOffset())
+	appendConvertTopNTraceStep(p, topn, opt)
+	return topn
+}

pkg/planner/core/operator/logicalop/logical_schema_producer.go

@@ -0,0 +1,107 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package logicalop
+
+import (
+	"github.com/pingcap/tidb/pkg/expression"
+	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
+	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace/logicaltrace"
+	"github.com/pingcap/tidb/pkg/types"
+)
+
+// LogicalSchemaProducer stores the schema for the logical plans who can produce schema directly.
+type LogicalSchemaProducer struct {
+	schema *expression.Schema
+	names  types.NameSlice
+	BaseLogicalPlan
+}
+
+// Schema implements the Plan.Schema interface.
+func (s *LogicalSchemaProducer) Schema() *expression.Schema {
+	if s.schema == nil {
+		if len(s.Children()) == 1 {
+			// default implementation for plans has only one child: proprgate child schema.
+			// multi-children plans are likely to have particular implementation.
+			s.schema = s.Children()[0].Schema().Clone()
+		} else {
+			s.schema = expression.NewSchema()
+		}
+	}
+	return s.schema
+}
+
+// OutputNames implements the Plan.OutputNames interface.
+func (s *LogicalSchemaProducer) OutputNames() types.NameSlice {
+	if s.names == nil && len(s.Children()) == 1 {
+		// default implementation for plans has only one child: proprgate child `OutputNames`.
+		// multi-children plans are likely to have particular implementation.
+		s.names = s.Children()[0].OutputNames()
+	}
+	return s.names
+}
+
+// SetOutputNames sets the output names for the plan.
+func (s *LogicalSchemaProducer) SetOutputNames(names types.NameSlice) {
+	s.names = names
+}
+
+// SetSchema sets the logical schema producer's schema.
+func (s *LogicalSchemaProducer) SetSchema(schema *expression.Schema) {
+	s.schema = schema
+}
+
+// SetSchemaAndNames sets the schema and names for the plan.
+func (s *LogicalSchemaProducer) SetSchemaAndNames(schema *expression.Schema, names types.NameSlice) {
+	s.schema = schema
+	s.names = names
+}
+
+// InlineProjection prunes unneeded columns inline an executor.
+func (s *LogicalSchemaProducer) InlineProjection(parentUsedCols []*expression.Column, opt *optimizetrace.LogicalOptimizeOp) {
+	if len(parentUsedCols) == 0 {
+		return
+	}
+	prunedColumns := make([]*expression.Column, 0)
+	used := expression.GetUsedList(s.SCtx().GetExprCtx().GetEvalCtx(), parentUsedCols, s.Schema())
+	for i := len(used) - 1; i >= 0; i-- {
+		if !used[i] {
+			prunedColumns = append(prunedColumns, s.Schema().Columns[i])
+			s.schema.Columns = append(s.Schema().Columns[:i], s.Schema().Columns[i+1:]...)
+		}
+	}
+	logicaltrace.AppendColumnPruneTraceStep(s.Self(), prunedColumns, opt)
+}
+
+// BuildKeyInfo implements LogicalPlan.BuildKeyInfo interface.
+func (s *LogicalSchemaProducer) BuildKeyInfo(selfSchema *expression.Schema, childSchema []*expression.Schema) {
+	selfSchema.Keys = nil
+	s.BaseLogicalPlan.BuildKeyInfo(selfSchema, childSchema)
+
+	// default implementation for plans has only one child: proprgate child keys
+	// multi-children plans are likely to have particular implementation.
+	if len(childSchema) == 1 {
+		for _, key := range childSchema[0].Keys {
+			indices := selfSchema.ColumnsIndices(key)
+			if indices == nil {
+				continue
+			}
+			newKey := make([]*expression.Column, 0, len(key))
+			for _, i := range indices {
+				newKey = append(newKey, selfSchema.Columns[i])
+			}
+			selfSchema.Keys = append(selfSchema.Keys, newKey)
+		}
+	}
+}