new tests

Author: ZhouXing19

pkg/sql/opt/memo/testdata/stats/canary-stats

@@ -0,0 +1,42 @@
+exec-ddl
+CREATE TABLE sel (a INT, b INT, c INT, d STRING, e STRING) WITH (canary_window = '30s')
+----
+
+exec-ddl
+ALTER TABLE sel INJECT STATISTICS '[
+		{
+		  "columns": ["a","b","c","d","e"],
+			"created_at": "2018-01-01 11:00:20.00000+00:00",
+			"row_count": 10,
+			"distinct_count": 10
+		},
+    {
+    	"columns": ["a","b","c","d","e"],
+      "created_at": "2018-01-01 11:00:00.00000+00:00",
+      "row_count": 20,
+      "distinct_count": 20
+    }
+	]'
+----
+
+opt stats-as-of=(2018-01-01 11:00:40.00000+00:00) canary-fraction=0.0
+SELECT a FROM sel ORDER BY a
+----
+sort
+ ├── columns: a:1(int)
+ ├── stats: [rows=20]
+ ├── ordering: +1
+ └── scan sel
+      ├── columns: a:1(int)
+      └── stats: [rows=20]
+
+opt stats-as-of=(2018-01-01 11:00:40.00000+00:00) canary-fraction=1.0
+SELECT a FROM sel ORDER BY a
+----
+sort
+ ├── columns: a:1(int)
+ ├── stats: [rows=10]
+ ├── ordering: +1
+ └── scan sel
+      ├── columns: a:1(int)
+      └── stats: [rows=10]

pkg/sql/opt/metadata.go

@@ -71,8 +71,12 @@ var canaryFraction = settings.RegisterFloatSetting(
 // should use the "canary path" for statistics.
 // This selection is atomic per query.
 func canaryRollDice(evalCtx *eval.Context) bool {
-	threshold := canaryFraction.Get(&evalCtx.Settings.SV)
-
+	var threshold float64
+	if evalCtx.TestingKnobs.CanaryFraction != nil {
+		threshold = *evalCtx.TestingKnobs.CanaryFraction
+	} else {
+		threshold = canaryFraction.Get(&evalCtx.Settings.SV)
+	}
 	// If the fraction is 0, never use canary stats. (should we even allow?)
 	if threshold == 0 {
 		return false

pkg/sql/opt/testutils/opttester/BUILD.bazel

@@ -42,6 +42,7 @@ go_library(
         "//pkg/sql/parser",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
+        "//pkg/sql/sem/asof",
         "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sem/volatility",

pkg/sql/opt/testutils/opttester/opt_tester.go

@@ -55,6 +55,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/asof"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
@@ -270,6 +271,10 @@ type Flags struct {
 	// MaxStackBytes specifies the number of bytes to limit the stack size to.
 	// If it is zero, the stack size has the default Go limit.
 	MaxStackBytes int
+
+	// CanaryFraction determines the fraction of queries that will have use
+	// canary stats rather than the stable stats.
+	CanaryFraction *float64
 }
 
 // New constructs a new instance of the OptTester for the given SQL statement.
@@ -591,6 +596,7 @@ func (ot *OptTester) RunCommand(tb testing.TB, d *datadriven.TestData) string {
 
 	ot.semaCtx.Placeholders = tree.PlaceholderInfo{}
 
+	ot.evalCtx.TestingKnobs.CanaryFraction = ot.Flags.CanaryFraction
 	ot.evalCtx.TestingKnobs.OptimizerCostPerturbation = ot.Flags.PerturbCost
 	ot.evalCtx.Locality = ot.Flags.Locality
 	ot.evalCtx.Placeholders = nil
@@ -1180,6 +1186,37 @@ func (f *Flags) Set(arg datadriven.CmdArg) error {
 		}
 		f.MaxStackBytes = int(bytes)
 
+	case "canary-fraction":
+		if len(arg.Vals) != 1 {
+			return fmt.Errorf("require one argument for canary-fraction")
+		}
+		canaryFrac, err := strconv.ParseFloat(arg.Vals[0], 64)
+		if err != nil {
+			return errors.Wrapf(err, "parse %s as float for canary fraction", arg.Vals[0])
+		}
+		f.CanaryFraction = &canaryFrac
+
+	case "stats-as-of":
+		// We set the session variable directly here rather than using the
+		// conventional "set" command because stats_as_of only implements
+		// SetWithPlanner, not Set. In production, this variable behaves
+		// like AS OF SYSTEM TIME, where timestamp evaluation through the
+		// planner is more robust. For testing purposes, we can safely
+		// bypass this requirement.
+		if len(arg.Vals) != 1 {
+			return fmt.Errorf("require one argument for stats-as-of")
+		}
+		input := arg.Vals[0]
+		if input == "null" {
+			f.evalCtx.SessionData().StatsAsOf = hlc.Timestamp{}
+		} else {
+			asOfTs, err := asof.Eval(context.Background(), tree.AsOfClause{Expr: tree.NewStrVal(input)}, &tree.SemaContext{}, &f.evalCtx)
+			if err != nil {
+				return errors.Wrapf(err, "parse %s as timestamp for stats-as-of", input)
+			}
+			f.evalCtx.SessionData().StatsAsOf = asOfTs.Timestamp
+		}
+
 	default:
 		return fmt.Errorf("unknown argument: %s", arg.Key)
 	}

pkg/sql/opt/testutils/testcat/BUILD.bazel

@@ -62,6 +62,7 @@ go_library(
         "//pkg/sql/types",
         "//pkg/sql/vecindex/vecpb",
         "//pkg/sql/vtable",
+        "//pkg/util/duration",
         "//pkg/util/intsets",
         "//pkg/util/treeprinter",
         "@com_github_cockroachdb_errors//:errors",

pkg/sql/opt/testutils/testcat/alter_table.go

@@ -47,7 +47,12 @@ func (tc *Catalog) AlterTable(stmt *tree.AlterTable) {
 			default:
 				panic(errors.AssertionFailedf("unsupported constraint type %v", d))
 			}
-
+		case *tree.AlterTableSetStorageParams:
+			if canaryWindowExpr := t.StorageParams.GetVal(canaryWindowStorageParamKey); canaryWindowExpr != nil {
+				tab.canaryWindowSize = getDurationFromStrExpr(canaryWindowExpr)
+			} else {
+				panic(errors.AssertionFailedf("unsupported AlterTableSetStorageParams stmt with storage params: %v", t.StorageParams))
+			}
 		default:
 			panic(errors.AssertionFailedf("unsupported ALTER TABLE command %T", t))
 		}

pkg/sql/opt/testutils/testcat/create_table.go

@@ -13,6 +13,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/geo/geopb"
@@ -30,7 +31,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecpb"
+	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
+	"github.com/cockroachdb/errors"
 )
 
 type indexType int
@@ -58,6 +61,8 @@ var uniqueRowIDString = "unique_rowid()"
 // referencing side of foreign keys (like it was required before 20.2).
 const createFKIndexes = false
 
+const canaryWindowStorageParamKey = "canary_window"
+
 // CreateTable creates a test table from a parsed DDL statement and adds it to
 // the catalog. This is intended for testing, and is not a complete (and
 // probably not fully correct) implementation. It just has to be "good enough".
@@ -441,6 +446,10 @@ OuterLoop:
 		}
 	}
 
+	if canaryWindowExpr := stmt.StorageParams.GetVal(canaryWindowStorageParamKey); canaryWindowExpr != nil {
+		tab.canaryWindowSize = getDurationFromStrExpr(canaryWindowExpr)
+	}
+
 	// Add the new table to the catalog.
 	tc.AddTable(tab)
 
@@ -1588,3 +1597,22 @@ func generateDefExprForSerialCol(
 			" %s", colName, tableName.String()))
 	}
 }
+
+func getDurationFromStrExpr(expr tree.Expr) time.Duration {
+	if expr == nil {
+		return 0
+	}
+	exprStr, ok := expr.(*tree.StrVal)
+	if !ok {
+		panic("expr must be of StrVal type")
+	}
+	d, err := tree.ParseDInterval(duration.IntervalStyle_POSTGRES, exprStr.RawString())
+	if err != nil {
+		panic(errors.Wrapf(err, "parsing %q", exprStr.RawString()))
+	}
+	durInSeconds, ok := d.Duration.AsInt64()
+	if !ok {
+		panic("unable to convert duration to seconds")
+	}
+	return time.Duration(durInSeconds) * time.Second
+}

pkg/sql/sem/eval/testing_knobs.go

@@ -41,6 +41,9 @@ type TestingKnobs struct {
 	// We use clusterversion.Key rather than a roachpb.Version because it will be used
 	// to get initial values to use during bootstrap.
 	TenantLogicalVersionKeyOverride clusterversion.Key
+	// CanaryFraction determines the portion of queries that would use canary
+	// stats for planning rather than using the stable stats.
+	CanaryFraction *float64
 }
 
 var _ base.ModuleTestingKnobs = &TestingKnobs{}