Merge pull request #112323 from yuzefovich/backport23.1-111303-112263

cockroachdb · Oct 13, 2023 · 539182c · 539182c
2 parents 72e5b81 + a958045
commit 539182c
Show file tree

Hide file tree

Showing 17 changed files with 854 additions and 461 deletions.
diff --git a/pkg/sql/BUILD.bazel b/pkg/sql/BUILD.bazel
@@ -829,6 +829,7 @@ go_test(
         "//pkg/sql/sqlliveness/sqllivenesstestutils",
         "//pkg/sql/sqlstats",
         "//pkg/sql/sqlstats/persistedsqlstats",
+        "//pkg/sql/sqlstats/persistedsqlstats/sqlstatsutil",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/stats",
         "//pkg/sql/stmtdiagnostics",

diff --git a/pkg/sql/appstatspb/app_stats.go b/pkg/sql/appstatspb/app_stats.go
@@ -138,6 +138,36 @@ func (t *TransactionStatistics) Add(other *TransactionStatistics) {
 	t.Count += other.Count
 }
 
+// Add combines CollectedStatementStatistics into a single AggregatedStatementMetadata.
+func (s *AggregatedStatementMetadata) Add(other *CollectedStatementStatistics) {
+	// Only set the value if it hasn't already been set.
+	if s.Query == "" || s.QuerySummary == "" {
+		s.ImplicitTxn = other.Key.ImplicitTxn
+		s.Query = other.Key.Query
+		s.QuerySummary = other.Key.QuerySummary
+		s.StmtType = other.Stats.SQLType
+	}
+
+	// Avoid creating the array if the db names match.
+	if len(s.Databases) != 1 || s.Databases[0] != other.Key.Database {
+		s.Databases = util.CombineUniqueString(s.Databases, []string{other.Key.Database})
+	}
+
+	if other.Key.DistSQL {
+		s.DistSQLCount++
+	}
+	if other.Key.Failed {
+		s.FailedCount++
+	}
+	if other.Key.FullScan {
+		s.FullScanCount++
+	}
+	if other.Key.Vec {
+		s.VecCount++
+	}
+	s.TotalCount++
+}
+
 // Add combines other into this StatementStatistics.
 func (s *StatementStatistics) Add(other *StatementStatistics) {
 	s.FirstAttemptCount += other.FirstAttemptCount

diff --git a/pkg/sql/distsql/columnar_operators_test.go b/pkg/sql/distsql/columnar_operators_test.go
@@ -105,6 +105,9 @@ var aggregateFuncToNumArguments = map[execinfrapb.AggregatorSpec_Func]int{
 	execinfrapb.FinalCorr:               1,
 	execinfrapb.FinalSqrdiff:            3,
 	execinfrapb.ArrayCatAgg:             1,
+	execinfrapb.MergeStatsMetadata:      1,
+	execinfrapb.MergeStatementStats:     1,
+	execinfrapb.MergeTransactionStats:   1,
 }
 
 // TestAggregateFuncToNumArguments ensures that all aggregate functions are
@@ -183,6 +186,11 @@ func TestAggregatorAgainstProcessor(t *testing.T) {
 				execinfrapb.PercentileContImpl:
 				// We skip percentile functions because those can only be
 				// planned as window functions.
+			case execinfrapb.MergeStatsMetadata,
+				execinfrapb.MergeStatementStats,
+				execinfrapb.MergeTransactionStats:
+				// We skip merge statistics functions because they
+				// require custom JSON objects.
 			default:
 				found = true
 			}

diff --git a/pkg/sql/execinfrapb/aggregate_funcs.go b/pkg/sql/execinfrapb/aggregate_funcs.go
@@ -74,4 +74,7 @@ const (
 	FinalCorr               = AggregatorSpec_FINAL_CORR
 	FinalSqrdiff            = AggregatorSpec_FINAL_SQRDIFF
 	ArrayCatAgg             = AggregatorSpec_ARRAY_CAT_AGG
+	MergeStatsMetadata      = AggregatorSpec_MERGE_STATS_METADATA
+	MergeStatementStats     = AggregatorSpec_MERGE_STATEMENT_STATS
+	MergeTransactionStats   = AggregatorSpec_MERGE_TRANSACTION_STATS
 )
diff --git a/pkg/sql/execinfrapb/processors_sql.proto b/pkg/sql/execinfrapb/processors_sql.proto
@@ -864,6 +864,9 @@ message AggregatorSpec {
     FINAL_CORR = 59;
     FINAL_SQRDIFF = 60;
     ARRAY_CAT_AGG = 61;
+    MERGE_STATS_METADATA = 62;
+    MERGE_STATEMENT_STATS = 63;
+    MERGE_TRANSACTION_STATS = 64;
   }
 
   enum Type {

diff --git a/pkg/sql/opt/exec/execbuilder/testdata/observability b/pkg/sql/opt/exec/execbuilder/testdata/observability
diff --git a/pkg/sql/opt/memo/typing.go b/pkg/sql/opt/memo/typing.go
@@ -208,6 +208,10 @@ func init() {
 	typingFuncMap[opt.LeadOp] = typeAsFirstArg
 	typingFuncMap[opt.NthValueOp] = typeAsFirstArg
 
+	typingFuncMap[opt.MergeStatsMetadataOp] = typeAsFirstArg
+	typingFuncMap[opt.MergeStatementStatsOp] = typeAsFirstArg
+	typingFuncMap[opt.MergeTransactionStatsOp] = typeAsFirstArg
+
 	// Modifiers for aggregations pass through their argument.
 	typingFuncMap[opt.AggDistinctOp] = typeAsFirstArg
 	typingFuncMap[opt.AggFilterOp] = typeAsFirstArg

diff --git a/pkg/sql/opt/operator.go b/pkg/sql/opt/operator.go
@@ -183,52 +183,55 @@ var UnaryOpReverseMap = map[Operator]tree.UnaryOperatorSymbol{
 // AggregateOpReverseMap maps from an optimizer operator type to the name of an
 // aggregation function.
 var AggregateOpReverseMap = map[Operator]string{
-	ArrayAggOp:            "array_agg",
-	ArrayCatAggOp:         "array_cat_agg",
-	AvgOp:                 "avg",
-	BitAndAggOp:           "bit_and",
-	BitOrAggOp:            "bit_or",
-	BoolAndOp:             "bool_and",
-	BoolOrOp:              "bool_or",
-	ConcatAggOp:           "concat_agg",
-	CountOp:               "count",
-	CorrOp:                "corr",
-	CountRowsOp:           "count_rows",
-	CovarPopOp:            "covar_pop",
-	CovarSampOp:           "covar_samp",
-	RegressionAvgXOp:      "regr_avgx",
-	RegressionAvgYOp:      "regr_avgy",
-	RegressionInterceptOp: "regr_intercept",
-	RegressionR2Op:        "regr_r2",
-	RegressionSlopeOp:     "regr_slope",
-	RegressionSXXOp:       "regr_sxx",
-	RegressionSXYOp:       "regr_sxy",
-	RegressionSYYOp:       "regr_syy",
-	RegressionCountOp:     "regr_count",
-	MaxOp:                 "max",
-	MinOp:                 "min",
-	SumIntOp:              "sum_int",
-	SumOp:                 "sum",
-	SqrDiffOp:             "sqrdiff",
-	VarianceOp:            "variance",
-	StdDevOp:              "stddev",
-	XorAggOp:              "xor_agg",
-	JsonAggOp:             "json_agg",
-	JsonbAggOp:            "jsonb_agg",
-	JsonObjectAggOp:       "json_object_agg",
-	JsonbObjectAggOp:      "jsonb_object_agg",
-	StringAggOp:           "string_agg",
-	ConstAggOp:            "any_not_null",
-	ConstNotNullAggOp:     "any_not_null",
-	AnyNotNullAggOp:       "any_not_null",
-	PercentileDiscOp:      "percentile_disc_impl",
-	PercentileContOp:      "percentile_cont_impl",
-	VarPopOp:              "var_pop",
-	StdDevPopOp:           "stddev_pop",
-	STMakeLineOp:          "st_makeline",
-	STUnionOp:             "st_union",
-	STCollectOp:           "st_collect",
-	STExtentOp:            "st_extent",
+	ArrayAggOp:              "array_agg",
+	ArrayCatAggOp:           "array_cat_agg",
+	AvgOp:                   "avg",
+	BitAndAggOp:             "bit_and",
+	BitOrAggOp:              "bit_or",
+	BoolAndOp:               "bool_and",
+	BoolOrOp:                "bool_or",
+	ConcatAggOp:             "concat_agg",
+	CountOp:                 "count",
+	CorrOp:                  "corr",
+	CountRowsOp:             "count_rows",
+	CovarPopOp:              "covar_pop",
+	CovarSampOp:             "covar_samp",
+	RegressionAvgXOp:        "regr_avgx",
+	RegressionAvgYOp:        "regr_avgy",
+	RegressionInterceptOp:   "regr_intercept",
+	RegressionR2Op:          "regr_r2",
+	RegressionSlopeOp:       "regr_slope",
+	RegressionSXXOp:         "regr_sxx",
+	RegressionSXYOp:         "regr_sxy",
+	RegressionSYYOp:         "regr_syy",
+	RegressionCountOp:       "regr_count",
+	MaxOp:                   "max",
+	MinOp:                   "min",
+	SumIntOp:                "sum_int",
+	SumOp:                   "sum",
+	SqrDiffOp:               "sqrdiff",
+	VarianceOp:              "variance",
+	StdDevOp:                "stddev",
+	XorAggOp:                "xor_agg",
+	JsonAggOp:               "json_agg",
+	JsonbAggOp:              "jsonb_agg",
+	JsonObjectAggOp:         "json_object_agg",
+	JsonbObjectAggOp:        "jsonb_object_agg",
+	StringAggOp:             "string_agg",
+	ConstAggOp:              "any_not_null",
+	ConstNotNullAggOp:       "any_not_null",
+	AnyNotNullAggOp:         "any_not_null",
+	PercentileDiscOp:        "percentile_disc_impl",
+	PercentileContOp:        "percentile_cont_impl",
+	VarPopOp:                "var_pop",
+	StdDevPopOp:             "stddev_pop",
+	STMakeLineOp:            "st_makeline",
+	STUnionOp:               "st_union",
+	STCollectOp:             "st_collect",
+	STExtentOp:              "st_extent",
+	MergeStatsMetadataOp:    "merge_stats_metadata",
+	MergeStatementStatsOp:   "merge_statement_stats",
+	MergeTransactionStatsOp: "merge_transaction_stats",
 }
 
 // WindowOpReverseMap maps from an optimizer operator type to the name of a
@@ -327,7 +330,8 @@ func AggregateIgnoresNulls(op Operator) bool {
 		PercentileContOp, STMakeLineOp, STCollectOp, STExtentOp, STUnionOp, StdDevPopOp,
 		VarPopOp, CovarPopOp, CovarSampOp, RegressionAvgXOp, RegressionAvgYOp,
 		RegressionInterceptOp, RegressionR2Op, RegressionSlopeOp, RegressionSXXOp,
-		RegressionSXYOp, RegressionSYYOp, RegressionCountOp:
+		RegressionSXYOp, RegressionSYYOp, RegressionCountOp, MergeStatsMetadataOp,
+		MergeStatementStatsOp, MergeTransactionStatsOp:
 		return true
 
 	case ArrayAggOp, ArrayCatAggOp, ConcatAggOp, ConstAggOp, CountRowsOp,
@@ -353,7 +357,8 @@ func AggregateIsNullOnEmpty(op Operator) bool {
 		JsonObjectAggOp, JsonbObjectAggOp, StdDevPopOp, STCollectOp, STExtentOp, STUnionOp,
 		VarPopOp, CovarPopOp, CovarSampOp, RegressionAvgXOp, RegressionAvgYOp,
 		RegressionInterceptOp, RegressionR2Op, RegressionSlopeOp, RegressionSXXOp,
-		RegressionSXYOp, RegressionSYYOp:
+		RegressionSXYOp, RegressionSYYOp, MergeStatsMetadataOp, MergeStatementStatsOp,
+		MergeTransactionStatsOp:
 		return true
 
 	case CountOp, CountRowsOp, RegressionCountOp:
@@ -381,7 +386,8 @@ func AggregateIsNeverNullOnNonNullInput(op Operator) bool {
 		StringAggOp, SumOp, SumIntOp, XorAggOp, PercentileDiscOp, PercentileContOp,
 		JsonObjectAggOp, JsonbObjectAggOp, StdDevPopOp, STCollectOp, STUnionOp,
 		VarPopOp, CovarPopOp, RegressionAvgXOp, RegressionAvgYOp, RegressionSXXOp,
-		RegressionSXYOp, RegressionSYYOp, RegressionCountOp:
+		RegressionSXYOp, RegressionSYYOp, RegressionCountOp, MergeStatsMetadataOp,
+		MergeStatementStatsOp, MergeTransactionStatsOp:
 		return true
 
 	case VarianceOp, StdDevOp, CorrOp, CovarSampOp, RegressionInterceptOp,
@@ -436,7 +442,8 @@ func AggregatesCanMerge(inner, outer Operator) bool {
 		SqrDiffOp, STCollectOp, StdDevOp, StringAggOp, VarianceOp, StdDevPopOp,
 		VarPopOp, CovarPopOp, CovarSampOp, RegressionAvgXOp, RegressionAvgYOp,
 		RegressionInterceptOp, RegressionR2Op, RegressionSlopeOp, RegressionSXXOp,
-		RegressionSXYOp, RegressionSYYOp, RegressionCountOp:
+		RegressionSXYOp, RegressionSYYOp, RegressionCountOp, MergeStatsMetadataOp,
+		MergeStatementStatsOp, MergeTransactionStatsOp:
 		return false
 
 	default:
@@ -458,7 +465,8 @@ func AggregateIgnoresDuplicates(op Operator) bool {
 		VarPopOp, JsonObjectAggOp, JsonbObjectAggOp, STCollectOp, CovarPopOp,
 		CovarSampOp, RegressionAvgXOp, RegressionAvgYOp, RegressionInterceptOp,
 		RegressionR2Op, RegressionSlopeOp, RegressionSXXOp, RegressionSXYOp,
-		RegressionSYYOp, RegressionCountOp:
+		RegressionSYYOp, RegressionCountOp, MergeStatsMetadataOp, MergeStatementStatsOp,
+		MergeTransactionStatsOp:
 		return false
 
 	default:

diff --git a/pkg/sql/opt/ops/scalar.opt b/pkg/sql/opt/ops/scalar.opt
@@ -1015,6 +1015,21 @@ define JsonbObjectAgg {
     Value ScalarExpr
 }
 
+[Scalar, Aggregate]
+define MergeStatsMetadata {
+    Input ScalarExpr
+}
+
+[Scalar, Aggregate]
+define MergeStatementStats {
+    Input ScalarExpr
+}
+
+[Scalar, Aggregate]
+define MergeTransactionStats {
+    Input ScalarExpr
+}
+
 [Scalar, Aggregate]
 define StringAgg {
     Input ScalarExpr

diff --git a/pkg/sql/opt/optbuilder/groupby.go b/pkg/sql/opt/optbuilder/groupby.go
@@ -891,6 +891,12 @@ func (b *Builder) constructAggregate(name string, args []opt.ScalarExpr) opt.Sca
 		return b.factory.ConstructJsonObjectAgg(args[0], args[1])
 	case "jsonb_object_agg":
 		return b.factory.ConstructJsonbObjectAgg(args[0], args[1])
+	case "merge_stats_metadata":
+		return b.factory.ConstructMergeStatsMetadata(args[0])
+	case "merge_statement_stats":
+		return b.factory.ConstructMergeStatementStats(args[0])
+	case "merge_transaction_stats":
+		return b.factory.ConstructMergeTransactionStats(args[0])
 	}
 
 	panic(errors.AssertionFailedf("unhandled aggregate: %s", name))