|
49 | 49 | is_uint32, |
50 | 50 | is_uint64, |
51 | 51 | ) |
| 52 | +from pymongo.collection import Collection |
52 | 53 | from pymongo.common import MAX_WRITE_BATCH_SIZE |
53 | 54 |
|
54 | 55 | from pymongoarrow.context import PyMongoArrowContext |
@@ -488,7 +489,9 @@ def transform_python(self, value): |
488 | 489 | return Decimal128(value) |
489 | 490 |
|
490 | 491 |
|
491 | | -def write(collection, tabular, *, exclude_none: bool = False, auto_convert: bool = True): |
| 492 | +def write( |
| 493 | + collection: Collection, tabular, *, exclude_none: bool = False, auto_convert: bool = True |
| 494 | +): |
492 | 495 | """Write data from `tabular` into the given MongoDB `collection`. |
493 | 496 |
|
494 | 497 | :Parameters: |
@@ -559,10 +562,13 @@ def write(collection, tabular, *, exclude_none: bool = False, auto_convert: bool |
559 | 562 |
|
560 | 563 | # Add handling for special case types. |
561 | 564 | codec_options = collection.codec_options |
| 565 | + base_codecs = [] |
| 566 | + if hasattr(codec_options.type_registry, "codecs"): |
| 567 | + base_codecs = codec_options.type_registry.codecs |
562 | 568 | if pd is not None: |
563 | | - type_registry = TypeRegistry([_PandasNACodec(), _DecimalCodec()]) |
| 569 | + type_registry = TypeRegistry([*base_codecs, _PandasNACodec(), _DecimalCodec()]) |
564 | 570 | else: |
565 | | - type_registry = TypeRegistry([_DecimalCodec()]) |
| 571 | + type_registry = TypeRegistry([*base_codecs, _DecimalCodec()]) |
566 | 572 | codec_options = codec_options.with_options(type_registry=type_registry) |
567 | 573 |
|
568 | 574 | while cur_offset < tab_size: |
|
0 commit comments