ARROW-72 Run performance tests on prototype and evaluate our approach (#68)

Author: juliusgeo

bindings/python/benchmark.py

@@ -11,7 +11,13 @@
 import pyarrow
 import pymongo
 from bson import BSON, Int64, ObjectId
-from pymongoarrow.api import Schema, find_arrow_all, find_numpy_all, find_pandas_all
+from pymongoarrow.api import (
+    Schema,
+    find_arrow_all,
+    find_numpy_all,
+    find_pandas_all,
+    write,
+)
 
 assert pymongo.has_c()
 
@@ -25,6 +31,7 @@
 dtypes = {}
 schemas = {}
 raw_bsons = {}
+arrow_tables = {}
 
 
 def _setup():
@@ -82,6 +89,8 @@ def _setup():
 
     raw_bsons[SMALL] = raw_bson_small
     raw_bsons[LARGE] = raw_bson_large
+    arrow_tables[SMALL] = find_arrow_all(db[collection_names[SMALL]], {}, schema=schemas[SMALL])
+    arrow_tables[LARGE] = find_arrow_all(db[collection_names[LARGE]], {}, schema=schemas[LARGE])
 
 
 def _teardown():
@@ -143,6 +152,17 @@ def to_arrow(use_large):
     find_arrow_all(c, {}, schema=schema)
 
 
+@bench("insert_arrow")
+def insert_arrow(use_large):
+    write(db[collection_names[use_large]], arrow_tables[use_large])
+
+
+@bench("insert_conventional")
+def insert_conventional(use_large):
+    tab = arrow_tables[use_large].to_pylist()
+    db[collection_names[use_large]].insert_many(tab)
+
+
 parser = argparse.ArgumentParser(
     formatter_class=argparse.RawTextHelpFormatter,
     epilog="""

bindings/python/pymongoarrow/api.py

@@ -279,20 +279,21 @@ def write(collection, tabular):
         "insertedCount": 0,
     }
     tabular_gen = _tabular_generator(tabular)
-    while cur_offset < len(tabular):
+    tab_size = len(tabular)
+    while cur_offset < tab_size:
         cur_size = 0
         cur_batch = []
         i = 0
         while (
             cur_size <= _MAX_MESSAGE_SIZE
             and len(cur_batch) <= _MAX_WRITE_BATCH_SIZE
-            and cur_offset + i < len(tabular)
+            and cur_offset + i < tab_size
         ):
             enc_tab = RawBSONDocument(
                 encode(next(tabular_gen), codec_options=collection.codec_options)
             )
             cur_batch.append(enc_tab)
-            cur_size += len(enc_tab)
+            cur_size += len(enc_tab.raw)
             i += 1
         try:
             collection.insert_many(cur_batch)