planner: fix null-reject proof for WEEK and YEARWEEK (#68817)

close #68816

Author: winoros

pkg/planner/util/null_misc.go

@@ -217,6 +217,14 @@ func proveNullRejectedScalarFunc(
 		return proveNullRejectedIn(ctx, innerSchema, expr, allowNullifiedFold)
 	case ast.IsNull:
 		return nullRejectProof{}
+	case ast.Week, ast.YearWeek:
+		// Only the date argument is NULL-preserving. A NULL mode argument is
+		// treated as mode 0 by MySQL/TiDB, so these functions cannot be listed
+		// in nullRejectNullPreservingFunctions.
+		if proveNullRejected(ctx, innerSchema, expr.GetArgs()[0], allowNullifiedFold).mustNull {
+			return nullRejectProof{nonTrue: true, mustNull: true}
+		}
+		return nullRejectProof{}
 	}
 
 	if mode, ok := nullRejectRejectNullTests[expr.FuncName.L]; ok {

pkg/planner/util/null_misc_builtins.go

@@ -32,7 +32,6 @@ var nullRejectNullPreservingFunctions = map[string]struct{}{
 	// Cast and unary operators.
 	ast.Cast:       {},
 	ast.UnaryNot:   {},
-	ast.UnaryPlus:  {},
 	ast.UnaryMinus: {},
 	ast.BitNeg:     {},
 
@@ -129,7 +128,6 @@ var nullRejectNullPreservingFunctions = map[string]struct{}{
 	ast.Right:           {},
 	ast.Rpad:            {},
 	ast.RTrim:           {},
-	ast.Soundex:         {},
 	ast.Space:           {},
 	ast.Substr:          {},
 	ast.Substring:       {},
@@ -185,11 +183,9 @@ var nullRejectNullPreservingFunctions = map[string]struct{}{
 	ast.ToDays:        {},
 	ast.ToSeconds:     {},
 	ast.UnixTimestamp: {},
-	ast.Week:          {},
 	ast.Weekday:       {},
 	ast.WeekOfYear:    {},
 	ast.Year:          {},
-	ast.YearWeek:      {},
 
 	// Encryption, hashing, and compression.
 	ast.Compress:           {},

pkg/planner/util/null_misc_test.go

@@ -40,6 +40,15 @@ func TestNullRejectBuiltinRegistrySnapshot(t *testing.T) {
 	require.NotEmpty(t, names)
 	require.Equal(t, "729f5252bcd91efe1a4bbf0c383a36c5a2e52ed2d90d7aab0a3e0b450322294c", hex.EncodeToString(sum[:]))
 
+	internalScalarNames := map[string]struct{}{
+		ast.Cast: {},
+	}
+	for name := range nullRejectNullPreservingFunctions {
+		if _, ok := internalScalarNames[name]; ok {
+			continue
+		}
+		require.Contains(t, names, name)
+	}
 	for name := range nullRejectRejectNullTests {
 		require.Contains(t, names, name)
 	}
@@ -55,7 +64,8 @@ func TestIsNullRejectedProofModes(t *testing.T) {
 	outerC := newNullRejectIntColumn(3)
 	innerS := newNullRejectStringColumn(4)
 	innerUnsignedD := newNullRejectUintColumn(5)
-	innerSchema := expression.NewSchema(innerA, innerB, innerS, innerUnsignedD)
+	innerDate := newNullRejectDatetimeColumn(6)
+	innerSchema := expression.NewSchema(innerA, innerB, innerS, innerUnsignedD, innerDate)
 
 	gtInnerAZero := newNullRejectFunc(t, exprCtx, ast.GT, types.NewFieldType(mysql.TypeTiny), innerA, expression.NewZero())
 	eqInnerAZero := newNullRejectFunc(t, exprCtx, ast.EQ, types.NewFieldType(mysql.TypeTiny), innerA, expression.NewZero())
@@ -181,6 +191,38 @@ func TestIsNullRejectedProofModes(t *testing.T) {
 		newNullRejectStringConst("abc"),
 		innerS,
 	)
+	weekWithNullableMode := newNullRejectFunc(
+		t,
+		exprCtx,
+		ast.Week,
+		types.NewFieldType(mysql.TypeLonglong),
+		newNullRejectStringConst("2024-01-08"),
+		innerA,
+	)
+	yearWeekWithNullableMode := newNullRejectFunc(
+		t,
+		exprCtx,
+		ast.YearWeek,
+		types.NewFieldType(mysql.TypeLonglong),
+		newNullRejectStringConst("2024-01-08"),
+		innerA,
+	)
+	weekWithNullableDateAndOuterMode := newNullRejectFunc(
+		t,
+		exprCtx,
+		ast.Week,
+		types.NewFieldType(mysql.TypeLonglong),
+		innerDate,
+		outerC,
+	)
+	yearWeekWithNullableDateAndOuterMode := newNullRejectFunc(
+		t,
+		exprCtx,
+		ast.YearWeek,
+		types.NewFieldType(mysql.TypeLonglong),
+		innerDate,
+		outerC,
+	)
 	deferredInnerGTZero := newNullRejectDeferredConst(exprCtx, gtInnerAZero)
 	deferredCoalesceInnerATwoGTTwo := newNullRejectDeferredConst(exprCtx,
 		newNullRejectFunc(t, exprCtx, ast.GT, types.NewFieldType(mysql.TypeTiny), coalesceInnerATwo, newNullRejectIntConst(2)),
@@ -342,6 +384,54 @@ func TestIsNullRejectedProofModes(t *testing.T) {
 			expr:     newNullRejectNotNull(t, exprCtx, jsonSearchNullableEscape),
 			expected: false,
 		},
+		{
+			name: "week_nullable_mode_uses_default_mode_zero",
+			expr: newNullRejectFunc(
+				t,
+				exprCtx,
+				ast.GE,
+				types.NewFieldType(mysql.TypeTiny),
+				weekWithNullableMode,
+				expression.NewZero(),
+			),
+			expected: false,
+		},
+		{
+			name: "yearweek_nullable_mode_uses_default_mode_zero",
+			expr: newNullRejectFunc(
+				t,
+				exprCtx,
+				ast.GE,
+				types.NewFieldType(mysql.TypeTiny),
+				yearWeekWithNullableMode,
+				expression.NewZero(),
+			),
+			expected: false,
+		},
+		{
+			name: "week_nullable_date_rejects_null_even_with_outer_mode",
+			expr: newNullRejectFunc(
+				t,
+				exprCtx,
+				ast.GE,
+				types.NewFieldType(mysql.TypeTiny),
+				weekWithNullableDateAndOuterMode,
+				expression.NewZero(),
+			),
+			expected: true,
+		},
+		{
+			name: "yearweek_nullable_date_rejects_null_even_with_outer_mode",
+			expr: newNullRejectFunc(
+				t,
+				exprCtx,
+				ast.GE,
+				types.NewFieldType(mysql.TypeTiny),
+				yearWeekWithNullableDateAndOuterMode,
+				expression.NewZero(),
+			),
+			expected: true,
+		},
 		{
 			name:     "deferred_expr_uses_symbolic_null_reject_proof",
 			expr:     deferredInnerGTZero,
@@ -393,6 +483,15 @@ func newNullRejectUintColumn(id int64) *expression.Column {
 	}
 }
 
+func newNullRejectDatetimeColumn(id int64) *expression.Column {
+	return &expression.Column{
+		UniqueID: id,
+		ID:       id,
+		Index:    int(id),
+		RetType:  types.NewFieldType(mysql.TypeDatetime),
+	}
+}
+
 func newNullRejectStringConst(value string) *expression.Constant {
 	return &expression.Constant{
 		Value:   types.NewStringDatum(value),