More refactor and optimization

Author: blink1073

bindings/python/pymongoarrow/api.py

@@ -38,11 +38,6 @@
 from pymongoarrow.schema import Schema
 from pymongoarrow.types import _validate_schema, get_numpy_type
 
-try:  # noqa: SIM105
-    from pymongoarrow.lib import process_bson_stream
-except ImportError:
-    pass
-
 __all__ = [
     "aggregate_arrow_all",
     "find_arrow_all",
@@ -108,7 +103,7 @@ def find_arrow_all(collection, query, *, schema=None, **kwargs):
 
     raw_batch_cursor = collection.find_raw_batches(query, **kwargs)
     for batch in raw_batch_cursor:
-        process_bson_stream(batch, context)
+        context.process_bson_stream(batch)
 
     return context.finish()
 
@@ -152,7 +147,7 @@ def aggregate_arrow_all(collection, pipeline, *, schema=None, **kwargs):
 
     raw_batch_cursor = collection.aggregate_raw_batches(pipeline, **kwargs)
     for batch in raw_batch_cursor:
-        process_bson_stream(batch, context)
+        context.process_bson_stream(batch)
 
     return context.finish()
 

bindings/python/pymongoarrow/context.py

@@ -21,6 +21,7 @@
     from pymongoarrow.lib import (
         BinaryBuilder,
         BoolBuilder,
+        BuilderManager,
         CodeBuilder,
         Date32Builder,
         Date64Builder,
@@ -75,6 +76,7 @@ def __init__(self, schema, builder_map, codec_options=None):
             self.tzinfo = codec_options.tzinfo
         else:
             self.tzinfo = None
+        self.manager = BuilderManager(builder_map, self.schema is not None, self.tzinfo)
 
     @classmethod
     def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
@@ -95,6 +97,9 @@ def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
         _parse_types(str_type_map, builder_map, tzinfo)
         return cls(schema, builder_map)
 
+    def process_bson_stream(self, stream):
+        self.manager.process_bson_stream(stream, len(stream))
+
     def finish(self):
         return self._finish(self.builder_map, self.schema)