continue refactoring

Author: blink1073

.pre-commit-config.yaml

@@ -94,6 +94,11 @@ repos:
     stages: [manual]
     args: ["--no-strict-imports"]
 
+- repo: https://github.com/MarcoGorelli/cython-lint
+  rev: v0.16.2
+  hooks:
+  - id: cython-lint
+
 - repo: https://github.com/codespell-project/codespell
   rev: "v2.2.6"
   hooks:

bindings/python/pymongoarrow/api.py

@@ -88,7 +88,7 @@ def find_arrow_all(collection, query, *, schema=None, **kwargs):
     :Returns:
       An instance of class:`pyarrow.Table`.
     """
-    context = PyMongoArrowContext.from_schema(schema, codec_options=collection.codec_options)
+    context = PyMongoArrowContext(schema, codec_options=collection.codec_options)
 
     for opt in ("cursor_type",):
         if kwargs.pop(opt, None):
@@ -126,7 +126,7 @@ def aggregate_arrow_all(collection, pipeline, *, schema=None, **kwargs):
     :Returns:
       An instance of class:`pyarrow.Table`.
     """
-    context = PyMongoArrowContext.from_schema(schema, codec_options=collection.codec_options)
+    context = PyMongoArrowContext(schema, codec_options=collection.codec_options)
 
     if pipeline and ("$out" in pipeline[-1] or "$merge" in pipeline[-1]):
         msg = (

bindings/python/pymongoarrow/context.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from bson.codec_options import DEFAULT_CODEC_OPTIONS
-from pyarrow import ListArray, StructArray, Table, timestamp
+from pyarrow import ListArray, StructArray, Table
 from pyarrow.types import is_struct
 
 from pymongoarrow.types import _BsonArrowTypes, _get_internal_typemap
@@ -54,15 +53,14 @@
         _BsonArrowTypes.date64: Date64Builder,
         _BsonArrowTypes.null: NullBuilder,
     }
-
 except ImportError:
     pass
 
 
 class PyMongoArrowContext:
     """A context for converting BSON-formatted data to an Arrow Table."""
 
-    def __init__(self, schema, builder_map, codec_options=None):
+    def __init__(self, schema, codec_options=None):
         """Initialize the context.
 
         :Parameters:
@@ -71,90 +69,70 @@ def __init__(self, schema, builder_map, codec_options=None):
             :class:`~pymongoarrow.builders._BuilderBase` instances.
         """
         self.schema = schema
-        self.builder_map = builder_map
         if self.schema is None and codec_options is not 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):
-        """Initialize the context from a :class:`~pymongoarrow.schema.Schema`
-        instance.
-
-        :Parameters:
-          - `schema`: Instance of :class:`~pymongoarrow.schema.Schema`.
-          - `codec_options` (optional): An instance of
-            :class:`~bson.codec_options.CodecOptions`.
-        """
-        if schema is None:
-            return cls(schema, {}, codec_options)
-
-        builder_map = {}
-        tzinfo = codec_options.tzinfo
-        str_type_map = _get_internal_typemap(schema.typemap)
-        _parse_types(str_type_map, builder_map, tzinfo)
-        return cls(schema, builder_map)
+        self.manager = BuilderManager(self.schema is not None, self.tzinfo)
+        if self.schema is not None:
+            schema_map = {}
+            str_type_map = _get_internal_typemap(schema.typemap)
+            _parse_types(str_type_map, schema_map, self.tzinfo)
+            self.manager.parse_types(schema_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)
+        builder_map = self.manager.finish().copy()
 
-    @staticmethod
-    def _finish(builder_map, schema):
+        # Handle nested builders.
         to_remove = []
         # Traverse the builder map right to left.
         for key, value in reversed(builder_map.items()):
             field = key.decode("utf-8")
-            arr = value.finish()
             if isinstance(value, DocumentBuilder):
+                arr = value.finish()
                 full_names = [f"{field}.{name.decode('utf-8')}" for name in arr]
                 arrs = [builder_map[c.encode("utf-8")] for c in full_names]
                 builder_map[field] = StructArray.from_arrays(arrs, names=arr)
                 to_remove.extend(full_names)
             elif isinstance(value, ListBuilder):
-                child = field + "[]"
-                to_remove.append(child)
-                builder_map[key] = ListArray.from_arrays(arr, builder_map.get(child, []))
+                arr = value.finish()
+                child_name = field + "[]"
+                to_remove.append(child_name)
+                child = builder_map[child_name.encode("utf-8")]
+                builder_map[key] = ListArray.from_arrays(arr, child)
             else:
-                builder_map[key] = arr
+                builder_map[key] = value.finish()
 
         for field in to_remove:
             key = field.encode("utf-8")
             if key in builder_map:
                 del builder_map[key]
 
         arrays = list(builder_map.values())
-        if schema is not None:
-            return Table.from_arrays(arrays=arrays, schema=schema.to_arrow())
+        if self.schema is not None:
+            return Table.from_arrays(arrays=arrays, schema=self.schema.to_arrow())
         return Table.from_arrays(arrays=arrays, names=list(builder_map.keys()))
 
 
-def _parse_types(str_type_map, builder_map, tzinfo):
+def _parse_types(str_type_map, schema_map, tzinfo):
     for fname, (ftype, arrow_type) in str_type_map.items():
         builder_cls = _TYPE_TO_BUILDER_CLS[ftype]
         encoded_fname = fname.encode("utf-8")
-        # special-case initializing builders for parameterized types
-        if builder_cls == DatetimeBuilder:
-            if tzinfo is not None and arrow_type.tz is None:
-                arrow_type = timestamp(arrow_type.unit, tz=tzinfo)  # noqa: PLW2901
-            builder_map[encoded_fname] = DatetimeBuilder(dtype=arrow_type)
-        elif builder_cls == DocumentBuilder:
-            builder_map[encoded_fname] = DocumentBuilder()
+        schema_map[encoded_fname] = (arrow_type, builder_cls)
+
+        # special-case nested builders
+        if builder_cls == DocumentBuilder:
             # construct a sub type map here
             sub_type_map = {}
             for i in range(arrow_type.num_fields):
                 field = arrow_type[i]
                 sub_name = f"{fname}.{field.name}"
                 sub_type_map[sub_name] = field.type
             sub_type_map = _get_internal_typemap(sub_type_map)
-            _parse_types(sub_type_map, builder_map, tzinfo)
-            continue
         elif builder_cls == ListBuilder:
-            builder_map[encoded_fname] = ListBuilder()
             if is_struct(arrow_type.value_type):
                 # construct a sub type map here
                 sub_type_map = {}
@@ -163,10 +141,4 @@ def _parse_types(str_type_map, builder_map, tzinfo):
                     sub_name = f"{fname}[].{field.name}"
                     sub_type_map[sub_name] = field.type
                 sub_type_map = _get_internal_typemap(sub_type_map)
-                _parse_types(sub_type_map, builder_map, tzinfo)
-                continue
-        elif builder_cls == BinaryBuilder:
-            subtype = arrow_type.subtype
-            builder_map[encoded_fname] = BinaryBuilder(subtype)
-        else:
-            builder_map[encoded_fname] = builder_cls()
+                _parse_types(sub_type_map, schema_map, tzinfo)