Add support for specifying bitcount unit as byte or bit, byte default (#2887)

* Add support for specifying bitcount unit as byte or bit, byte default

* Add bitcount test

* Test bitcount without unit specified

---------

Co-authored-by: wanghongwei5 <wanghongwei5@360.cn>
Co-authored-by: ofekshenawa <104765379+ofekshenawa@users.noreply.github.com>
This commit is contained in:
Nanthen Hale 2024-02-15 04:15:30 +08:00 committed by GitHub
parent 36bab9c8dc
commit 8d2a022fd5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 112 additions and 0 deletions

View File

@ -2,6 +2,7 @@ package redis
import (
"context"
"errors"
)
type BitMapCmdable interface {
@ -37,15 +38,28 @@ func (c cmdable) SetBit(ctx context.Context, key string, offset int64, value int
type BitCount struct {
Start, End int64
Unit string // BYTE(default) | BIT
}
const BitCountIndexByte string = "BYTE"
const BitCountIndexBit string = "BIT"
func (c cmdable) BitCount(ctx context.Context, key string, bitCount *BitCount) *IntCmd {
args := []interface{}{"bitcount", key}
if bitCount != nil {
if bitCount.Unit == "" {
bitCount.Unit = "BYTE"
}
if bitCount.Unit != BitCountIndexByte && bitCount.Unit != BitCountIndexBit {
cmd := NewIntCmd(ctx)
cmd.SetErr(errors.New("redis: invalid bitcount index"))
return cmd
}
args = append(
args,
bitCount.Start,
bitCount.End,
string(bitCount.Unit),
)
}
cmd := NewIntCmd(ctx, args...)

98
bitmap_commands_test.go Normal file
View File

@ -0,0 +1,98 @@
package redis_test
import (
. "github.com/bsm/ginkgo/v2"
. "github.com/bsm/gomega"
"github.com/redis/go-redis/v9"
)
type bitCountExpected struct {
Start int64
End int64
Expected int64
}
var _ = Describe("BitCountBite", func() {
var client *redis.Client
key := "bit_count_test"
BeforeEach(func() {
client = redis.NewClient(redisOptions())
Expect(client.FlushDB(ctx).Err()).NotTo(HaveOccurred())
values := []int{0, 1, 0, 0, 1, 0, 1, 0, 1, 1}
for i, v := range values {
cmd := client.SetBit(ctx, key, int64(i), v)
Expect(cmd.Err()).NotTo(HaveOccurred())
}
})
AfterEach(func() {
Expect(client.Close()).NotTo(HaveOccurred())
})
It("bit count bite", func() {
var expected = []bitCountExpected{
{0, 0, 0},
{0, 1, 1},
{0, 2, 1},
{0, 3, 1},
{0, 4, 2},
{0, 5, 2},
{0, 6, 3},
{0, 7, 3},
{0, 8, 4},
{0, 9, 5},
}
for _, e := range expected {
cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End, Unit: redis.BitCountIndexBit})
Expect(cmd.Err()).NotTo(HaveOccurred())
Expect(cmd.Val()).To(Equal(e.Expected))
}
})
})
var _ = Describe("BitCountByte", func() {
var client *redis.Client
key := "bit_count_test"
BeforeEach(func() {
client = redis.NewClient(redisOptions())
Expect(client.FlushDB(ctx).Err()).NotTo(HaveOccurred())
values := []int{0, 0, 0, 0, 0, 0, 0, 1, 1, 1}
for i, v := range values {
cmd := client.SetBit(ctx, key, int64(i), v)
Expect(cmd.Err()).NotTo(HaveOccurred())
}
})
AfterEach(func() {
Expect(client.Close()).NotTo(HaveOccurred())
})
It("bit count byte", func() {
var expected = []bitCountExpected{
{0, 0, 1},
{0, 1, 3},
}
for _, e := range expected {
cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End, Unit: redis.BitCountIndexByte})
Expect(cmd.Err()).NotTo(HaveOccurred())
Expect(cmd.Val()).To(Equal(e.Expected))
}
})
It("bit count byte with no unit specified", func() {
var expected = []bitCountExpected{
{0, 0, 1},
{0, 1, 3},
}
for _, e := range expected {
cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End})
Expect(cmd.Err()).NotTo(HaveOccurred())
Expect(cmd.Val()).To(Equal(e.Expected))
}
})
})