diff --git a/src/commands/cmd_bit.cc b/src/commands/cmd_bit.cc index ba9f721e937..2cf273f576d 100644 --- a/src/commands/cmd_bit.cc +++ b/src/commands/cmd_bit.cc @@ -219,11 +219,22 @@ class CommandBitOp : public Commander { op_flag_ = kBitOpXor; else if (opname == "not") op_flag_ = kBitOpNot; + else if (opname == "diff") + op_flag_ = kBitOpDiff; + else if (opname == "diff1") + op_flag_ = kBitOpDiff1; + else if (opname == "andor") + op_flag_ = kBitOpAndOr; + else if (opname == "one") + op_flag_ = kBitOpOne; else return {Status::RedisInvalidCmd, errInvalidSyntax}; if (op_flag_ == kBitOpNot && args.size() != 4) { return {Status::RedisInvalidCmd, "BITOP NOT must be called with a single source key."}; } + if ((op_flag_ == kBitOpDiff || op_flag_ == kBitOpDiff1 || op_flag_ == kBitOpAndOr) && args.size() < 5) { + return {Status::RedisInvalidCmd, "BITOP DIFF/DIFF1/ANDOR must be called with at least two source keys."}; + } return Commander::Parse(args); } diff --git a/src/types/redis_bitmap.cc b/src/types/redis_bitmap.cc index 8a7cee9193d..495c35436cb 100644 --- a/src/types/redis_bitmap.cc +++ b/src/types/redis_bitmap.cc @@ -483,6 +483,10 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st } size_t num_keys = meta_pairs.size(); + // Determine if the first source key (X) exists in meta_pairs. + // meta_pairs preserves op_keys order for existing keys, so meta_pairs[0] is X iff X exists. + const bool first_key_exists = !meta_pairs.empty() && meta_pairs[0].first == AppendNamespacePrefix(op_keys[0]); + auto batch = storage_->GetWriteBatchBase(); if (max_bitmap_size == 0) { /* Compute the bit operation, if all bitmap is empty. cleanup the dest bitmap. */ @@ -499,9 +503,10 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st if (!s.ok()) return s; BitmapMetadata res_metadata; - // If the operation is AND and the number of keys is less than the number of op_keys, - // we can skip setting the subkeys of the result bitmap and just set the metadata. - const bool can_skip_op = op_flag == kBitOpAnd && num_keys != op_keys.size(); + // AND: any missing key means result is all zeros. + // DIFF/ANDOR: missing X means result is all zeros (X=0 → X & anything = 0). + const bool can_skip_op = (op_flag == kBitOpAnd && num_keys != op_keys.size()) || + ((op_flag == kBitOpDiff || op_flag == kBitOpAndOr) && !first_key_exists); if (!can_skip_op) { uint64_t stop_index = (max_bitmap_size - 1) / kBitmapSegmentBytes; std::unique_ptr frag_res(new unsigned char[kBitmapSegmentBytes]); @@ -510,6 +515,9 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st for (uint64_t frag_index = 0; frag_index <= stop_index; frag_index++) { std::vector fragments; uint16_t frag_maxlen = 0, frag_minlen = 0; + // Tracks whether fragments[0] is X's fragment (only relevant for DIFF/DIFF1/ANDOR). + bool x_frag_is_first = false; + bool is_first_meta_pair = true; for (const auto &meta_pair : meta_pairs) { std::string sub_key = InternalKey(meta_pair.first, std::to_string(frag_index * kBitmapSegmentBytes), meta_pair.second.version, storage_->IsSlotIdEncoded()) @@ -521,16 +529,24 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st } if (s.IsNotFound()) { if (op_flag == kBitOpAnd) { - // If any of the input bitmaps is empty, the result of AND - // is empty. + // If any of the input bitmaps is empty, the result of AND is empty. + frag_maxlen = 0; + break; + } + // For DIFF/ANDOR: X's segment missing means result is 0 for this segment. + if ((op_flag == kBitOpDiff || op_flag == kBitOpAndOr) && first_key_exists && is_first_meta_pair) { frag_maxlen = 0; break; } } else { if (frag_maxlen < fragment.size()) frag_maxlen = fragment.size(); if (fragment.size() < frag_minlen || frag_minlen == 0) frag_minlen = fragment.size(); + if (is_first_meta_pair && first_key_exists && fragments.empty()) { + x_frag_is_first = true; + } fragments.emplace_back(std::move(fragment)); } + is_first_meta_pair = false; } size_t frag_numkeys = fragments.size(); @@ -548,7 +564,8 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st * result in GCC compiling the code using multiple-words load/store * operations that are not supported even in ARM >= v6. */ #ifndef USE_ALIGNED_ACCESS - if (frag_minlen >= sizeof(uint64_t) * 4 && frag_numkeys <= 16) { + if (frag_minlen >= sizeof(uint64_t) * 4 && frag_numkeys <= 16 && op_flag != kBitOpDiff && + op_flag != kBitOpDiff1 && op_flag != kBitOpAndOr && op_flag != kBitOpOne) { auto *lres = reinterpret_cast(frag_res.get()); const uint64_t *lp[16]; for (uint64_t i = 0; i < frag_numkeys; i++) { @@ -593,24 +610,69 @@ rocksdb::Status Bitmap::BitOp(engine::Context &ctx, BitOpFlags op_flag, const st } #endif + // For DIFF/DIFF1/ANDOR: y_start is where Y fragments begin in fragments[]. + // When x_frag_is_first=true, fragments[0]=X and Ys start at 1. + // When x_frag_is_first=false, X=0 (missing) and all fragments are Ys. + const uint64_t y_start = x_frag_is_first ? 1 : 0; + uint8_t output = 0, byte = 0; for (; j < frag_maxlen; j++) { - output = (fragments[0].size() <= j) ? 0 : fragments[0][j]; - if (op_flag == kBitOpNot) output = ~output; - for (uint64_t i = 1; i < frag_numkeys; i++) { - byte = (fragments[i].size() <= j) ? 0 : fragments[i][j]; - switch (op_flag) { - case kBitOpAnd: - output &= byte; - break; - case kBitOpOr: - output |= byte; - break; - case kBitOpXor: - output ^= byte; - break; - default: - break; + output = (fragments[0].size() <= j) ? 0 : static_cast(fragments[0][j]); + if (op_flag == kBitOpNot) { + output = ~output; + } else if (op_flag == kBitOpDiff1) { + // DIFF1: bits set in any Y but not in X + uint8_t x_byte = x_frag_is_first ? output : 0; + uint8_t or_rest = 0; + for (uint64_t i = y_start; i < frag_numkeys; i++) { + byte = (fragments[i].size() <= j) ? 0 : static_cast(fragments[i][j]); + or_rest |= byte; + } + output = or_rest & ~x_byte; + } else if (op_flag == kBitOpAndOr) { + // ANDOR: bits set in X AND in at least one Y + // (ANDOR with missing X is already handled by can_skip_op) + uint8_t or_rest = 0; + for (uint64_t i = y_start; i < frag_numkeys; i++) { + byte = (fragments[i].size() <= j) ? 0 : static_cast(fragments[i][j]); + or_rest |= byte; + } + output = output & or_rest; + } else if (op_flag == kBitOpOne) { + // ONE: bits set in exactly one key across all inputs + // xor_acc tracks odd parity, and_acc tracks bits set in 2+ keys + uint8_t xor_acc = output, and_acc = 0; + for (uint64_t i = 1; i < frag_numkeys; i++) { + byte = (fragments[i].size() <= j) ? 0 : static_cast(fragments[i][j]); + and_acc |= (xor_acc & byte); + xor_acc ^= byte; + } + output = xor_acc & ~and_acc; + } else { + // For DIFF: X = fragments[0] if x_frag_is_first, else X = 0. + // DIFF with missing X is handled by can_skip_op, so x_frag_is_first=true here. + if (op_flag == kBitOpDiff && !x_frag_is_first) { + output = 0; + } + for (uint64_t i = (op_flag == kBitOpDiff ? y_start : 1); i < frag_numkeys; i++) { + byte = (fragments[i].size() <= j) ? 0 : static_cast(fragments[i][j]); + switch (op_flag) { + case kBitOpAnd: + output &= byte; + break; + case kBitOpOr: + output |= byte; + break; + case kBitOpXor: + output ^= byte; + break; + case kBitOpDiff: + // DIFF: bits set in X but not in any Y + output &= ~byte; + break; + default: + break; + } } } frag_res[j] = output; diff --git a/src/types/redis_bitmap.h b/src/types/redis_bitmap.h index 32a53cc72ab..6fa0cc16233 100644 --- a/src/types/redis_bitmap.h +++ b/src/types/redis_bitmap.h @@ -34,6 +34,10 @@ enum BitOpFlags { kBitOpOr, kBitOpXor, kBitOpNot, + kBitOpDiff, + kBitOpDiff1, + kBitOpAndOr, + kBitOpOne, }; namespace redis { diff --git a/tests/gocase/unit/type/bitmap/bitmap_test.go b/tests/gocase/unit/type/bitmap/bitmap_test.go index 0f3399316d9..0d2e601724d 100644 --- a/tests/gocase/unit/type/bitmap/bitmap_test.go +++ b/tests/gocase/unit/type/bitmap/bitmap_test.go @@ -39,10 +39,14 @@ import ( type BITOP int32 const ( - AND BITOP = 0 - OR BITOP = 1 - XOR BITOP = 2 - NOT BITOP = 3 + AND BITOP = 0 + OR BITOP = 1 + XOR BITOP = 2 + NOT BITOP = 3 + DIFF BITOP = 4 + DIFF1 BITOP = 5 + ANDOR BITOP = 6 + ONE BITOP = 7 ) func Set2SetBit(t *testing.T, rdb *redis.Client, ctx context.Context, key string, bs []byte) { @@ -88,22 +92,69 @@ func SimulateBitOp(op BITOP, values ...[]byte) string { } else { x = '0' } - } - for j := 1; j < len(binaryArray); j++ { - left := int(x - '0') - right := int(binaryArray[j][i] - '0') - switch op { - case AND: - left = left & right - case XOR: - left = left ^ right - case OR: - left = left | right + } else if op == DIFF { + // bits in X but not in any Y + for j := 1; j < len(binaryArray); j++ { + if binaryArray[j][i] == '1' { + x = '0' + } + } + } else if op == DIFF1 { + // bits in any Y but not in X + orRest := byte('0') + for j := 1; j < len(binaryArray); j++ { + if binaryArray[j][i] == '1' { + orRest = '1' + } } - if left == 0 { + if orRest == '1' && x == '0' { + x = '1' + } else { x = '0' + } + } else if op == ANDOR { + // bits in X AND at least one Y + orRest := byte('0') + for j := 1; j < len(binaryArray); j++ { + if binaryArray[j][i] == '1' { + orRest = '1' + } + } + if x == '1' && orRest == '1' { + x = '1' } else { + x = '0' + } + } else if op == ONE { + // bits set in exactly one key + count := 0 + for j := 0; j < len(binaryArray); j++ { + if binaryArray[j][i] == '1' { + count++ + } + } + if count == 1 { x = '1' + } else { + x = '0' + } + } else { + for j := 1; j < len(binaryArray); j++ { + left := int(x - '0') + right := int(binaryArray[j][i] - '0') + switch op { + case AND: + left = left & right + case XOR: + left = left ^ right + case OR: + left = left | right + } + if left == 0 { + x = '0' + } else { + x = '1' + } } } binaryResult = append(binaryResult, x) @@ -357,6 +408,170 @@ func TestBitmap(t *testing.T) { require.EqualValues(t, 32, rdb.BitOpOr(ctx, "x", "a", "b").Val()) }) + t.Run("BITOP DIFF basic", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + // X=0xff, Y=0x0f -> DIFF = 0xf0 (bits in X not in Y) + Set2SetBit(t, rdb, ctx, "x", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF", "dest", "x", "y").Err()) + require.EqualValues(t, SimulateBitOp(DIFF, []byte("\xff"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP DIFF with multiple Y keys", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "y1", []byte("\x0f")) + Set2SetBit(t, rdb, ctx, "y2", []byte("\xf0")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF", "dest", "x", "y1", "y2").Err()) + require.EqualValues(t, SimulateBitOp(DIFF, []byte("\xff"), []byte("\x0f"), []byte("\xf0")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP DIFF missing key treated as zero", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xaa")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF", "dest", "x", "no-such-key").Err()) + require.EqualValues(t, SimulateBitOp(DIFF, []byte("\xaa"), []byte("\x00")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP DIFF1 basic", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + // X=0xff, Y=0x0f -> DIFF1 = 0x00 (bits in Y not in X, but X has all bits set) + Set2SetBit(t, rdb, ctx, "x", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF1", "dest", "x", "y").Err()) + require.EqualValues(t, SimulateBitOp(DIFF1, []byte("\xff"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP DIFF1 with partial overlap", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + // X=0x0f, Y=0xff -> DIFF1 = 0xf0 (bits in Y not in X) + Set2SetBit(t, rdb, ctx, "x", []byte("\x0f")) + Set2SetBit(t, rdb, ctx, "y", []byte("\xff")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF1", "dest", "x", "y").Err()) + require.EqualValues(t, SimulateBitOp(DIFF1, []byte("\x0f"), []byte("\xff")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP ANDOR basic", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + // X=0xff, Y=0x0f -> ANDOR = 0x0f (bits in X AND at least one Y) + Set2SetBit(t, rdb, ctx, "x", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ANDOR", "dest", "x", "y").Err()) + require.EqualValues(t, SimulateBitOp(ANDOR, []byte("\xff"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP ANDOR with multiple Y keys", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "y1", []byte("\x0f")) + Set2SetBit(t, rdb, ctx, "y2", []byte("\xf0")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ANDOR", "dest", "x", "y1", "y2").Err()) + require.EqualValues(t, SimulateBitOp(ANDOR, []byte("\xff"), []byte("\x0f"), []byte("\xf0")), rdb.Get(ctx, "dest").Val()) + }) + + // Redis semantics: when X (first source key) does not exist, it is treated as + // a stream of zero bytes. So DIFF(nosuch, y) = 0 & ~y = 0, not y. + t.Run("BITOP DIFF missing first key X treated as zero (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF", "dest", "nosuch", "y").Err()) + // X=0x00, Y=0x0f -> DIFF = 0x00 & ~0x0f = 0x00 + require.EqualValues(t, SimulateBitOp(DIFF, []byte("\x00"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP DIFF1 missing first key X treated as zero (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "DIFF1", "dest", "nosuch", "y").Err()) + // X=0x00, Y=0x0f -> DIFF1 = 0x0f & ~0x00 = 0x0f + require.EqualValues(t, SimulateBitOp(DIFF1, []byte("\x00"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP ANDOR missing first key X treated as zero (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "y", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ANDOR", "dest", "nosuch", "y").Err()) + // X=0x00, Y=0x0f -> ANDOR = 0x00 & 0x0f = 0x00 + require.EqualValues(t, SimulateBitOp(ANDOR, []byte("\x00"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + // Redis requires at least 2 source keys for DIFF, DIFF1, ANDOR (X + at least one Y). + // Calling with only X and no Y keys should return an error. + t.Run("BITOP DIFF requires at least one Y key (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xaa")) + util.ErrorRegexp(t, rdb.Do(ctx, "BITOP", "DIFF", "dest", "x").Err(), ".*") + }) + + t.Run("BITOP DIFF1 requires at least one Y key (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xaa")) + util.ErrorRegexp(t, rdb.Do(ctx, "BITOP", "DIFF1", "dest", "x").Err(), ".*") + }) + + t.Run("BITOP ANDOR requires at least one Y key (Redis semantics)", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "x", []byte("\xaa")) + util.ErrorRegexp(t, rdb.Do(ctx, "BITOP", "ANDOR", "dest", "x").Err(), ".*") + }) + + t.Run("BITOP ONE basic", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + // A=0xff, B=0x0f -> ONE = 0xf0 (bits set in exactly one key) + Set2SetBit(t, rdb, ctx, "a", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "b", []byte("\x0f")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ONE", "dest", "a", "b").Err()) + require.EqualValues(t, SimulateBitOp(ONE, []byte("\xff"), []byte("\x0f")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP ONE with three keys", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "a", []byte("\xff")) + Set2SetBit(t, rdb, ctx, "b", []byte("\x0f")) + Set2SetBit(t, rdb, ctx, "c", []byte("\xf0")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ONE", "dest", "a", "b", "c").Err()) + require.EqualValues(t, SimulateBitOp(ONE, []byte("\xff"), []byte("\x0f"), []byte("\xf0")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP ONE single key returns same key", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + Set2SetBit(t, rdb, ctx, "a", []byte("\xaa")) + require.NoError(t, rdb.Do(ctx, "BITOP", "ONE", "dest", "a").Err()) + require.EqualValues(t, SimulateBitOp(ONE, []byte("\xaa")), rdb.Get(ctx, "dest").Val()) + }) + + t.Run("BITOP new ops fuzzing", func(t *testing.T) { + require.NoError(t, rdb.FlushDB(ctx).Err()) + for i := 0; i < 10; i++ { + numKeys := 2 + i%3 + vec := make([][]byte, numKeys) + veckeys := make([]string, numKeys) + for k := 0; k < numKeys; k++ { + vec[k] = []byte(util.RandStringWithSeed(1, 10, util.Binary, int64(i*100+k))) + veckeys[k] = fmt.Sprintf("fuzz-%d-%d", i, k) + Set2SetBit(t, rdb, ctx, veckeys[k], vec[k]) + } + doArgs := func(op string) []interface{} { + args := []interface{}{"BITOP", op, "target"} + for _, k := range veckeys { + args = append(args, k) + } + return args + } + require.NoError(t, rdb.Do(ctx, doArgs("DIFF")...).Err()) + require.EqualValues(t, SimulateBitOp(DIFF, vec...), rdb.Get(ctx, "target").Val()) + + require.NoError(t, rdb.Do(ctx, doArgs("DIFF1")...).Err()) + require.EqualValues(t, SimulateBitOp(DIFF1, vec...), rdb.Get(ctx, "target").Val()) + + require.NoError(t, rdb.Do(ctx, doArgs("ANDOR")...).Err()) + require.EqualValues(t, SimulateBitOp(ANDOR, vec...), rdb.Get(ctx, "target").Val()) + + require.NoError(t, rdb.Do(ctx, doArgs("ONE")...).Err()) + require.EqualValues(t, SimulateBitOp(ONE, vec...), rdb.Get(ctx, "target").Val()) + } + }) + t.Run("BITFIELD and BITFIELD_RO on string type", func(t *testing.T) { str := "zhe ge ren hen lan, shen me dou mei you liu xia." require.NoError(t, rdb.Set(ctx, "str", str, 0).Err())