|
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