diff --git a/src/redis-roaring.c b/src/redis-roaring.c index 7aa1ef3..63230db 100644 --- a/src/redis-roaring.c +++ b/src/redis-roaring.c @@ -22,9 +22,11 @@ void BitmapFree(void* value); /* === Bitmap type methods === */ void BitmapRdbSave(RedisModuleIO* rdb, void* value) { Bitmap* bitmap = value; - size_t serialized_max_size = roaring_bitmap_size_in_bytes(bitmap); + roaring_bitmap_shrink_to_fit(bitmap); + roaring_bitmap_run_optimize(bitmap); + size_t serialized_max_size = roaring_bitmap_portable_size_in_bytes(bitmap); char* serialized_bitmap = RedisModule_Alloc(serialized_max_size); - size_t serialized_size = roaring_bitmap_serialize(bitmap, serialized_bitmap); + size_t serialized_size = roaring_bitmap_portable_serialize(bitmap, serialized_bitmap); RedisModule_SaveStringBuffer(rdb, serialized_bitmap, serialized_size); RedisModule_Free(serialized_bitmap); } @@ -36,7 +38,7 @@ void* BitmapRdbLoad(RedisModuleIO* rdb, int encver) { } size_t size; char* serialized_bitmap = RedisModule_LoadStringBuffer(rdb, &size); - Bitmap* bitmap = roaring_bitmap_deserialize(serialized_bitmap); + Bitmap* bitmap = roaring_bitmap_portable_deserialize(serialized_bitmap); RedisModule_Free(serialized_bitmap); return bitmap; } @@ -926,5 +928,81 @@ int RedisModule_OnLoad(RedisModuleCtx* ctx) { return REDISMODULE_ERR; } + +/** + * R.Export + * */ +int RExportBinCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 2) { + return RedisModule_WrongArity(ctx); + } + // Detect Key is exsit + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); + int type = RedisModule_KeyType(key); + if (type == REDISMODULE_KEYTYPE_EMPTY) { + return RedisModule_ReplyWithError(ctx, "bitmap key not exist"); + } + + // Detect value is bitmap type + if (RedisModule_ModuleTypeGetType(key) != BitmapType) { + return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + } + + RedisModule_AutoMemory(ctx); + + Bitmap *bitmap = RedisModule_ModuleTypeGetValue(key); + roaring_bitmap_shrink_to_fit(bitmap); + roaring_bitmap_run_optimize(bitmap); + size_t serialized_max_size = roaring_bitmap_portable_size_in_bytes(bitmap); + char *serialized_bitmap = RedisModule_Alloc(serialized_max_size); + size_t serialized_size = roaring_bitmap_portable_serialize(bitmap, serialized_bitmap); + + // Redis will handle free-memroy phase + RedisModule_ReplyWithStringBuffer(ctx, serialized_bitmap, serialized_size); + return REDISMODULE_OK; +} + +/** + * R.ImportBin + * */ +int RImportBinCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 3) { + return RedisModule_WrongArity(ctx); + } + size_t len; + const char *bin = RedisModule_StringPtrLen(argv[2], &len); + if (len < 10) { + return RedisModule_ReplyWithError(ctx, "bad binary data for roaring"); + } + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + int type = RedisModule_KeyType(key); + if (type != REDISMODULE_KEYTYPE_EMPTY && + RedisModule_ModuleTypeGetType(key) != BitmapType) { + return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + } + + RedisModule_AutoMemory(ctx); + + Bitmap *newbitmap = roaring_bitmap_portable_deserialize(bin); + uint64_t count; + if (type == REDISMODULE_KEYTYPE_EMPTY) { + count = bitmap_get_cardinality(newbitmap); + RedisModule_ModuleTypeSetValue(key, BitmapType, newbitmap); + } else { + Bitmap* bitmap = RedisModule_ModuleTypeGetValue(key); + roaring_bitmap_or_inplace(bitmap, newbitmap); + count = bitmap_get_cardinality(bitmap); + bitmap_free(newbitmap); + } + RedisModule_ReplicateVerbatim(ctx); + RedisModule_ReplyWithLongLong(ctx, count); + return REDISMODULE_OK; +} + if (RedisModule_CreateCommand(ctx, "R.EXPORT", RExportBinCommand, "readonly", 1, 1, 1) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + if (RedisModule_CreateCommand(ctx, "R.IMPORT", RImportBinCommand, "write", 0, 0, 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } return REDISMODULE_OK; }