ARROW-167 int32 should coerce BSON types the same as int64 (#153)

Author: ShaneHarvey

bindings/python/pymongoarrow/lib.pyx

@@ -172,6 +172,8 @@ cdef void process_raw_bson_stream(const uint8_t * docstream, size_t length, obje
     cdef map[cstring, void *] builder_map
     cdef map[cstring, void*].iterator it
     cdef bson_subtype_t subtype
+    cdef int32_t val32
+    cdef int64_t val64
 
     cdef _ArrayBuilderBase builder = None
     cdef Int32Builder int32_builder
@@ -262,8 +264,24 @@ cdef void process_raw_bson_stream(const uint8_t * docstream, size_t length, obje
                 value_t = bson_iter_type(&doc_iter)
                 if ftype == BSON_TYPE_INT32:
                     int32_builder = builder
-                    if value_t == BSON_TYPE_INT32:
-                        int32_builder.append_raw(bson_iter_int32(&doc_iter))
+                    if (value_t == BSON_TYPE_INT32 or value_t == BSON_TYPE_BOOL):
+                        int32_builder.append_raw(bson_iter_as_int64(&doc_iter))
+                    elif value_t == BSON_TYPE_INT64:
+                        val64 = bson_iter_as_int64(&doc_iter)
+                        val32 = <int32_t> val64
+                        if val64 == val32:
+                            int32_builder.append_raw(val32)
+                        else:
+                            # Use append (not append_raw) to surface overflow errors.
+                            int32_builder.append(val64)
+                    elif value_t == BSON_TYPE_DOUBLE:
+                        # Treat nan as null.
+                        val = bson_iter_as_double(&doc_iter)
+                        if isnan(val):
+                            int32_builder.append_null()
+                        else:
+                            # Use append (not append_raw) to surface overflow errors.
+                            int32_builder.append(bson_iter_as_int64(&doc_iter))
                     else:
                         int32_builder.append_null()
                 elif ftype == BSON_TYPE_INT64:

bindings/python/test/test_arrow.py

@@ -487,7 +487,7 @@ def test_auto_schema_heterogeneous(self):
         self.coll.insert_many(data)
         for func in [find_arrow_all, aggregate_arrow_all]:
             out = func(self.coll, {} if func == find_arrow_all else []).drop(["_id"])
-            self.assertEqual(out["a"].to_pylist(), [1, None, None, 4])
+            self.assertEqual(out["a"].to_pylist(), [1, None, 1, 4])
 
     def test_auto_schema_tz(self):
         # Create table with random data of various types.
@@ -564,7 +564,7 @@ def test_malformed_embedded_documents(self):
             dict(data=dict(a=1, b=True)),
             dict(data=dict(a=1, b=True, c="bar")),
             dict(data=dict(a=1)),
-            dict(data=dict(a=True, b=False)),
+            dict(data=dict(a="str", b=False)),
         ]
         self.coll.drop()
         self.coll.insert_many(data)
@@ -590,6 +590,56 @@ def test_mixed_subtype(self):
         res = find_arrow_all(coll, {}, schema=schema)
         self.assertEqual(res["data"].to_pylist(), [Binary(b"1", 10), None])
 
+    def _test_mixed_types_int(self, inttype):
+        docs = [
+            {"a": 1},
+            {"a": 2.9},  # float should be truncated.
+            {"a": True},  # True should be 1.
+            {"a": False},  # False should be 0.
+            {"a": float("nan")},  # Should be null.
+            {"a": None},  # Should be null.
+            {},  # Should be null.
+            {"a": "string"},  # Should be null.
+        ]
+        self.coll.delete_many({})
+        self.coll.insert_many(docs)
+        table = find_arrow_all(self.coll, {}, projection={"_id": 0}, schema=Schema({"a": inttype}))
+        expected = Table.from_pylist(
+            [
+                {"a": 1},
+                {"a": 2},
+                {"a": 1},
+                {"a": 0},
+                {"a": None},
+                {"a": None},
+                {},
+                {"a": None},
+            ],
+            schema=ArrowSchema([field("a", inttype)]),
+        )
+        self.assertEqual(table, expected)
+
+    def test_mixed_types_int32(self):
+        self._test_mixed_types_int(int32())
+        # Value too large to fit in int32 should cause an overflow error.
+        self.coll.delete_many({})
+        self.coll.insert_one({"a": 2 << 34})
+        with self.assertRaises(OverflowError):
+            find_arrow_all(self.coll, {}, projection={"_id": 0}, schema=Schema({"a": int32()}))
+        # Test double overflowing int32
+        self.coll.delete_many({})
+        self.coll.insert_one({"a": float(2 << 34)})
+        with self.assertRaises(OverflowError):
+            find_arrow_all(self.coll, {}, projection={"_id": 0}, schema=Schema({"a": int32()}))
+
+    def test_mixed_types_int64(self):
+        self._test_mixed_types_int(int64())
+        # Test double overflowing int64
+        self.coll.delete_many({})
+        self.coll.insert_one({"a": float(2 << 65)})
+        with self.assertRaises(OverflowError):
+            find_arrow_all(self.coll, {}, projection={"_id": 0}, schema=Schema({"a": int32()}))
+
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):
     def run_find(self, *args, **kwargs):

bindings/python/test/test_numpy.py

@@ -235,7 +235,7 @@ def test_auto_schema_heterogeneous(self):
             with self.subTest(func.__name__):
                 out = func(self.coll, {} if func == find_numpy_all else [])
                 del out["_id"]
-                np.equal(out, [1.0, np.nan, np.nan, 4.0])
+                np.equal(out, [1.0, np.nan, 1.0, 4.0])
 
     def test_auto_schema_tz(self):
         schema = {

bindings/python/test/test_pandas.py

@@ -269,7 +269,7 @@ def test_auto_schema_heterogeneous(self):
         self.coll.insert_many(data)
         for func in [find_pandas_all, aggregate_pandas_all]:
             out = func(self.coll, {} if func == find_pandas_all else []).drop(columns=["_id"])
-            np.equal(out["a"], [1.0, np.nan, np.nan, 4.0])
+            np.equal(out["a"], [1.0, np.nan, 1, 4.0])
 
     def test_auto_schema_tz(self):
         schema = {