introduce expr/constant into fdSet with uniqueID and refine some op's logic

expression/util.go

@@ -166,6 +166,161 @@ func extractColumns(result []*Column, expr Expression, filter func(*Column) bool
 	return result
 }
 
+// ExtractEquivalenceColumns detects the equivalence from CNF exprs.
+func ExtractEquivalenceColumns(result [][]Expression, exprs []Expression) [][]Expression {
+	// exprs are CNF expressions, EQ condition only make sense in the top level of every expr.
+	for _, expr := range exprs {
+		result = extractEquivalenceColumns(result, expr)
+	}
+	return result
+}
+
+func extractEquivalenceColumns(result [][]Expression, expr Expression) [][]Expression {
+	switch v := expr.(type) {
+	case *ScalarFunction:
+		// a==b, a<=>b, the latter one is evaluated to true when a,b are both null.
+		if v.FuncName.L == ast.EQ || v.FuncName.L == ast.NullEQ {
+			args := v.GetArgs()
+			if len(args) == 2 {
+				col1, ok1 := args[0].(*Column)
+				col2, ok2 := args[1].(*Column)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col1, col2})
+				}
+				col, ok1 := args[0].(*Column)
+				scl, ok2 := args[1].(*ScalarFunction)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col, scl})
+				}
+				col, ok1 = args[1].(*Column)
+				scl, ok2 = args[0].(*ScalarFunction)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col, scl})
+				}
+			}
+			return result
+		}
+		if v.FuncName.L == ast.In {
+			args := v.GetArgs()
+			// only `col in (only 1 element)`, can we build an equivalence here.
+			if len(args[1:]) == 1 {
+				col1, ok1 := args[0].(*Column)
+				col2, ok2 := args[1].(*Column)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col1, col2})
+				}
+				col, ok1 := args[0].(*Column)
+				scl, ok2 := args[1].(*ScalarFunction)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col, scl})
+				}
+				col, ok1 = args[1].(*Column)
+				scl, ok2 = args[0].(*ScalarFunction)
+				if ok1 && ok2 {
+					result = append(result, []Expression{col, scl})
+				}
+			}
+			return result
+		}
+		// For Non-EQ function, we don't have to traverse down.
+		// eg: (a=b or c=d) doesn't make any definitely equivalence assertion.
+	}
+	return result
+}
+
+// ExtractConstantEqColumnsOrScalar detects the constant equal relationship from CNF exprs.
+func ExtractConstantEqColumnsOrScalar(ctx sessionctx.Context, result []Expression, exprs []Expression) []Expression {
+	// exprs are CNF expressions, EQ condition only make sense in the top level of every expr.
+	for _, expr := range exprs {
+		result = extractConstantEqColumnsOrScalar(ctx, result, expr)
+	}
+	return result
+}
+
+func extractConstantEqColumnsOrScalar(ctx sessionctx.Context, result []Expression, expr Expression) []Expression {
+	switch v := expr.(type) {
+	case *ScalarFunction:
+		if v.FuncName.L == ast.EQ || v.FuncName.L == ast.NullEQ {
+			args := v.GetArgs()
+			if len(args) == 2 {
+				col, ok1 := args[0].(*Column)
+				_, ok2 := args[1].(*Constant)
+				if ok1 && ok2 {
+					result = append(result, col)
+				}
+				col, ok1 = args[1].(*Column)
+				_, ok2 = args[0].(*Constant)
+				if ok1 && ok2 {
+					result = append(result, col)
+				}
+				// take the correlated column as constant here.
+				col, ok1 = args[0].(*Column)
+				_, ok2 = args[1].(*CorrelatedColumn)
+				if ok1 && ok2 {
+					result = append(result, col)
+				}
+				col, ok1 = args[1].(*Column)
+				_, ok2 = args[0].(*CorrelatedColumn)
+				if ok1 && ok2 {
+					result = append(result, col)
+				}
+				scl, ok1 := args[0].(*ScalarFunction)
+				_, ok2 = args[1].(*Constant)
+				if ok1 && ok2 {
+					result = append(result, scl)
+				}
+				scl, ok1 = args[1].(*ScalarFunction)
+				_, ok2 = args[0].(*Constant)
+				if ok1 && ok2 {
+					result = append(result, scl)
+				}
+				// take the correlated column as constant here.
+				scl, ok1 = args[0].(*ScalarFunction)
+				_, ok2 = args[1].(*CorrelatedColumn)
+				if ok1 && ok2 {
+					result = append(result, scl)
+				}
+				scl, ok1 = args[1].(*ScalarFunction)
+				_, ok2 = args[0].(*CorrelatedColumn)
+				if ok1 && ok2 {
+					result = append(result, scl)
+				}
+			}
+			return result
+		}
+		if v.FuncName.L == ast.In {
+			args := v.GetArgs()
+			allArgsIsConst := true
+			// only `col in (all same const)`, can col be the constant column.
+			// eg: a in (1, "1") does, while a in (1, '2') doesn't.
+			guard := args[1]
+			for i, v := range args[1:] {
+				if _, ok := v.(*Constant); !ok {
+					allArgsIsConst = false
+					break
+				}
+				if i == 0 {
+					continue
+				}
+				if !guard.Equal(ctx, v) {
+					allArgsIsConst = false
+					break
+				}
+			}
+			if allArgsIsConst {
+				if col, ok := args[0].(*Column); ok {
+					result = append(result, col)
+				} else if scl, ok := args[0].(*ScalarFunction); ok {
+					result = append(result, scl)
+				}
+			}
+			return result
+		}
+		// For Non-EQ function, we don't have to traverse down.
+	}
+	return result
+}
+
 // ExtractColumnSet extracts the different values of `UniqueId` for columns in expressions.
 func ExtractColumnSet(exprs []Expression) *intsets.Sparse {
 	set := &intsets.Sparse{}

planner/core/logical_plan_builder.go

@@ -1158,10 +1158,15 @@ func (b *PlanBuilder) buildProjectionField(ctx context.Context, p LogicalPlan, f
 	if expr == nil {
 		return nil, name, nil
 	}
+	// for expr projection, we should record the map relationship <hashcode, uniqueID> down.
 	newCol := &expression.Column{
 		UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(),
 		RetType:  expr.GetType(),
 	}
+	if b.ctx.GetSessionVars().MapHashCode2UniqueID4ExtendedCol == nil {
+		b.ctx.GetSessionVars().MapHashCode2UniqueID4ExtendedCol = make(map[string]int, 1)
+	}
+	b.ctx.GetSessionVars().MapHashCode2UniqueID4ExtendedCol[string(expr.HashCode(b.ctx.GetSessionVars().StmtCtx))] = int(newCol.UniqueID)
 	newCol.SetCoercibility(expr.Coercibility())
 	return newCol, name, nil
 }
@@ -2732,6 +2737,7 @@ func checkColFuncDepend(
 			continue
 		}
 		funcDepend := true
+		// if all columns of some unique/pri indexes are determined, all columns left are check-passed.
 		for _, indexCol := range index.Columns {
 			iColInfo := tblInfo.Columns[indexCol.Offset]
 			if !mysql.HasNotNullFlag(iColInfo.Flag) {
@@ -4245,16 +4251,29 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as
 	return result, nil
 }
 
-func (ds *DataSource) extractFD() *fd.FDSet {
+func (ds *DataSource) ExtractFD() *fd.FDSet {
 	// FD in datasource (leaf node) can be cached and reused.
-	if ds.fdSet == nil {
-		fds := &fd.FDSet{}
+	// Once the all conditions are not equal to nil, built it again.
+	if ds.fdSet == nil || ds.allConds != nil {
+		fds := &fd.FDSet{HashCodeToUniqueID: make(map[string]int)}
 		allCols := fd.NewFastIntSet()
 		// should use the column's unique ID avoiding fdSet conflict.
 		for _, col := range ds.TblCols {
 			// todo: change it to int64
 			allCols.Insert(int(col.UniqueID))
 		}
+		// int pk doesn't store its index column in indexInfo.
+		if ds.tableInfo.PKIsHandle {
+			keyCols := fd.NewFastIntSet()
+			for _, col := range ds.TblCols {
+				if mysql.HasPriKeyFlag(col.RetType.Flag) {
+					keyCols.Insert(int(col.UniqueID))
+				}
+			}
+			fds.AddStrictFunctionalDependency(keyCols, allCols)
+			fds.MakeNotNull(keyCols)
+		}
+		// other indices including common handle.
 		for _, idx := range ds.tableInfo.Indices {
 			keyCols := fd.NewFastIntSet()
 			allColIsNotNull := true
@@ -4282,6 +4301,24 @@ func (ds *DataSource) extractFD() *fd.FDSet {
 				fds.AddLaxFunctionalDependency(keyCols, allCols)
 			}
 		}
+		// handle the datasource conditions (maybe pushed down from upper layer OP)
+		if ds.allConds != nil {
+			// extract the not null attributes from selection conditions.
+			notnullColsUniqueIDs := extractNotNullFromConds(ds.allConds, ds)
+
+			// extract the constant cols from selection conditions.
+			constUniqueIDs := extractConstantCols(ds.allConds, ds, fds)
+
+			// extract equivalence cols.
+			equivUniqueIDs := extractEquivalenceCols(ds.allConds, ds, fds)
+
+			// apply conditions to FD.
+			fds.MakeNotNull(notnullColsUniqueIDs)
+			fds.AddConstants(constUniqueIDs)
+			for _, equiv := range equivUniqueIDs {
+				fds.AddEquivalence(equiv[0], equiv[1])
+			}
+		}
 		ds.fdSet = fds
 	}
 	return ds.fdSet