Add sample option for large RDB files

README.md

@@ -74,6 +74,7 @@ Options:
     1. '1KB~1MB' get keys with size in range [1KB, 1MB]
     2. '10MB~inf' magic variable 'inf' represents the Infinity
     3. '1024~10KB' get keys with size in range [0Bytes, 10KB]
+  -sample sample rate between (0, 1] for sampling keys
   -concurrent The number of concurrent json converters. 4 by default.
   -show-global-meta Show global meta likes redis-verion/ctime/functions
   -no-expired filter expired keys(deprecated, please use 'expire' option)
@@ -618,6 +619,19 @@ rdb -c json -o dump.json -size 1KB~1MB cases/memory.rdb
 rdb -c memory -o memory.csv -size 10MB~inf cases/memory.rdb
 ```
 
+# Sample Filter
+
+The `-sample` parameter can be configured to sample keys from large RDB files with a sampling rate in `(0, 1]`.
+
+Examples:
+```bash
+# Sample 10% of keys to generate JSON
+rdb -c json -o dump.json -sample 0.1 dump.rdb
+
+# Sample 20% of keys to generate memory report
+rdb -c memory -o memory.csv -sample 0.2 dump.rdb
+```
+
 # Expiration Filter
 
 The `-expire` parameter can be configured to filter based on the expiration time.

README_CN.md

@@ -60,6 +60,7 @@ Options:
     1. '1KB~1MB' get keys with size in range [1KB, 1MB]
     2. '10MB~inf' magic variable 'inf' represents the Infinity
     3. '1024~10KB' get keys with size in range [0Bytes, 10KB]
+  -sample sample rate between (0, 1] for sampling keys
   -concurrent The number of concurrent json converters. 4 by default.
   -show-global-meta Show global meta likes redis-verion/ctime/functions
   -no-expired filter expired keys(deprecated, please use 'expire' option)
@@ -544,6 +545,19 @@ rdb -c json -o dump.json -size 1KB~1MB cases/memory.rdb
 rdb -c memory -o memory.csv -size 10MB~inf cases/memory.rdb
 ```
 
+# 抽样过滤器
+
+`-sample` 参数可以配置对大型 RDB 文件中的键进行抽样处理，取值范围为 `(0, 1]`。
+
+示例：
+```bash
+# 抽样 10% 的键生成 JSON
+rdb -c json -o dump.json -sample 0.1 dump.rdb
+
+# 抽样 20% 的键生成内存报告
+rdb -c memory -o memory.csv -sample 0.2 dump.rdb
+```
+
 # Expiration 过滤器
 
 `-expire` 参数可以配置根据过期时间过滤.

cmd.go

@@ -32,6 +32,7 @@ Options:
     1. '1KB~1MB' get keys with size in range [1KB, 1MB]
     2. '10MB~inf' magic variable 'inf' represents the Infinity
     3. '1024~10KB' get keys with size in range [0Bytes, 10KB]
+  -sample sample rate between (0, 1] for sampling keys
   -concurrent The number of concurrent json converters. 4 by default.
   -show-global-meta Show global meta likes redis-verion/ctime/functions
   -no-expired filter expired keys(deprecated, please use 'expire' option)
@@ -82,6 +83,7 @@ func main() {
 	var concurrent int
 	var showGlobalMeta bool
 	var prefixSeps separators
+	var sample float64
 	var err error
 	flagSet.StringVar(&cmd, "c", "", "command for rdb: json")
 	flagSet.StringVar(&output, "o", "", "output file path")
@@ -93,6 +95,7 @@ func main() {
 	flagSet.StringVar(&regexExpr, "regex", "", "regex expression")
 	flagSet.StringVar(&expirationExpr, "expire", "", "expiration filter expression")
 	flagSet.StringVar(&sizeExpr, "size", "", "size filter expression")
+	flagSet.Float64Var(&sample, "sample", 0, "sample rate between (0, 1] for sampling keys")
 	flagSet.BoolVar(&noExpired, "no-expired", false, "filter expired keys(deprecated, please use expire)")
 	flagSet.Var(&prefixSeps, "prefix-sep", "separator for prefix analysis (flat-map mode, constant memory)")
 	flagSet.BoolVar(&showGlobalMeta, "show-global-meta", false, "Show global meta likes redis-verion/ctime/functions")
@@ -124,6 +127,15 @@ func main() {
 	if sizeExpr != "" {
 		options = append(options, helper.WithSizeOption(sizeExpr))
 	}
+	var hasSample bool
+	flagSet.Visit(func(f *flag.Flag) {
+		if f.Name == "sample" {
+			hasSample = true
+		}
+	})
+	if hasSample {
+		options = append(options, helper.WithSampleOption(sample))
+	}
 	if showGlobalMeta {
 		options = append(options, helper.WithGlobalMeta())
 	}
@@ -147,7 +159,7 @@ func main() {
 	case "memory":
 		err = helper.MemoryProfile(src, output, options...)
 	case "aof":
-		err = helper.ToAOF(src, output, options)
+		err = helper.ToAOF(src, output, options...)
 	case "bigkey":
 		err = helper.FindBiggestKeys(src, n, outputFile, options...)
 	case "hotkey":

cmd_test.go

@@ -69,6 +69,32 @@ func TestCmd(t *testing.T) {
 		t.Error("command prefix with prefix-sep failed")
 	}
 
+	os.Args = []string{"", "-c", "memory", "-sample", "0.5", "-o", "tmp/memory_sample.csv", "cases/memory.rdb"}
+	main()
+	if f, _ := os.Stat("tmp/memory_sample.csv"); f == nil {
+		t.Error("command memory with sample failed")
+	}
+
+	os.Args = []string{"", "-c", "json", "-sample", "0.5", "-o", "tmp/cmd_sample.json", "cases/memory.rdb"}
+	main()
+	if f, _ := os.Stat("tmp/cmd_sample.json"); f == nil {
+		t.Error("command json with sample failed")
+	}
+
+	os.Args = []string{"", "-c", "bigkey", "-sample", "0.5", "-n", "10", "-o", "tmp/bigkey_sample.csv", "cases/memory.rdb"}
+	main()
+	if f, _ := os.Stat("tmp/bigkey_sample.csv"); f == nil {
+		t.Error("command bigkey with sample failed")
+	}
+
+	// invalid sample rate commands
+	os.Args = []string{"", "-c", "memory", "-sample", "-0.1", "-o", "tmp/memory_invalid1.csv", "cases/memory.rdb"}
+	main()
+	os.Args = []string{"", "-c", "memory", "-sample", "0", "-o", "tmp/memory_invalid2.csv", "cases/memory.rdb"}
+	main()
+	os.Args = []string{"", "-c", "memory", "-sample", "1.5", "-o", "tmp/memory_invalid3.csv", "cases/memory.rdb"}
+	main()
+
 	// test error command line
 	os.Args = []string{"", "-c", "json", "-o", "tmp/output", "/none/a"}
 	main()

helper/regex.go

@@ -211,8 +211,15 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) {
 	var expirationOpt ExpirationOption
 	var sizeOpt SizeOption
 	var globalMetaOpt GlobalMetaOption
+	var sampleOpt *SampleOption
 	for _, opt := range options {
 		switch o := opt.(type) {
+		case []interface{}:
+			var err error
+			dec, err = wrapDecoder(dec, o...)
+			if err != nil {
+				return nil, err
+			}
 		case RegexOption:
 			regexOpt = o
 		case NoExpiredOption:
@@ -223,6 +230,14 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) {
 			sizeOpt = o
 		case GlobalMetaOption:
 			globalMetaOpt = o
+		case SampleOption:
+			sampleOpt = &o
+		}
+	}
+	if sampleOpt != nil {
+		rate := float64(*sampleOpt)
+		if rate <= 0 || rate > 1 {
+			return nil, fmt.Errorf("illegal sample rate: %v, should be in range (0, 1]", rate)
 		}
 	}
 	if globalMetaOpt {
@@ -274,5 +289,11 @@ func wrapDecoder(dec decoder, options ...interface{}) (decoder, error) {
 			sizeRange: rng,
 		}
 	}
+	if sampleOpt != nil {
+		rate := float64(*sampleOpt)
+		if rate < 1.0 {
+			dec = newSampleDecoder(dec, rate)
+		}
+	}
 	return dec, nil
 }

helper/sample.go

@@ -0,0 +1,52 @@
+package helper
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/hdt3213/rdb/model"
+)
+
+// SampleOption enables sampling keys with a rate in (0, 1]
+type SampleOption float64
+
+// WithSampleOption creates a SampleOption from sample rate
+func WithSampleOption(rate float64) SampleOption {
+	return SampleOption(rate)
+}
+
+// WithSample creates a SampleOption from sample rate
+func WithSample(rate float64) SampleOption {
+	return SampleOption(rate)
+}
+
+type sampleDecoder struct {
+	dec  decoder
+	rate float64
+	rnd  *rand.Rand
+}
+
+func newSampleDecoder(dec decoder, rate float64) *sampleDecoder {
+	return &sampleDecoder{
+		dec:  dec,
+		rate: rate,
+		rnd:  rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+}
+
+func newSampleDecoderWithRand(dec decoder, rate float64, rnd *rand.Rand) *sampleDecoder {
+	return &sampleDecoder{
+		dec:  dec,
+		rate: rate,
+		rnd:  rnd,
+	}
+}
+
+func (d *sampleDecoder) Parse(cb func(object model.RedisObject) bool) error {
+	return d.dec.Parse(func(object model.RedisObject) bool {
+		if d.rnd.Float64() < d.rate {
+			return cb(object)
+		}
+		return true
+	})
+}