Remove use of GET_ITERATOR_CHUNK_SIZE in favor of MongoDB native document fetching policy

Author: Jibola

django_mongodb_backend/compiler.py

@@ -1,4 +1,5 @@
 import itertools
+import sys
 from collections import defaultdict
 
 from bson import SON, json_util
@@ -11,13 +12,16 @@
 from django.db.models.functions.math import Power
 from django.db.models.lookups import IsNull
 from django.db.models.sql import compiler
-from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, SINGLE
+from django.db.models.sql.constants import MULTI, SINGLE
 from django.db.models.sql.datastructures import BaseTable
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
 from .query import MongoQuery, wrap_database_errors
 
+# Maximum document batch size for MongoDB cursor responses.
+MAX_BATCH_SIZE_MB = 1024 * 1024 * 16
+
 
 class SQLCompiler(compiler.SQLCompiler):
     """Base class for all Mongo compilers."""
@@ -235,9 +239,7 @@ def pre_sql_setup(self, with_col_aliases=False):
         self.order_by_objs = [expr.replace_expressions(all_replacements) for expr, _ in order_by]
         return extra_select, order_by, group_by
 
-    def execute_sql(
-        self, result_type=MULTI, chunked_fetch=False, chunk_size=GET_ITERATOR_CHUNK_SIZE
-    ):
+    def execute_sql(self, result_type=MULTI, chunked_fetch=False, chunk_size=None):
         self.pre_sql_setup()
         try:
             query = self.build_query(
@@ -258,7 +260,8 @@ def execute_sql(
             else:
                 return self._make_result(obj, self.columns)
         # result_type is MULTI
-        cursor.batch_size(chunk_size)
+        # if chunk_size is not None:
+        #     cursor.batch_size(chunk_size)
         result = self.cursor_iter(cursor, chunk_size, self.columns)
         if not chunked_fetch:
             # If using non-chunked reads, read data into memory.
@@ -270,7 +273,7 @@ def results_iter(
         results=None,
         tuple_expected=False,
         chunked_fetch=False,
-        chunk_size=GET_ITERATOR_CHUNK_SIZE,
+        chunk_size=None,
     ):
         """
         Return an iterator over the results from executing query given
@@ -318,12 +321,25 @@ def _make_result(self, entity, columns):
             result.append(obj.get(name))
         return result
 
-    def cursor_iter(self, cursor, chunk_size, columns):
-        """Yield chunks of results from cursor."""
+    def cursor_iter(self, cursor, _, columns):
+        """
+        Yield chunks of results from cursor.
+        MongoDB ignores all chunk_size overrides. Cursor iteration abides by
+        MongoDB's default cursor batch size response.
+        Read more here: https://www.mongodb.com/docs/manual/core/cursors/#cursor-batches
+        """
         chunk = []
-        for row in cursor:
+        chunk_size = 101  # MongoDB's default initial batch size
+
+        for i, row in enumerate(cursor):
             chunk.append(self._make_result(row, columns))
+
             if len(chunk) == chunk_size:
+                if i == chunk_size - 1:  # First chunk complete
+                    # Using current row as representation, approximate
+                    # how many rows can fit in a 16MB payload (MongoDB batch_size max)
+                    # then set that as the new chunk size.
+                    chunk_size = MAX_BATCH_SIZE_MB // sys.getsizeof(row)
                 yield chunk
                 chunk = []
         yield chunk

django_mongodb_backend/query.py

@@ -76,7 +76,8 @@ def get_cursor(self):
         """
         self.compiler.connection.validate_no_broken_transaction()
         return self.compiler.collection.aggregate(
-            self.get_pipeline(), session=self.compiler.connection.session
+            self.get_pipeline(),
+            session=self.compiler.connection.session,
         )
 
     def get_pipeline(self):

docs/source/releases/5.1.x.rst

@@ -11,6 +11,10 @@ Django MongoDB Backend 5.1.x
   the ``base_field`` uses a database converter.
 - Fixed ``RecursionError`` when using ``Trunc`` database functions on non-MongoDB
   databases.
+- Removed ``GET_ITERATOR_CHUNK_SIZE`` as default from
+  :meth:`SQLCompiler.execute_sql() <django_mongodb_backend.compiler.SQLCompiler.execute_sql>`.
+  Cursor iteration will now use MongoDB default batch sizing.
+
 
 5.1.0 beta 3
 ============

docs/source/releases/5.2.x.rst

@@ -18,12 +18,16 @@ New features
   instance that may be of more than one model class.
 
 Bug fixes
----------
+--------
 
 - Fixed ``RecursionError`` when using ``Trunc`` database functions on non-MongoDB
   databases.
 - :meth:`QuerySet.explain() <django.db.models.query.QuerySet.explain>` now
   :ref:`returns a string that can be parsed as JSON <queryset-explain>`.
+- Removed ``GET_ITERATOR_CHUNK_SIZE`` as default from
+  :meth:`SQLCompiler.execute_sql() <django_mongodb_backend.compiler.SQLCompiler.execute_sql>`.
+  Cursor iteration will now use MongoDB default batch sizing.
+
 
 5.2.0 beta 1
 ============