adding support for importing/exporting directories of docs

voxel51 · brimoor · Aug 30, 2022 · Aug 19, 2022 · Aug 19, 2022 · Aug 19, 2022
commit 613007ffda05762cf3f5d3a6ee4e9debc7e7e771
diff --git a/fiftyone/core/odm/database.py b/fiftyone/core/odm/database.py
@@ -534,19 +534,39 @@ def export_document(doc, json_path):
     etau.write_file(json_util.dumps(doc), json_path)
 
 
-def export_collection(docs, json_path, key="documents", num_docs=None):
+def export_collection(
+    docs,
+    json_dir_or_path,
+    key="documents",
+    patt="{idx:06d}-{id}.json",
+    num_docs=None,
+):
     """Exports the collection to disk in JSON format.
 
     Args:
-        docs: an iteraable containing the documents to export
-        json_path: the path to write the JSON file
+        docs: an iterable containing the documents to export
+        json_dir_or_path: the path to write a single JSON file containing the
+            entire collection, or a directory in which to write per-document
+            JSON files
         key ("documents"): the field name under which to store the documents
+            when ``json_path`` is a single JSON file
+        patt ("{idx:06d}-{id}.json"): a filename pattern to use when
+            ``json_path`` is a directory. The pattern may contain ``idx`` to
+            refer to the index of the document in ``docs`` or ``id`` to refer
+            to the document's ID
         num_docs (None): the total number of documents. If omitted, this must
             be computable via ``len(docs)``
     """
     if num_docs is None:
         num_docs = len(docs)
 
+    if json_dir_or_path.endswith(".json"):
+        _export_collection_single(docs, json_dir_or_path, key, num_docs)
+    else:
+        _export_collection_multi(docs, json_dir_or_path, patt, num_docs)
+
+
+def _export_collection_single(docs, json_path, key, num_docs):
     etau.ensure_basedir(json_path)
 
     with open(json_path, "w") as f:
@@ -560,6 +580,16 @@ def export_collection(docs, json_path, key="documents", num_docs=None):
         f.write("]}")
 
 
+def _export_collection_multi(docs, json_dir, patt, num_docs):
+    etau.ensure_dir(json_dir)
+
+    json_patt = os.path.join(json_dir, patt)
+    with fou.ProgressBar(total=num_docs, iters_str="docs") as pb:
+        for idx, doc in pb(enumerate(docs, 1)):
+            json_path = json_patt.format(idx=idx, id=str(doc["_id"]))
+            export_document(doc, json_path)
+
+
 def import_document(json_path):
     """Imports a document from JSON on disk.
 
@@ -573,27 +603,49 @@ def import_document(json_path):
         return json_util.loads(f.read())
 
 
-def import_collection(json_path, key="documents"):
+def import_collection(json_dir_or_path, key="documents"):
     """Imports the collection from JSON on disk.
 
     Args:
-        json_path: the path to the collection on disk
+        json_dir_or_path: the path to a JSON file on disk, or a directory
+            containing per-document JSON files
         key ("documents"): the field name under which the documents are stored
+            when ``json_path`` is a single JSON file
 
     Returns:
         a tuple of
 
-        -   the list of BSON documents
+        -   an iterable of BSON documents
         -   the number of documents
     """
+    if json_dir_or_path.endswith(".json"):
+        return _import_collection_single(json_dir_or_path, key)
+
+    return _import_collection_multi(json_dir_or_path)
+
+
+def _import_collection_single(json_path, key):
     with open(json_path, "r") as f:
         docs = json_util.loads(f.read()).get(key, [])
 
-    return docs, len(docs)
+    num_docs = len(docs)
 
+    return docs, num_docs
 
-def insert_documents(docs, coll, ordered=False):
-    """Inserts a list of documents into a collection.
+
+def _import_collection_multi(json_dir):
+    json_paths = [
+        p
+        for p in etau.list_files(json_dir, abs_paths=True)
+        if p.endswith(".json")
+    ]
+    docs = map(import_document, json_paths)
+
+    return docs, len(json_paths)
+
+
+def insert_documents(docs, coll, ordered=False, progress=False, num_docs=None):
+    """Inserts documents into a collection.
 
     The ``_id`` field of the input documents will be populated if it is not
     already set.
@@ -602,14 +654,39 @@ def insert_documents(docs, coll, ordered=False):
         docs: an iterable of BSON document dicts
         coll: a pymongo collection
         ordered (False): whether the documents must be inserted in order
+        progress (False): whether to render a progress bar tracking the
+            insertion
+        num_docs (None): the total number of documents. Only used when
+            ``progress=True``. If omitted, this will be computed via
+            ``len(docs)``, if possible
+
+    Returns:
+        a list of IDs of the inserted documents
     """
+    ids = []
+
     try:
-        for batch in fou.iter_batches(docs, 100000):  # mongodb limit
-            coll.insert_many(list(batch), ordered=ordered)
+        batcher = fou.DynamicBatcher(
+            docs,
+            target_latency=0.2,
+            init_batch_size=1,
+            max_batch_beta=2.0,
+            max_batch_size=100000,  # mongodb limit
+            progress=progress,
+            total=num_docs,
+        )
+
+        with batcher:
+            for batch in batcher:
+                batch = list(batch)
+                coll.insert_many(batch, ordered=ordered)
+                ids.extend(b["_id"] for b in batch)
     except BulkWriteError as bwe:
         msg = bwe.details["writeErrors"][0]["errmsg"]
         raise ValueError(msg) from bwe
 
+    return ids
+
 
 def bulk_write(ops, coll, ordered=False):
     """Performs a batch of write operations on a collection.