add optional convert param to write

Author: sleepyStick

bindings/python/pymongoarrow/api.py

@@ -35,7 +35,20 @@
 from numpy import ndarray
 from pyarrow import Schema as ArrowSchema
 from pyarrow import Table, timestamp
-from pyarrow.types import is_date32, is_date64
+from pyarrow.types import (
+    is_date32,
+    is_date64,
+    is_duration,
+    is_float16,
+    is_float32,
+    is_int8,
+    is_int16,
+    is_list,
+    is_uint8,
+    is_uint16,
+    is_uint32,
+    is_uint64,
+)
 from pymongo.common import MAX_WRITE_BATCH_SIZE
 
 from pymongoarrow.context import PyMongoArrowContext
@@ -475,14 +488,15 @@ def transform_python(self, value):
         return Decimal128(value)
 
 
-def write(collection, tabular, *, exclude_none: bool = False):
+def write(collection, tabular, *, exclude_none: bool = False, auto_convert: bool = True):
     """Write data from `tabular` into the given MongoDB `collection`.
 
     :Parameters:
       - `collection`: Instance of :class:`~pymongo.collection.Collection`.
         against which to run the operation.
       - `tabular`: A tabular data store to use for the write operation.
       - `exclude_none`: Whether to skip writing `null` fields in documents.
+      - `auto_convert` (optional): Whether to attempt a best-effort conversion of unsupported types.
 
     :Returns:
       An instance of :class:`result.ArrowWriteResult`.
@@ -500,9 +514,24 @@ def write(collection, tabular, *, exclude_none: bool = False):
             if is_date32(dtype) or is_date64(dtype):
                 changed = True
                 dtype = timestamp("ms")  # noqa: PLW2901
+            elif auto_convert:
+                if is_uint8(dtype) or is_uint16(dtype) or is_int8(dtype) or is_int16(dtype):
+                    changed = True
+                    dtype = pa.int32()  # noqa: PLW2901
+                elif is_uint32(dtype) or is_uint64(dtype) or is_duration(dtype):
+                    changed = True
+                    dtype = pa.int64()  # noqa: PLW2901
+                elif is_float16(dtype) or is_float32(dtype):
+                    changed = True
+                    dtype = pa.float64()  # noqa: PLW2901
             new_types.append(dtype)
         if changed:
-            cols = [tabular.column(i).cast(new_types[i]) for i in range(tabular.num_columns)]
+            cols = [
+                tabular.column(i).cast(new_types[i])
+                if not is_list(new_types[i])
+                else tabular.column(i)
+                for i in range(tabular.num_columns)
+            ]
             tabular = Table.from_arrays(cols, names=tabular.column_names)
         _validate_schema(tabular.schema.types)
     elif pd is not None and isinstance(tabular, pd.DataFrame):

bindings/python/test/test_arrow.py

@@ -19,7 +19,7 @@
 import threading
 import unittest
 import unittest.mock as mock
-from datetime import date, datetime
+from datetime import date, datetime, timedelta
 from pathlib import Path
 from test import client_context
 from test.utils import AllowListEventListener, NullsTestMixin
@@ -1095,22 +1095,22 @@ def alltypes_sample(self, size=10000, seed=0, categorical=False):
 
         np.random.seed(seed)
         arrays = {
-            "uint8": pa.array(np.arange(size, dtype=np.uint8), type=pa.int32()),
-            "uint16": pa.array(np.arange(size, dtype=np.uint16), type=pa.int32()),
-            "uint32": pa.array(np.arange(size, dtype=np.uint32), type=pa.int64()),
-            "uint64": pa.array(np.arange(size, dtype=np.uint64), type=pa.int64()),
-            "int8": pa.array(np.arange(size, dtype=np.int8), type=pa.int32()),
-            "int16": pa.array(np.arange(size, dtype=np.int16), type=pa.int32()),
+            "uint8": np.arange(size, dtype=np.uint8),
+            "uint16": np.arange(size, dtype=np.uint16),
+            "uint32": np.arange(size, dtype=np.uint32),
+            "uint64": np.arange(size, dtype=np.uint64),
+            "int8": np.arange(size, dtype=np.int8),
+            "int16": np.arange(size, dtype=np.int16),
             "int32": np.arange(size, dtype=np.int32),
             "int64": np.arange(size, dtype=np.int64),
-            "float16": pa.array(np.arange(size, dtype=np.float16), type=pa.float64()),
-            "float32": pa.array(np.arange(size, dtype=np.float32), type=pa.float64()),
-            "float64": pa.array(np.arange(size, dtype=np.float64), type=pa.float64()),
+            "float16": np.arange(size, dtype=np.float16),
+            "float32": np.arange(size, dtype=np.float32),
+            "float64": np.arange(size, dtype=np.float64),
             "bool": np.random.randn(size) > 0,
             "datetime_ms": np.arange("2016-01-01T00:00:00.001", size, dtype="datetime64[ms]"),
             "datetime_us": np.arange("2016-01-01T00:00:00.000001", size, dtype="datetime64[us]"),
             "datetime_ns": np.arange("2016-01-01T00:00:00.000000001", size, dtype="datetime64[ns]"),
-            "timedelta": pa.array(np.arange(size, dtype="timedelta64[s]"), type=pa.int64()),
+            "timedelta": np.arange(size, dtype="timedelta64[s]"),
             "str": pd.Series([str(x) for x in range(size)]),
             "empty_str": [""] * size,
             "str_with_nulls": [None] + [str(x) for x in range(size - 2)] + [None],
@@ -1181,6 +1181,10 @@ def compare_arrow_mongodb_data(self, arrow_table, mongo_data):
                     assert (
                         arrow_value == mongo_value
                     ), f"List mismatch in field '{column_name}' for row {row_idx}. Expected {arrow_value}, got {mongo_value}."
+                elif isinstance(arrow_value, timedelta):
+                    assert (
+                        arrow_value == timedelta(seconds=mongo_value)
+                    ), f"Timedelta mismatch in field '{column_name}' for row {row_idx}. Expected {arrow_value}, got {mongo_value}."
                 else:
                     assert (
                         arrow_value == mongo_value