sql: implement stats_as_of session variable for time-based statistics selection

This commit adds a new session variable `stats_as_of` that allows
controlling statistics selection based on a specific timestamp rather
than the current time. Previously, statistics selection was always
relative to the current wall clock time, making it difficult to get
consistent query plans for historical analysis or testing.

This feature is only for debugging and troubleshooting, and should not
be used in production.

The implementation is also integrated into the existing canary
statistics logic to respect the as-of timestamp when determining canary
window boundaries.

Release note (sql change): adds a new session variable `stats_as_of`
that allows controlling statistics selection based on a specific
timestamp rather than the current time.

Author: ZhouXing19

pkg/sql/opt/memo/BUILD.bazel

@@ -52,6 +52,7 @@ go_library(
         "//pkg/util/buildutil",
         "//pkg/util/duration",
         "//pkg/util/encoding",
+        "//pkg/util/hlc",
         "//pkg/util/intsets",
         "//pkg/util/iterutil",
         "//pkg/util/json",

pkg/sql/opt/memo/statistics_builder.go

@@ -686,29 +686,40 @@ func (sb *statisticsBuilder) makeTableStatistics(tabID opt.TableID) *props.Stati
 	useCanary := sb.mem.Metadata().UseCanaryStats()
 	var skippedStatsCreationTimestamp time.Time
 	canaryWindowSize := tabMeta.CanaryWindowSize
+	// TODO(janexing): should we use clock.Now() or the StmtTimestamp as
+	// the default? Or avoid this issue by only setting it if the session var
+	// is set?
+	asOfTs := hlc.Timestamp{WallTime: sb.evalCtx.StmtTimestamp.UnixNano()}
+	if asOf := sb.evalCtx.SessionData().StatsAsOf; !asOf.IsEmpty() {
+		asOfTs = asOf
+	}
 	// Find the most recent full statistic. (Stats are ordered with most recent first.)
 	var first int
 	sd := sb.evalCtx.SessionData()
 	for first < tab.StatisticCount() {
 		stat := tab.Statistic(first)
+		createdAtTS := hlc.Timestamp{WallTime: stat.CreatedAt().UnixNano()}
 		if stat.IsPartial() ||
 			stat.IsMerged() && !sd.OptimizerUseMergedPartialStatistics ||
 			stat.IsForecast() && !sd.OptimizerUseForecasts {
 			first++
 			continue
+		} else if createdAtTS.After(asOfTs) {
+			// The stats is too new, skip it.
+			first++
+			continue
 		} else if canaryWindowSize > 0 && !useCanary && first < tab.StatisticCount()-1 {
 			// The following, we are getting full statistics for the stable stats,
 			// in contrast to the canary stats. The canary stats is skipped.
 			// If there remains only one full statistics, don't skip.
-			createdAtTS := hlc.Timestamp{WallTime: stat.CreatedAt().UnixNano()}
 			if stat.CreatedAt() == skippedStatsCreationTimestamp && !skippedStatsCreationTimestamp.IsZero() {
 				// We've already seen this canary stat, so skip it.
 				first++
 				continue
 			}
 			// Found a canary stats (defined as creation time within the canary window size). Register the
 			// creation timestamp and move on the next older one.
-			if createdAtTS.AddDuration(canaryWindowSize).After(hlc.Timestamp{WallTime: time.Now().UnixNano()}) {
+			if createdAtTS.AddDuration(canaryWindowSize).After(asOfTs) {
 				// If there is already a canary stats skipped, we don't skip again.
 				if skippedStatsCreationTimestamp.IsZero() {
 					skippedStatsCreationTimestamp = stat.CreatedAt()

pkg/sql/sessiondatapb/local_only_session_data.proto

@@ -751,6 +751,10 @@ message LocalOnlySessionData {
   // clamp the selectivity of open-ended inequality filters (e.g. <, >, !=)
   // but not (=, BETWEEN etc.) to a minimum threshold.
   bool optimizer_clamp_inequality_selectivity = 190;
+  // StatsAsOf is an optional timestamp that, when set, stats selection
+  // will be based on the age gap between the stats creation timestamp and
+  // the StatsAsOf timestamp, rather than the current time.
+  util.hlc.Timestamp stats_as_of = 191 [(gogoproto.nullable) = false];
 
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //

pkg/sql/sessionmutator/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
         "//pkg/sql/sessiondata",
         "//pkg/sql/sessiondatapb",
         "//pkg/util/duration",
+        "//pkg/util/hlc",
         "//pkg/util/timeutil",
         "//pkg/util/timeutil/pgdate",
         "@com_github_cockroachdb_redact//:redact",

pkg/sql/sessionmutator/mutator.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
 	"github.com/cockroachdb/redact"
@@ -1102,3 +1103,7 @@ func (m *SessionDataMutator) SetOptimizerClampLowHistogramSelectivity(val bool)
 func (m *SessionDataMutator) SetOptimizerClampInequalitySelectivity(val bool) {
 	m.Data.OptimizerClampInequalitySelectivity = val
 }
+
+func (m *SessionDataMutator) SetStatsAsOf(val hlc.Timestamp) {
+	m.Data.StatsAsOf = val
+}

pkg/sql/vars.go

@@ -4444,6 +4444,38 @@ var varGen = map[string]sessionVar{
 		},
 		GlobalDefault: globalFalse,
 	},
+	`stats_as_of`: {
+		GetStringVal: makePostgresBoolGetStringValFn(`stats_as_of`),
+		SetWithPlanner: func(ctx context.Context, p *planner, local bool, s string) error {
+			ts := hlc.Timestamp{}
+			if s != "" {
+				asOfTimestamp, err := p.EvalAsOfTimestamp(ctx, tree.AsOfClause{Expr: tree.NewStrVal(s)})
+				if err != nil {
+					return errors.Wrap(err, "could not parse stats_as_of")
+				}
+				ts = asOfTimestamp.Timestamp
+			}
+
+			return p.applyOnSessionDataMutators(
+				ctx,
+				local,
+				func(m sessionmutator.SessionDataMutator) error {
+					m.SetStatsAsOf(ts)
+					return nil
+				},
+			)
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			asOfTs := evalCtx.SessionData().StatsAsOf
+			if asOfTs.IsEmpty() {
+				return "", nil
+			}
+			return asOfTs.AsOfSystemTime(), nil
+		},
+		GlobalDefault: func(sv *settings.Values) string {
+			return ""
+		},
+	},
 }
 
 func ReplicationModeFromString(s string) (sessiondatapb.ReplicationMode, error) {