planner: The length function could not be substitute when collation of mapped column is utfxxx_bin

pkg/expression/constant_propagation.go

@@ -15,6 +15,7 @@
 package expression
 
 import (
+	exprctx "github.com/pingcap/tidb/pkg/expression/context"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
@@ -35,7 +36,7 @@ type basePropConstSolver struct {
 	eqList    []*Constant         // if eqList[i] != nil, it means col_i = eqList[i]
 	unionSet  *disjointset.IntSet // unionSet stores the relations like col_i = col_j
 	columns   []*Column           // columns stores all columns appearing in the conditions
-	ctx       BuildContext
+	ctx       exprctx.ExprContext
 }
 
 func (s *basePropConstSolver) getColID(col *Column) int {
@@ -357,8 +358,8 @@ func (s *propConstSolver) solve(conditions []Expression) []Expression {
 
 // PropagateConstant propagate constant values of deterministic predicates in a condition.
 // This is a constant propagation logic for expression list such as ['a=1', 'a=b']
-func PropagateConstant(ctx BuildContext, conditions []Expression) []Expression {
-	return newPropConstSolver().PropagateConstant(ctx, conditions)
+func PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression {
+	return newPropConstSolver().PropagateConstant(exprctx.WithConstantPropagateCheck(ctx), conditions)
 }
 
 type propOuterJoinConstSolver struct {
@@ -662,7 +663,7 @@ func (s *propOuterJoinConstSolver) solve(joinConds, filterConds []Expression) ([
 }
 
 // propagateConstantDNF find DNF item from CNF, and propagate constant inside DNF.
-func propagateConstantDNF(ctx BuildContext, conds []Expression) []Expression {
+func propagateConstantDNF(ctx exprctx.ExprContext, conds []Expression) []Expression {
 	for i, cond := range conds {
 		if dnf, ok := cond.(*ScalarFunction); ok && dnf.FuncName.L == ast.LogicOr {
 			dnfItems := SplitDNFItems(cond)
@@ -681,7 +682,7 @@ func propagateConstantDNF(ctx BuildContext, conds []Expression) []Expression {
 // Second step is to extract `outerCol = innerCol` from join conditions, and derive new join
 // conditions based on this column equal condition and `outerCol` related
 // expressions in join conditions and filter conditions;
-func PropConstOverOuterJoin(ctx BuildContext, joinConds, filterConds []Expression,
+func PropConstOverOuterJoin(ctx exprctx.ExprContext, joinConds, filterConds []Expression,
 	outerSchema, innerSchema *Schema, nullSensitive bool) ([]Expression, []Expression) {
 	solver := &propOuterJoinConstSolver{
 		outerSchema:   outerSchema,
@@ -695,7 +696,7 @@ func PropConstOverOuterJoin(ctx BuildContext, joinConds, filterConds []Expressio
 
 // PropagateConstantSolver is a constant propagate solver.
 type PropagateConstantSolver interface {
-	PropagateConstant(ctx BuildContext, conditions []Expression) []Expression
+	PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression
 }
 
 // newPropConstSolver returns a PropagateConstantSolver.
@@ -706,7 +707,7 @@ func newPropConstSolver() PropagateConstantSolver {
 }
 
 // PropagateConstant propagate constant values of deterministic predicates in a condition.
-func (s *propConstSolver) PropagateConstant(ctx BuildContext, conditions []Expression) []Expression {
+func (s *propConstSolver) PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression {
 	s.ctx = ctx
 	return s.solve(conditions)
 }

pkg/expression/context/context.go

@@ -114,6 +114,9 @@ type BuildContext interface {
 	// in most cases except for the method `isNullRejected` in planner.
 	// See the comments for `isNullRejected` in planner for more details.
 	IsInNullRejectCheck() bool
+	// IsConstantPropagateCheck returns the flag to indicate whether the expression is in constant propagate check.
+	// It should be true only when we are doing constant propagation in rule_predicate_push_down.
+	IsConstantPropagateCheck() bool
 	// ConnectionID indicates the connection ID of the current session.
 	// If the context is not in a session, it should return 0.
 	ConnectionID() uint64
@@ -145,6 +148,21 @@ func (ctx *NullRejectCheckExprContext) IsInNullRejectCheck() bool {
 	return true
 }
 
+// ConstantPropagateCheckContext is a wrapper to return true for `IsConstantPropagateCheck`.
+type ConstantPropagateCheckContext struct {
+	ExprContext
+}
+
+// WithConstantPropagateCheck returns a new `ConstantPropagateCheckContext` with the given `ExprContext`.
+func WithConstantPropagateCheck(ctx ExprContext) *ConstantPropagateCheckContext {
+	return &ConstantPropagateCheckContext{ExprContext: ctx}
+}
+
+// IsConstantPropagateCheck always returns true for `ConstantPropagateCheckContext`
+func (ctx *ConstantPropagateCheckContext) IsConstantPropagateCheck() bool {
+	return true
+}
+
 type innerOverrideEvalContext struct {
 	EvalContext
 	typeCtx types.Context

pkg/expression/contextsession/sessionctx.go

@@ -114,6 +114,11 @@ func (ctx *SessionExprContext) IsInNullRejectCheck() bool {
 	return false
 }
 
+// IsConstantPropagateCheck returns whether the ctx is in constant propagate check.
+func (ctx *SessionExprContext) IsConstantPropagateCheck() bool {
+	return false
+}
+
 // GetWindowingUseHighPrecision determines whether to compute window operations without loss of precision.
 // see https://dev.mysql.com/doc/refman/8.0/en/window-function-optimization.html for more details.
 func (ctx *SessionExprContext) GetWindowingUseHighPrecision() bool {

pkg/expression/contextstatic/exprctx.go

@@ -268,6 +268,11 @@ func (ctx *StaticExprContext) IsInNullRejectCheck() bool {
 	return false
 }
 
+// IsConstantPropagateCheck implements the `ExprContext.IsConstantPropagateCheck` and should always return false.
+func (ctx *StaticExprContext) IsConstantPropagateCheck() bool {
+	return false
+}
+
 // ConnectionID implements the `ExprContext.ConnectionID`.
 func (ctx *StaticExprContext) ConnectionID() uint64 {
 	return ctx.connectionID

pkg/expression/util.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/charset"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/opcode"
 	"github.com/pingcap/tidb/pkg/parser/terror"
@@ -467,6 +468,28 @@ func ColumnSubstituteImpl(ctx BuildContext, expr Expression, schema *Schema, new
 			}
 			return false, false, v
 		}
+		// If the collation of the column is PAD SPACE,
+		// we can't propagate the constant to the length function.
+		// For example, schema = ['name'], newExprs = ['a'], v = length(name).
+		// We can't substitute name with 'a' in length(name) because the collation of name is PAD SPACE.
+		// TODO: We will fix it here temporarily, and redesign the logic if we encounter more similar functions or situations later.
+		// Fixed issue #53730
+		if ctx.IsConstantPropagateCheck() && v.FuncName.L == ast.Length {
+			arg0, isColumn := v.GetArgs()[0].(*Column)
+			if isColumn {
+				id := schema.ColumnIndex(arg0)
+				if id != -1 {
+					_, isConstant := newExprs[id].(*Constant)
+					if isConstant {
+						mappedNewColumnCollate := schema.Columns[id].GetStaticType().GetCollate()
+						if mappedNewColumnCollate == charset.CollationUTF8MB4 ||
+							mappedNewColumnCollate == charset.CollationUTF8 {
+							return false, false, v
+						}
+					}
+				}
+			}
+		}
 		// cowExprRef is a copy-on-write util, args array allocation happens only
 		// when expr in args is changed
 		refExprArr := cowExprRef{v.GetArgs(), nil}

pkg/planner/core/casetest/rule/BUILD.bazel

@@ -9,10 +9,11 @@ go_test(
         "rule_inject_extra_projection_test.go",
         "rule_join_reorder_test.go",
         "rule_outer2inner_test.go",
+        "rule_predicate_pushdown_test.go",
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 6,
+    shard_count = 7,
     deps = [
         "//pkg/domain",
         "//pkg/expression",

pkg/planner/core/casetest/rule/main_test.go

@@ -32,6 +32,7 @@ func TestMain(m *testing.M) {
 	testDataMap.LoadTestSuiteData("testdata", "outer2inner")
 	testDataMap.LoadTestSuiteData("testdata", "derive_topn_from_window")
 	testDataMap.LoadTestSuiteData("testdata", "join_reorder_suite")
+	testDataMap.LoadTestSuiteData("testdata", "predicate_pushdown_suite")
 	opts := []goleak.Option{
 		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
 		goleak.IgnoreTopFunction("github.com/bazelbuild/rules_go/go/tools/bzltestutil.RegisterTimeoutHandler.func1"),
@@ -61,3 +62,7 @@ func GetDerivedTopNSuiteData() testdata.TestData {
 func GetJoinReorderSuiteData() testdata.TestData {
 	return testDataMap["join_reorder_suite"]
 }
+
+func GetPredicatePushdownSuiteData() testdata.TestData {
+	return testDataMap["predicate_pushdown_suite"]
+}

pkg/planner/core/casetest/rule/rule_predicate_pushdown_test.go

@@ -0,0 +1,59 @@
+// 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 rule
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testdata"
+	"github.com/stretchr/testify/require"
+)
+
+func runPredicatePushdownTestData(t *testing.T, tk *testkit.TestKit, name string) {
+	var input []string
+	var output []struct {
+		SQL     string
+		Plan    []string
+		Warning []string
+	}
+	predicatePushdownSuiteData := GetPredicatePushdownSuiteData()
+	predicatePushdownSuiteData.LoadTestCasesByName(name, t, &input, &output)
+	require.Equal(t, len(input), len(output))
+	for i := range input {
+		if strings.Contains(input[i], "set") {
+			tk.MustExec(input[i])
+			continue
+		}
+		testdata.OnRecord(func() {
+			output[i].SQL = input[i]
+			output[i].Plan = testdata.ConvertRowsToStrings(tk.MustQuery("explain format = 'brief' " + input[i]).Rows())
+			output[i].Warning = testdata.ConvertRowsToStrings(tk.MustQuery("show warnings").Rows())
+		})
+		tk.MustQuery("explain format = 'brief' " + input[i]).Check(testkit.Rows(output[i].Plan...))
+		tk.MustQuery("show warnings").Check(testkit.Rows(output[i].Warning...))
+	}
+}
+
+func TestConstantPropagateWithCollation(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	// create table
+	tk.MustExec("use test")
+	tk.MustExec("create table t (id int primary key, name varchar(20));")
+
+	runPredicatePushdownTestData(t, tk, "TestConstantPropagateWithCollation")
+}

pkg/planner/core/casetest/rule/testdata/predicate_pushdown_suite_in.json

@@ -0,0 +1,12 @@
+[
+  {
+    "name": "TestConstantPropagateWithCollation",
+    "cases": [
+      "set names utf8",
+      "select * from t where name='a' and length(name)=1; -- without constant propagated",
+      "set names utf8mb4",
+      "select * from t where name='a' and length(name)=1; -- without constant propagated",
+      "select * from (select 'test' as b from t) n where length(b) > 2; -- can be substituted"
+    ]
+  }
+]

pkg/planner/core/casetest/rule/testdata/predicate_pushdown_suite_out.json

@@ -0,0 +1,44 @@
+[
+  {
+    "Name": "TestConstantPropagateWithCollation",
+    "Cases": [
+      {
+        "SQL": "",
+        "Plan": null,
+        "Warning": null
+      },
+      {
+        "SQL": "select * from t where name='a' and length(name)=1; -- without constant propagated",
+        "Plan": [
+          "TableReader 8.00 root  data:Selection",
+          "└─Selection 8.00 cop[tikv]  eq(length(test.t.name), 1), eq(test.t.name, \"a\")",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo"
+        ],
+        "Warning": null
+      },
+      {
+        "SQL": "",
+        "Plan": null,
+        "Warning": null
+      },
+      {
+        "SQL": "select * from t where name='a' and length(name)=1; -- without constant propagated",
+        "Plan": [
+          "TableReader 8.00 root  data:Selection",
+          "└─Selection 8.00 cop[tikv]  eq(length(test.t.name), 1), eq(test.t.name, \"a\")",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo"
+        ],
+        "Warning": null
+      },
+      {
+        "SQL": "select * from (select 'test' as b from t) n where length(b) > 2; -- can be substituted",
+        "Plan": [
+          "Projection 10000.00 root  test->Column#3",
+          "└─TableReader 10000.00 root  data:TableFullScan",
+          "  └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo"
+        ],
+        "Warning": null
+      }
+    ]
+  }
+]