sql/opt: add data-driven test framework for canary statistics

This commit introduces testing infrastructure for the canary statistics
feature. Previously, there was no systematic way to test the canary
statistics selection logic and the stats_as_of functionality. The
implementation adds a new data-driven test framework in
`pkg/sql/opt/memo/statistics_builder_test.go` that supports commands
for setting canary fraction, stats-as-of timestamps, and making
statistics.

The motivation for adding this test infrastructure is to unit-test the
canary full statistics rollout logic in
`statisticsBuilder.makeTableStatistics()`. While end-to-end testing can
be accomplished via `EXPLAIN ANALYZE` statements, such logictests would
also be influenced by `Builder.maybeAnnotateWithEstimates()`, which
performs its own filtering of full statistics. This dedicated test
infrastructure isolates the testing of the `makeTableStatistics` logic
to ensure the canary selection behavior can be verified independently.

Release note: None

Author: ZhouXing19

pkg/sql/opt/memo/statistics_builder_test.go

@@ -7,15 +7,20 @@ package memo
 
 import (
 	"context"
+	"strconv"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
+	"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/util/hlc"
+	"github.com/cockroachdb/datadriven"
+	"github.com/stretchr/testify/require"
 )
 
 // Most of the functionality in statistics.go is tested by data-driven
@@ -238,3 +243,54 @@ func testStats(t *testing.T, s *props.Statistics, expectedStats string) {
 		t.Fatalf("\nexpected: %s\nactual  : %s", expectedStats, actual)
 	}
 }
+
+func TestDataDriven(t *testing.T) {
+	evalCtx := eval.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
+	catalog := testcat.New()
+	var mem Memo
+
+	statsFunc := func(tab *testcat.Table, tabID opt.TableID) string {
+		t.Helper()
+		sb := &statisticsBuilder{}
+		sb.init(context.Background(), &evalCtx, &mem)
+
+		// TODO: the resStats should contain the createdAt timestamp.
+		resStats := sb.makeTableStatistics(tabID)
+		return resStats.String()
+	}
+
+	datadriven.Walk(t, "testdata/stats_builder", func(t *testing.T, path string) {
+		datadriven.RunTest(t, path, func(t *testing.T, td *datadriven.TestData) string {
+			switch td.Cmd {
+			case "exec-ddl":
+				_, err := catalog.ExecuteDDL(td.Input)
+				require.NoError(t, err)
+			case "canary-fraction":
+				require.NotEmpty(t, td.Input)
+				fraction, err := strconv.ParseFloat(td.Input, 64)
+				require.NoError(t, err)
+				evalCtx.TestingKnobs.CanaryFraction = &fraction
+			case "stats-as-of":
+				require.NotEmpty(t, td.Input)
+				if td.Input == "null" {
+					evalCtx.SessionData().StatsAsOf = hlc.Timestamp{}
+				} else {
+					asOfTs, err := asof.Eval(context.Background(), tree.AsOfClause{Expr: tree.NewStrVal(td.Input)}, &tree.SemaContext{}, &evalCtx)
+					require.NoError(t, err)
+					evalCtx.SessionData().StatsAsOf = asOfTs.Timestamp
+				}
+			case "make-stats":
+				mem.Init(context.Background(), &evalCtx)
+				require.GreaterOrEqual(t, len(td.CmdArgs), 1)
+				tblName := td.CmdArgs[0].String()
+				tn := tree.NewUnqualifiedTableName(tree.Name(tblName))
+				tab := catalog.Table(tn)
+				tabID := mem.Metadata().AddTable(tab, tn)
+				return statsFunc(tab, tabID)
+			default:
+				t.Fatalf("unknown command: %s", td.Cmd)
+			}
+			return ""
+		})
+	})
+}

pkg/sql/opt/memo/testdata/stats_builder/canary_stats

@@ -0,0 +1,91 @@
+exec-ddl
+CREATE TABLE sel (a INT, b INT, c INT, d STRING, e STRING) WITH (canary_window = '30s')
+----
+
+stats-as-of
+2018-01-01 1:00:40.00000+00:00
+----
+
+exec-ddl
+ALTER TABLE sel INJECT STATISTICS '[
+		{
+		  "columns": ["a","b","c","d","e"],
+			"created_at": "2018-01-01 1:00:20.00000+00:00",
+			"row_count": 10,
+			"distinct_count": 10
+		},
+    {
+    	"columns": ["a","b","c","d","e"],
+      "created_at": "2018-01-01 1:00:00.00000+00:00",
+      "row_count": 20,
+      "distinct_count": 20
+    }
+	]'
+----
+
+# We should pick the canary stats created at 1:00:20 since they are within the 30s canary window.
+canary-fraction
+1.0
+----
+
+make-stats sel
+----
+[rows=10]
+
+# We should pick the stable stats created at 1:00:00 which is the second most recent stats.
+canary-fraction
+0.0
+----
+
+make-stats sel
+----
+[rows=20]
+
+
+exec-ddl
+ALTER TABLE sel SET (canary_window = '10s')
+----
+
+# We should pick the stable stats created at 1:00:20 since there doesn't
+# exist any canary stats within the 10s canary window.
+canary-fraction
+0.0
+----
+
+make-stats sel
+----
+[rows=10]
+
+subtest only-one-full-stats-available
+
+exec-ddl
+ALTER TABLE sel INJECT STATISTICS '[
+		{
+		  "columns": ["a","b","c","d","e"],
+			"created_at": "2018-01-01 1:00:20.00000+00:00",
+			"row_count": 30,
+			"distinct_count": 30
+		}
+	]'
+----
+
+# For both canary and stable stats, we will pick the only one available,
+# regardless if it is within the canary window range.
+canary-fraction
+1.0
+----
+
+make-stats sel
+----
+[rows=30]
+
+# We should pick the stable stats created at 1:00:00 which is the second most recent stats.
+canary-fraction
+0.0
+----
+
+make-stats sel
+----
+[rows=30]
+
+subtest end

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/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,8 @@ 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 *float64
 }
 
 var _ base.ModuleTestingKnobs = &TestingKnobs{}