Merge #152469 #153344

152469: sql/stats: support collecting partial statistics on arbitrary constraints r=Uzair5162 a=Uzair5162

This commit adds support to manually create partial statistics with a `WHERE` clause, where the predicate fully constrains the first column of an index. Additionally, the predicate must reference a single outer column, which must be the same as the single column that stats are being collected on. These stats are stored in `system.table_statistics` with their predicate.

Part of: #93998

Release note (sql change): Users can now manually create single-column partial statistics on boolean predicate expressions that can become simple index scans. These statistics can be created by adding a constraining `WHERE` expression to `CREATE STATISTICS`.

For example:

```
CREATE TABLE t (a INT PRIMARY KEY);
INSERT INTO t VALUES (1), (2), (3), (4), (5);
CREATE STATISTICS constrained_stat ON a FROM t WHERE a > 2;
```

153344: kvstorage: add echotest for DestroyReplica r=arulajmani a=pav-kv

This PR introduces a basic echotest for `DestroyReplica` storage operation. To be extended with a similar test for replica destruction with separated engines.

Part of #152845

Co-authored-by: Uzair Ahmad <[email protected]>
Co-authored-by: Pavel Kalinnikov <[email protected]>

Author: 3 people

pkg/jobs/jobspb/jobs.proto

@@ -1264,6 +1264,21 @@ message CreateStatsDetails {
 
   // If true, will collect partial table statistics at extreme values.
   bool using_extremes = 9;
+
+  // WHERE clause for partial statistics collection. This field is only used
+  // to populate the predicate in the system.table_statistics table and to
+  // determine if this is a constrained stats collection. The actual constrained
+  // scan is done over the spans in the where_spans field.
+  string where_clause = 10;
+
+  // Spans over which to collect partial statistics with a WHERE clause.
+  repeated roachpb.Span where_spans = 11 [(gogoproto.nullable) = false];
+
+  // The ID of the index used to collect partial statistics with a WHERE clause.
+  int64 where_index_id = 12 [
+    (gogoproto.customname) = "WhereIndexID",
+    (gogoproto.casttype) = "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb.IndexID"
+  ];
 }
 
 message CreateStatsProgress {

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -38,26 +38,34 @@ go_test(
     srcs = [
         "cluster_version_test.go",
         "datadriven_test.go",
+        "destroy_test.go",
     ],
     data = glob(["testdata/**"]),
     embed = [":kvstorage"],
     deps = [
         "//pkg/clusterversion",
         "//pkg/keys",
+        "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/concurrency/lock",
+        "//pkg/kv/kvserver/kvserverpb",
+        "//pkg/kv/kvserver/logstore",
+        "//pkg/kv/kvserver/print",
         "//pkg/kv/kvserver/stateloader",
         "//pkg/raft/raftpb",
         "//pkg/roachpb",
         "//pkg/settings/cluster",
         "//pkg/storage",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",
+        "//pkg/testutils/echotest",
         "//pkg/util/hlc",
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "//pkg/util/stop",
         "//pkg/util/tracing",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_datadriven//:datadriven",
+        "@com_github_cockroachdb_redact//:redact",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/kv/kvserver/kvstorage/datadriven_test.go

@@ -8,6 +8,7 @@ package kvstorage
 import (
 	"context"
 	"fmt"
+	"path/filepath"
 	"regexp"
 	"strings"
 	"testing"
@@ -100,7 +101,8 @@ func TestDataDriven(t *testing.T) {
 	// Scan stats (shown after loading the range descriptors) can be non-deterministic.
 	reStripScanStats := regexp.MustCompile(`stats: .*$`)
 
-	datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
+	dir := filepath.Join(datapathutils.TestDataPath(t), t.Name())
+	datadriven.Walk(t, dir, func(t *testing.T, path string) {
 		e := newEnv(t)
 		defer e.close()
 		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) (output string) {

pkg/kv/kvserver/kvstorage/destroy_test.go

@@ -0,0 +1,144 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package kvstorage
+
+import (
+	"context"
+	"fmt"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/stateloader"
+	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/echotest"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/cockroachdb/redact"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDestroyReplica(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	storage.DisableMetamorphicSimpleValueEncoding(t) // for deterministic output
+	eng := storage.NewDefaultInMemForTesting()
+	defer eng.Close()
+
+	var sb redact.StringBuilder
+	ctx := context.Background()
+	mutate := func(name string, write func(storage.ReadWriter)) {
+		b := eng.NewBatch()
+		defer b.Close()
+		write(b)
+		str, err := print.DecodeWriteBatch(b.Repr())
+		require.NoError(t, err)
+		_, err = sb.WriteString(fmt.Sprintf(">> %s:\n%s", name, str))
+		require.NoError(t, err)
+		require.NoError(t, b.Commit(false))
+	}
+
+	r := replicaInfo{
+		id:      roachpb.FullReplicaID{RangeID: 123, ReplicaID: 3},
+		hs:      raftpb.HardState{Term: 5, Commit: 14},
+		ts:      kvserverpb.RaftTruncatedState{Index: 10, Term: 5},
+		keys:    roachpb.RSpan{Key: []byte("a"), EndKey: []byte("z")},
+		last:    15,
+		applied: 12,
+	}
+	mutate("raft", func(rw storage.ReadWriter) {
+		r.createRaftState(ctx, t, rw)
+	})
+	mutate("state", func(rw storage.ReadWriter) {
+		r.createStateMachine(ctx, t, rw)
+	})
+	mutate("destroy", func(rw storage.ReadWriter) {
+		require.NoError(t, DestroyReplica(ctx, r.id, rw, rw, r.id.ReplicaID+1, ClearRangeDataOptions{
+			ClearUnreplicatedByRangeID: true,
+			ClearReplicatedByRangeID:   true,
+			ClearReplicatedBySpan:      r.keys,
+		}))
+	})
+
+	str := strings.ReplaceAll(sb.String(), "\n\n", "\n")
+	echotest.Require(t, str, filepath.Join(datapathutils.TestDataPath(t), t.Name()+".txt"))
+}
+
+// replicaInfo contains the basic info about the replica, used for generating
+// its storage counterpart.
+//
+// TODO(pav-kv): make it reusable for other tests.
+type replicaInfo struct {
+	id      roachpb.FullReplicaID
+	hs      raftpb.HardState
+	ts      kvserverpb.RaftTruncatedState
+	keys    roachpb.RSpan
+	last    kvpb.RaftIndex
+	applied kvpb.RaftIndex
+}
+
+func (r *replicaInfo) createRaftState(ctx context.Context, t *testing.T, w storage.Writer) {
+	sl := logstore.NewStateLoader(r.id.RangeID)
+	require.NoError(t, sl.SetHardState(ctx, w, r.hs))
+	require.NoError(t, sl.SetRaftTruncatedState(ctx, w, &r.ts))
+	for i := r.ts.Index + 1; i <= r.last; i++ {
+		require.NoError(t, storage.MVCCBlindPutProto(
+			ctx, w,
+			sl.RaftLogKey(i), hlc.Timestamp{}, /* timestamp */
+			&raftpb.Entry{Index: uint64(i), Term: 5},
+			storage.MVCCWriteOptions{},
+		))
+	}
+}
+
+func (r *replicaInfo) createStateMachine(ctx context.Context, t *testing.T, rw storage.ReadWriter) {
+	sl := stateloader.Make(r.id.RangeID)
+	require.NoError(t, sl.SetRangeTombstone(ctx, rw, kvserverpb.RangeTombstone{
+		NextReplicaID: r.id.ReplicaID,
+	}))
+	require.NoError(t, sl.SetRaftReplicaID(ctx, rw, r.id.ReplicaID))
+	// TODO(pav-kv): figure out whether LastReplicaGCTimestamp should be in the
+	// log or state engine.
+	require.NoError(t, storage.MVCCBlindPutProto(
+		ctx, rw,
+		keys.RangeLastReplicaGCTimestampKey(r.id.RangeID),
+		hlc.Timestamp{}, /* timestamp */
+		&hlc.Timestamp{WallTime: 12345678},
+		storage.MVCCWriteOptions{},
+	))
+	createRangeData(t, rw, r.keys)
+}
+
+func createRangeData(t *testing.T, rw storage.ReadWriter, span roachpb.RSpan) {
+	ts := hlc.Timestamp{WallTime: 1}
+	for _, k := range []roachpb.Key{
+		keys.RangeDescriptorKey(span.Key),   // system
+		span.Key.AsRawKey(),                 // user
+		roachpb.Key(span.EndKey).Prevish(2), // user
+	} {
+		// Put something under the system or user key.
+		require.NoError(t, rw.PutMVCC(
+			storage.MVCCKey{Key: k, Timestamp: ts}, storage.MVCCValue{},
+		))
+		// Put something under the corresponding lock key.
+		ek, _ := storage.LockTableKey{
+			Key: k, Strength: lock.Intent, TxnUUID: uuid.UUID{},
+		}.ToEngineKey(nil)
+		require.NoError(t, rw.PutEngineKey(ek, nil))
+	}
+}