[feat]Rocksdb backend

WORKSPACE

@@ -54,6 +54,22 @@ http_archive(
     url = "https://github.com/sewenew/redis-plus-plus/archive/refs/tags/1.2.3.zip",
 )
 
+# Enable to build RocksDB from source.
+# http_archive(
+#     name = "rocksdb",
+#     build_file = "//build_deps/toolchains/rocksdb:rocksdb.BUILD",
+#     sha256 = "2df8f34a44eda182e22cf84dee7a14f17f55d305ff79c06fb3cd1e5f8831e00d",
+#     strip_prefix = "rocksdb-6.22.1",
+#     urls = [
+#         "https://github.com/facebook/rocksdb/archive/refs/tags/v6.22.1.tar.gz",
+#     ],
+# )
+new_local_repository(
+    name = "rocksdb",
+    build_file = "//build_deps/toolchains/rocksdb:rocksdb.BUILD",
+    path = "/usr/local"
+)
+
 tf_configure(
     name = "local_config_tf",
 )

build_deps/toolchains/rocksdb/rocksdb.BUILD

@@ -0,0 +1,33 @@
+load("@rules_foreign_cc//foreign_cc:defs.bzl", "make")
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])  # BSD
+
+filegroup(
+    name = "all_srcs",
+    srcs = glob(["**"]),
+    visibility = ["//visibility:public"],
+)
+
+# Enable this to compile RocksDB from source instead.
+#make(
+#    name = "rocksdb",
+#    args = [
+#        "EXTRA_CXXFLAGS=\"-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\"",
+#        "-j6",
+#    ],
+#    targets = ["static_lib", "install-static"],
+#    lib_source = "@rocksdb//:all_srcs",
+#    out_static_libs = ["librocksdb.a"],
+#)
+
+cc_library(
+    name = "rocksdb",
+    srcs = ["lib/librocksdb.a"],
+    includes = ["./include"],
+    hdrs = glob(["rocksdb/*.h"]),
+    visibility = ["//visibility:public"],
+)

tensorflow_recommenders_addons/dynamic_embedding/__init__.py

@@ -21,6 +21,9 @@
     'RedisTable',
     'RedisTableConfig',
     'RedisTableCreator',
+    'RocksDBTable',
+    'RocksDBTableConfig',
+    'RocksDBTableCreator',
     'Variable',
     'TrainableWrapper',
     'DynamicEmbeddingOptimizer',
@@ -50,11 +53,15 @@
     CuckooHashTableCreator,
     RedisTableConfig,
     RedisTableCreator,
+    RocksDBTableConfig,
+    RocksDBTableCreator,
 )
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.cuckoo_hashtable_ops import (
     CuckooHashTable,)
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.redis_table_ops import (
     RedisTable,)
+from tensorflow_recommenders_addons.dynamic_embedding.python.ops.rocksdb_table_ops import (
+    RocksDBTable,)
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.dynamic_embedding_ops import (
     embedding_lookup,)
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.dynamic_embedding_ops import (

tensorflow_recommenders_addons/dynamic_embedding/core/BUILD

@@ -68,6 +68,30 @@ custom_op_library(
     ],
 )
 
+custom_op_library(
+    name = "_rocksdb_table_ops.so",
+    srcs = [
+        "kernels/rocksdb_table_op.cc",
+        "kernels/rocksdb_table_op.h",
+        "ops/rocksdb_table_ops.cc",
+        "utils/types.h",
+        "utils/utils.h",
+    ],
+    # Hack: To allow allow locating <rocksdb/xyz.h>.
+    includes = [
+        ".",
+    ],
+    linkopts = [
+        "-L/usr/local/lib",
+        "-lbz2",
+        "-llz4",
+        "-lzstd",
+    ],
+    deps = [
+        "@rocksdb",
+    ],
+)
+
 custom_op_library(
     name = "_math_ops.so",
     srcs = [