basic support for column families

cozo-rocks/include/cozorocks.h

@@ -128,16 +128,28 @@ inline void write_status(RDB::Status &&rstatus, Status &status) {
 struct DBBridge {
     mutable std::unique_ptr <RDB::DB> inner;
 
+    mutable std::vector<RDB::ColumnFamilyHandle *> handles;
+
     DBBridge(RDB::DB *inner_) : inner(inner_) {}
 
+    inline std::unique_ptr<std::vector<std::string>> cf_names() const {
+        auto ret = std::make_unique<std::vector<std::string>>();
+        for (auto h : handles) {
+            ret->push_back(h->GetName());
+        }
+        return ret;
+    }
+
     inline void put(
             const WriteOptionsBridge &options,
+            std::size_t cf_id,
             rust::Slice<const uint8_t> key,
             rust::Slice<const uint8_t> val,
             Status &status
     ) const {
         write_status(
                 inner->Put(options.inner,
+                           handles[cf_id],
                            RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size()),
                            RDB::Slice(reinterpret_cast<const char *>(val.data()), val.size())),
                 status
@@ -146,13 +158,14 @@ struct DBBridge {
 
     inline std::unique_ptr <PinnableSliceBridge> get(
             const ReadOptionsBridge &options,
+            std::size_t cf_id,
             rust::Slice<const uint8_t> key,
             Status &status
     ) const {
         auto pinnable_val = std::make_unique<PinnableSliceBridge>();
         write_status(
                 inner->Get(options.inner,
-                           inner->DefaultColumnFamily(),
+                           handles[cf_id],
                            RDB::Slice(reinterpret_cast<const char *>(key.data()), key.size()),
                            &pinnable_val->inner),
                 status
@@ -170,7 +183,8 @@ inline std::unique_ptr <std::vector<std::string>> list_column_families(const Opt
     return column_families;
 }
 
-inline std::unique_ptr <DBBridge> open_db(const OptionsBridge &options, const rust::Slice<const uint8_t> path, Status &status) {
+inline std::unique_ptr <DBBridge>
+open_db(const OptionsBridge &options, const rust::Slice<const uint8_t> path, Status &status) {
     auto old_column_families = std::vector<std::string>();
     RDB::DB::ListColumnFamilies(options.inner,
                                 std::string(reinterpret_cast<const char *>(path.data()), path.size()),
@@ -195,5 +209,7 @@ inline std::unique_ptr <DBBridge> open_db(const OptionsBridge &options, const ru
                                   &handles,
                                   &db_ptr);
     write_status(std::move(s), status);
-    return std::unique_ptr<DBBridge>(new DBBridge(db_ptr));
+    auto ret = std::unique_ptr<DBBridge>(new DBBridge(db_ptr));
+    ret->handles = std::move(handles);
+    return ret;
 }

cozo-rocks/src/lib.rs

@@ -86,12 +86,14 @@ mod ffi {
         type DBBridge;
         fn list_column_families(options: &OptionsBridge, path: &[u8]) -> UniquePtr<CxxVector<CxxString>>;
         fn open_db(options: &OptionsBridge, path: &[u8], status: &mut Status) -> UniquePtr<DBBridge>;
-        fn put(self: &DBBridge, options: &WriteOptionsBridge, key: &[u8], val: &[u8], status: &mut Status);
-        fn get(self: &DBBridge, options: &ReadOptionsBridge, key: &[u8], status: &mut Status) -> UniquePtr<PinnableSliceBridge>;
+        fn cf_names(self: &DBBridge) -> UniquePtr<CxxVector<CxxString>>;
+        fn put(self: &DBBridge, options: &WriteOptionsBridge, cf_id: usize, key: &[u8], val: &[u8], status: &mut Status);
+        fn get(self: &DBBridge, options: &ReadOptionsBridge, cf_id: usize, key: &[u8], status: &mut Status) -> UniquePtr<PinnableSliceBridge>;
     }
 }
 
 
+use std::collections::BTreeMap;
 use cxx::UniquePtr;
 pub use ffi::*;
 
@@ -177,10 +179,10 @@ impl Default for WriteOptions {
 
 pub struct DB {
     bridge: UniquePtr<DBBridge>,
-    #[allow(dead_code)]
-    options: Options,
-    default_read_options: ReadOptions,
-    default_write_options: WriteOptions,
+    pub options: Options,
+    pub default_read_options: ReadOptions,
+    pub default_write_options: WriteOptions,
+    pub column_families: BTreeMap<String, usize>
 }
 
 fn get_path_bytes(path: &std::path::Path) -> &[u8] {
@@ -211,11 +213,13 @@ impl DB {
         );
 
         if status.code == StatusCode::kOk {
+            let column_families = bridge.cf_names().iter().enumerate().map(|(i, v)| (v.to_string_lossy().into_owned(), i)).collect();
             Ok(Self {
                 bridge,
                 default_read_options: ReadOptions::default(),
                 default_write_options: WriteOptions::default(),
                 options,
+                column_families
             })
         } else {
             Err(status)
@@ -225,9 +229,9 @@ impl DB {
     }
 
     #[inline]
-    pub fn put(&self, key: impl AsRef<[u8]>, val: impl AsRef<[u8]>, options: Option<&WriteOptions>) -> Result<Status, Status> {
+    pub fn put(&self, key: impl AsRef<[u8]>, val: impl AsRef<[u8]>, cf: usize, options: Option<&WriteOptions>) -> Result<Status, Status> {
         let mut status = Status::default();
-        self.bridge.put(&options.unwrap_or(&self.default_write_options).bridge,
+        self.bridge.put(&options.unwrap_or(&self.default_write_options).bridge, cf,
                         key.as_ref(), val.as_ref(),
                         &mut status);
         if status.code == StatusCode::kOk {
@@ -238,10 +242,10 @@ impl DB {
     }
 
     #[inline]
-    pub fn get(&self, key: impl AsRef<[u8]>, options: Option<&ReadOptions>) -> Result<Option<PinnableSlice>, Status> {
+    pub fn get(&self, key: impl AsRef<[u8]>, cf: usize, options: Option<&ReadOptions>) -> Result<Option<PinnableSlice>, Status> {
         let mut status = Status::default();
         let slice = self.bridge.get(
-            &options.unwrap_or(&self.default_read_options).bridge,
+            &options.unwrap_or(&self.default_read_options).bridge, cf,
             key.as_ref(), &mut status);
         match status.code {
             StatusCode::kOk => Ok(Some(PinnableSlice { bridge: slice })),

src/storage.rs

@@ -68,7 +68,6 @@ mod tests {
             .optimize_level_style_compaction()
             .set_create_if_missing(true)
             .set_comparator("cozo_comparator_v1", cozo_comparator_v1);
-        println!("{:?}", DB::list_column_families(&options, "xxyyzz.db"));
 
         let db = DB::open(options,
                           "xxyyzz.db").unwrap();
@@ -83,20 +82,20 @@ mod tests {
         builder.build_value(&Value::RefString("Another key"));
         let key2 = builder.get();
 
-        let val = db.get(&key, None).unwrap();
+        let val = db.get(&key, 0, None).unwrap();
         // let val = val.as_bytes();
         println!("before anything {}", val.is_none());
 
-        db.put(&key, "A motherfucking value!!! 👋👋👋", None).unwrap();
-        db.put(&key2, "Another motherfucking value!!! 👋👋👋", None).unwrap();
+        db.put(&key, "A motherfucking value!!! 👋👋👋", 0, None).unwrap();
+        db.put(&key2, "Another motherfucking value!!! 👋👋👋", 0, None).unwrap();
         // db.put("Yes man", "A motherfucking value!!! 👋👋👋", None).unwrap();
-        let val = db.get(&key, None).unwrap().unwrap();
+        let val = db.get(&key, 0, None).unwrap().unwrap();
         let val = val.as_bytes();
         println!("{}", from_utf8(val).unwrap());
-        let val = db.get(&key2, None).unwrap().unwrap();
+        let val = db.get(&key2, 0, None).unwrap().unwrap();
         let val = val.as_bytes();
         println!("{}", from_utf8(val).unwrap());
-        let val = db.get(&key, None).unwrap().unwrap();
+        let val = db.get(&key, 0, None).unwrap().unwrap();
         let val = val.as_bytes();
         println!("{}", from_utf8(val).unwrap());
     }