Box bulk (many) put & get

Author: Ivan Dlugos

objectbox/box.py

@@ -22,33 +22,104 @@ def count(self, limit: int = 0) -> int:
         obx_box_count(self._c_box, limit, ctypes.byref(count))
         return int(count.value)
 
-    def put(self, object) -> int:
-        id = object_id = self._entity.get_object_id(object)
+    def put(self, *objects):
+        """Puts an object (or a list of objects) and returns its ID (or nothing for a list objects)"""
+
+        if len(objects) != 1:
+            self._put_many(objects)
+        elif isinstance(objects[0], list):
+            self._put_many(objects[0])
+        else:
+            return self._put_one(objects[0])
+
+    def _put_one(self, obj) -> int:
+        id = object_id = self._entity.get_object_id(obj)
 
         if not id:
             id = obx_box_id_for_put(self._c_box, 0)
 
-        data = self._entity.marshal(object, id)
+        data = self._entity.marshal(obj, id)
         obx_box_put(self._c_box, id, bytes(data), len(data), OBXPutMode_PUT)
 
         if id != object_id:
-            self._entity.set_object_id(object, id)
+            self._entity.set_object_id(obj, id)
+
         return id
 
+    def _put_many(self, objects) -> None:
+        # retrieve IDs from the objects (to distinguish new objects and updates)
+        new = {}
+        ids = {}
+        for k in range(len(objects)):
+            id = self._entity.get_object_id(objects[k])
+            if not id:
+                new[k] = 0
+            ids[k] = id
+
+        # acquire IDs for the new objects and set them
+        if len(new) > 0:
+            c_next_id = obx_id()
+            obx_box_ids_for_put(self._c_box, len(new), ctypes.byref(c_next_id))
+            next_id = c_next_id.value
+            for k in new.keys():
+                ids[k] = next_id
+                next_id += 1
+
+        # allocate a C bytes array structure where we will push the object data
+        # OBX_bytes_array with .count = len(objects)
+        c_bytes_array_p = obx_bytes_array_create(len(objects))
+
+        try:
+            # we need to keep the data around until put_many is executed because obx_bytes_array_set doesn't do a copy
+            data = {}
+            for k in range(len(objects)):
+                data[k] = bytes(self._entity.marshal(objects[k], ids[k]))
+                key = ctypes.c_size_t(k)
+
+                # OBX_bytes_array.data[k] = data
+                obx_bytes_array_set(c_bytes_array_p, key, data[k], len(data[k]))
+
+            c_ids = (obx_id * len(ids))(*ids.values())
+            obx_box_put_many(self._c_box, c_bytes_array_p, c_ids, OBXPutMode_PUT)
+
+        finally:
+            obx_bytes_array_free(c_bytes_array_p)
+
+        # assign new IDs on the object
+        for k in new.keys():
+            self._entity.set_object_id(objects[k], ids[k])
+
     def get(self, id: int):
         c_data = ctypes.c_void_p()
         c_size = ctypes.c_size_t()
         obx_box_get(self._c_box, id, ctypes.byref(c_data), ctypes.byref(c_size))
 
-        # TODO verify which of the following two approaches is better.
+        data = c_voidp_as_bytes(c_data, c_size.value)
+        return self._entity.unmarshal(data)
 
-        # slice the data from the pointer
-        # data = ctypes.cast(c_data, ctypes.POINTER(ctypes.c_char))[:c_size.value]
+    def get_all(self) -> list:
+        # OBX_bytes_array*
+        c_bytes_array_p = obx_box_get_all(self._c_box)
 
-        # create a memory view
-        data = memoryview(ctypes.cast(c_data, ctypes.POINTER(ctypes.c_ubyte * c_size.value))[0]).tobytes()
+        try:
+            # OBX_bytes_array
+            c_bytes_array = c_bytes_array_p.contents
 
-        return self._entity.unmarshal(data)
+            result = list()
+            for i in range(c_bytes_array.count):
+                # OBX_bytes
+                c_bytes = c_bytes_array.data[i]
+                data = c_voidp_as_bytes(c_bytes.data, c_bytes.size)
+                result.append(self._entity.unmarshal(data))
+
+            return result
+        finally:
+            obx_bytes_array_free(c_bytes_array_p)
 
     def remove(self, id: int):
-        obx_box_remove(self._c_box, id)
+        obx_box_remove(self._c_box, id)
+
+    def remove_all(self) -> int:
+        count = ctypes.c_uint64()
+        obx_box_remove_all(self._c_box, ctypes.byref(count))
+        return int(count.value)

objectbox/c.py

@@ -61,7 +61,7 @@ class OBX_bytes(ctypes.Structure):
 
 class OBX_bytes_array(ctypes.Structure):
     _fields_ = [
-        ('data', ctypes.POINTER(OBX_bytes)),
+        ('data', OBX_bytes_p),
         ('count', ctypes.c_size_t),
     ]
 
@@ -194,6 +194,18 @@ def c_str(string: str) -> ctypes.c_char_p:
     return string.encode('utf-8')
 
 
+def c_voidp_as_bytes(voidp, size):
+    # TODO verify which of the following two approaches is better. Performance-wise, it seems the same.
+
+    # slice the data from the pointer
+    # return ctypes.cast(voidp, ctypes.POINTER(ctypes.c_char))[:size]
+
+    # create a memory view
+    return memoryview(ctypes.cast(voidp, ctypes.POINTER(ctypes.c_ubyte * size))[0]).tobytes()
+
+
+
+
 # OBX_model* (void);
 obx_model_create = fn('obx_model_create', OBX_model_p, [])
 
@@ -233,21 +245,43 @@ def c_str(string: str) -> ctypes.c_char_p:
 obx_box_get = fn('obx_box_get', obx_err,
                  [OBX_box_p, obx_id, ctypes.POINTER(ctypes.c_void_p), ctypes.POINTER(ctypes.c_size_t)])
 
+# OBX_bytes_array* (OBX_box* box);
+obx_box_get_all = fn('obx_box_get_all', OBX_bytes_array_p, [OBX_box_p])
+
 # obx_id (OBX_box* box, obx_id id_or_zero);
 obx_box_id_for_put = fn('obx_box_id_for_put', obx_id, [OBX_box_p, obx_id])
 
+# obx_err (OBX_box* box, uint64_t count, obx_id* out_first_id);
+obx_box_ids_for_put = fn('obx_box_ids_for_put', obx_err, [OBX_box_p, ctypes.c_uint64, ctypes.POINTER(obx_id)])
+
 # obx_err (OBX_box* box, obx_id id, const void* data, size_t size, OBXPutMode mode);
 obx_box_put = fn('obx_box_put', obx_err, [OBX_box_p, obx_id, ctypes.c_void_p, ctypes.c_size_t, OBXPutMode])
 
+# obx_err (OBX_box* box, const OBX_bytes_array* objects, const obx_id* ids, OBXPutMode mode);
+obx_box_put_many = fn('obx_box_put_many', obx_err, [OBX_box_p, OBX_bytes_array_p, ctypes.POINTER(obx_id), OBXPutMode])
+
 # obx_err (OBX_box* box, obx_id id);
 obx_box_remove = fn('obx_box_remove', obx_err, [OBX_box_p, obx_id])
 
+# obx_err (OBX_box* box, uint64_t* out_count);
+obx_box_remove_all = fn('obx_box_remove_all', obx_err, [OBX_box_p, ctypes.POINTER(ctypes.c_uint64)])
+
 # obx_err (OBX_box* box, bool* out_is_empty);
 obx_box_is_empty = fn('obx_box_is_empty', obx_err, [OBX_box_p, ctypes.POINTER(ctypes.c_bool)])
 
 # obx_err obx_box_count(OBX_box* box, uint64_t limit, uint64_t* out_count);
 obx_box_count = fn('obx_box_count', obx_err, [OBX_box_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64)])
 
+# OBX_bytes_array* (size_t count);
+obx_bytes_array_create = fn('obx_bytes_array_create', OBX_bytes_array_p, [ctypes.c_size_t])
+
+# obx_err (OBX_bytes_array* array, size_t index, const void* data, size_t size);
+obx_bytes_array_set = fn('obx_bytes_array_set', obx_err,
+                         [OBX_bytes_array_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t])
+
+# void (OBX_bytes_array * array);
+obx_bytes_array_free = fn('obx_bytes_array_free', None, [OBX_bytes_array_p])
+
 OBXPropertyType_Bool = 1
 OBXPropertyType_Byte = 2
 OBXPropertyType_Short = 3

objectbox/model/entity.py

@@ -26,8 +26,8 @@ def __init__(self, cls, id: int, uid: int):
         self.id_property: Property = None
         self.fill_properties()
 
-    def __call__(self):
-        return self.cls()
+    def __call__(self, *args):
+        return self.cls(*args)
 
     def fill_properties(self):
         # TODO allow subclassing and support entities with __slots__ defined

tests/model.py

@@ -10,3 +10,6 @@ class TestEntity:
     float = Property(float, id=5, uid=1005)
     bytes = Property(bytes, id=6, uid=1006)
     transient: str = ""  # not "Property" so it's not stored
+
+    def __init__(self, string: str = ""):
+        self.str = string

tests/test_basics.py

@@ -29,7 +29,17 @@ def test_open():
     load_empty_test_objectbox()
 
 
-def test_box():
+def assert_equal(actual, expected):
+    """Check that two TestEntity objects have the same property data"""
+    assert actual.id == expected.id
+    assert isinstance(expected.bool, objectbox.model.Property) or actual.bool == expected.bool
+    assert isinstance(expected.int, objectbox.model.Property) or actual.int == expected.int
+    assert isinstance(expected.str, objectbox.model.Property) or actual.str == expected.str
+    assert isinstance(expected.float, objectbox.model.Property) or actual.float == expected.float
+    assert isinstance(expected.bytes, objectbox.model.Property) or actual.bytes == expected.bytes
+
+
+def test_box_basics():
     ob = load_empty_test_objectbox()
     box = objectbox.Box(ob, TestEntity)
 
@@ -60,18 +70,10 @@ def test_box():
     assert not box.is_empty()
     assert box.count() == 2
 
-    def verifyObject(read, object):
-        assert read.id == object.id
-        assert read.bool == object.bool
-        assert read.int == object.int
-        assert read.str == object.str
-        assert read.float == object.float
-        assert read.bytes == object.bytes
-        assert read.transient != object.transient  # !=
-
     # read
     read = box.get(object.id)
-    verifyObject(read, object)
+    assert_equal(read, object)
+    assert read.transient != object.transient  # !=
 
     # update
     object.str = "bar"
@@ -80,7 +82,7 @@ def verifyObject(read, object):
 
     # read again
     read = box.get(object.id)
-    verifyObject(read, object)
+    assert_equal(read, object)
 
     # remove
     box.remove(object.id)
@@ -91,3 +93,33 @@ def verifyObject(read, object):
         box.get(object.id)
 
 
+def test_box_bulk():
+    ob = load_empty_test_objectbox()
+    box = objectbox.Box(ob, TestEntity)
+
+    box.put(TestEntity("first"))
+
+    objects = [TestEntity("second"), TestEntity("third"), TestEntity("fourth"), box.get(1)]
+    box.put(objects)
+    assert box.count() == 4
+    assert objects[0].id == 2
+    assert objects[1].id == 3
+    assert objects[2].id == 4
+    assert objects[3].id == 1
+
+    assert_equal(box.get(objects[0].id), objects[0])
+    assert_equal(box.get(objects[1].id), objects[1])
+    assert_equal(box.get(objects[2].id), objects[2])
+    assert_equal(box.get(objects[3].id), objects[3])
+
+    objects_read = box.get_all()
+    assert len(objects_read) == 4
+    assert_equal(objects_read[0], objects[3])
+    assert_equal(objects_read[1], objects[0])
+    assert_equal(objects_read[2], objects[1])
+    assert_equal(objects_read[3], objects[2])
+
+    # remove all
+    removed = box.remove_all()
+    assert removed == 4
+    assert box.count() == 0