ARROW-94 Test for nulls (#83)

Author: bcwarner

bindings/python/test/test_arrow.py

@@ -15,8 +15,9 @@
 import unittest
 import unittest.mock as mock
 from test import client_context
-from test.utils import AllowListEventListener
+from test.utils import AllowListEventListener, TestNullsBase
 
+import pyarrow
 import pymongo
 from bson import Decimal128, ObjectId
 from pyarrow import Table, binary, bool_, decimal256, float64, int32, int64
@@ -384,3 +385,20 @@ def test_find_decimal128(self):
         )
         table = find_arrow_all(coll, {}, schema=schema)
         self.assertEqual(table, expected)
+
+
+class TestNulls(TestNullsBase):
+    def find_fn(self, coll, query, schema):
+        return find_arrow_all(coll, query, schema=schema)
+
+    def equal_fn(self, left, right):
+        self.assertEqual(left, right)
+
+    def table_from_dict(self, dict, schema=None):
+        return pyarrow.Table.from_pydict(dict, schema)
+
+    def assert_in_idx(self, table, col_name):
+        self.assertTrue(col_name in table.column_names)
+
+    def na_safe(self, atype):
+        return True

bindings/python/test/test_numpy.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # from datetime import datetime, timedelta
+import datetime
 import unittest
 from test import client_context
-from test.utils import AllowListEventListener
+from test.utils import AllowListEventListener, TestNullsBase
 from unittest import mock
 
 import numpy as np
@@ -242,3 +243,50 @@ def test_find_decimal128(self):
         }
         actual = find_numpy_all(self.coll, {}, schema=self.schema)
         self.assert_numpy_equal(actual, expected)
+
+
+# The spec for pyarrow says to_numpy is experimental, so we should expect
+# this to change in the future.
+class TestNulls(TestNullsBase, NumpyTestBase):
+    def table_from_dict(self, d, schema=None):
+        out = {}
+        for k, v in d.items():
+            if any(isinstance(x, int) for x in v) and None in v:
+                out[k] = np.array(v, dtype=np.float_)
+            else:
+                out[k] = np.array(v, dtype=np.dtype(type(v[0])))  # Infer
+        return out
+
+    def equal_fn(self, left, right):
+        left = np.nan_to_num(left)
+        right = np.nan_to_num(left)
+        self.assertTrue(np.all(np.equal(left, right)))
+
+    def find_fn(self, coll, query, schema=None):
+        return find_numpy_all(coll, query, schema=schema)
+
+    def assert_in_idx(self, table, col_name):
+        self.assertTrue(col_name in table)
+
+    pytype_tab_map = {
+        str: "str128",
+        int: ["int64", "float64"],
+        float: "float64",
+        datetime.datetime: "datetime64[ms]",
+        ObjectId: "object",
+        Decimal128: "object",
+        bool: "object",
+    }
+
+    pytype_writeback_exc_map = {
+        str: None,
+        int: None,
+        float: None,
+        datetime.datetime: ValueError,  # TypeError,
+        ObjectId: ValueError,  # TypeError,
+        Decimal128: ValueError,  # TypeError,
+        bool: None,
+    }
+
+    def na_safe(self, atype):
+        return atype != TestNulls.pytype_tab_map[str]

bindings/python/test/test_pandas.py

@@ -12,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # from datetime import datetime, timedelta
+import datetime
 import unittest
 import unittest.mock as mock
 from test import client_context
-from test.utils import AllowListEventListener
+from test.utils import AllowListEventListener, TestNullsBase
 
 import numpy as np
 import pandas as pd
+import pandas.testing
+import pyarrow
 from bson import Decimal128, ObjectId
 from pyarrow import decimal256, int32, int64
 from pymongo import DESCENDING, WriteConcern
@@ -219,3 +222,45 @@ def test_find_decimal128(self):
 
         table = find_pandas_all(self.coll, {}, schema=self.schema)
         pd.testing.assert_frame_equal(expected, table)
+
+
+class TestNulls(TestNullsBase):
+    def find_fn(self, coll, query, schema):
+        return find_pandas_all(coll, query, schema=schema)
+
+    def equal_fn(self, left, right):
+        left = left.fillna(0).replace(-0b1 << 63, 0)  # NaN is sometimes this
+        right = right.fillna(0).replace(-0b1 << 63, 0)
+        if type(left) == pandas.DataFrame:
+            pandas.testing.assert_frame_equal(left, right, check_dtype=False)
+        else:
+            pandas.testing.assert_series_equal(left, right, check_dtype=False)
+
+    def table_from_dict(self, dict, schema=None):
+        return pandas.DataFrame(dict)
+
+    def assert_in_idx(self, table, col_name):
+        self.assertTrue(col_name in table.columns)
+
+    pytype_tab_map = {
+        str: "object",
+        int: ["int64", "float64"],
+        float: "float64",
+        datetime.datetime: "datetime64[ns]",
+        ObjectId: "object",
+        Decimal128: "object",
+        bool: "object",
+    }
+
+    pytype_writeback_exc_map = {
+        str: None,
+        int: None,
+        float: None,
+        datetime.datetime: ValueError,
+        ObjectId: ValueError,
+        Decimal128: pyarrow.lib.ArrowInvalid,
+        bool: None,
+    }
+
+    def na_safe(self, atype):
+        return True

bindings/python/test/utils.py

@@ -11,9 +11,25 @@
 # 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.
+import datetime
+import unittest
 from collections import defaultdict
+from test import client_context
 
-from pymongo import monitoring
+import numpy as np
+import pyarrow
+from bson import Decimal128, ObjectId
+from pandas import isna
+from pyarrow import bool_, float64, int64, string, timestamp
+from pymongo import WriteConcern, monitoring
+from pymongoarrow.api import write
+from pymongoarrow.schema import Schema
+from pymongoarrow.types import (
+    _TYPE_NORMALIZER_FACTORY,
+    Decimal128StringType,
+    ObjectIdType,
+    _in_type_map,
+)
 
 
 class EventListener(monitoring.CommandListener):
@@ -54,3 +70,190 @@ def succeeded(self, event):
     def failed(self, event):
         if event.command_name in self.commands:
             super(AllowListEventListener, self).failed(event)
+
+
+class TestNullsBase(unittest.TestCase):
+    def find_fn(self, coll, query, schema):
+        raise NotImplementedError
+
+    def equal_fn(self, left, right):
+        raise NotImplementedError
+
+    def table_from_dict(self, dict, schema=None):
+        raise NotImplementedError
+
+    def assert_in_idx(self, table, col_name):
+        raise NotImplementedError
+
+    def na_safe(self, atype):
+        raise NotImplementedError
+
+    # Map Python types to constructors
+    pytype_cons_map = {
+        str: str,
+        int: int,
+        float: float,
+        datetime.datetime: lambda x: datetime.datetime(x + 1970, 1, 1),
+        ObjectId: lambda _: ObjectId(),
+        Decimal128: lambda x: Decimal128(str(x)),
+    }
+
+    # Map Python types to types for table we are comparing.
+    pytype_tab_map = {
+        str: string(),
+        int: int64(),
+        float: float64(),
+        datetime.datetime: timestamp("ms"),
+        ObjectId: ObjectIdType(),
+        Decimal128: Decimal128StringType(),
+        bool: bool_(),
+    }
+
+    pytype_writeback_exc_map = {
+        str: None,
+        int: None,
+        float: None,
+        datetime.datetime: None,
+        ObjectId: pyarrow.lib.ArrowInvalid,
+        Decimal128: pyarrow.lib.ArrowInvalid,
+        bool: None,
+    }
+
+    @classmethod
+    def setUpClass(cls):
+        if cls is TestNullsBase:
+            raise unittest.SkipTest("Base class")
+
+        if not client_context.connected:
+            raise unittest.SkipTest("cannot connect to MongoDB")
+        cls.cmd_listener = AllowListEventListener("find", "aggregate")
+        cls.getmore_listener = AllowListEventListener("getMore")
+        cls.client = client_context.get_client(
+            event_listeners=[cls.getmore_listener, cls.cmd_listener]
+        )
+
+        cls.oids = [ObjectId() for _ in range(4)]
+        cls.coll = cls.client.pymongoarrow_test.get_collection(
+            "test", write_concern=WriteConcern(w="majority")
+        )
+
+    def setUp(self):
+        self.coll.drop()
+
+        self.cmd_listener.reset()
+        self.getmore_listener.reset()
+
+    def assertType(self, obj1, arrow_type):
+        if isinstance(obj1, pyarrow.ChunkedArray):
+            if "storage_type" in dir(arrow_type):
+                self.assertEqual(obj1.type, arrow_type.storage_type)
+            else:
+                self.assertEqual(obj1.type, arrow_type)
+        else:
+            if isinstance(arrow_type, list):
+                self.assertTrue(obj1.dtype.name in arrow_type)
+            else:
+                self.assertEqual(obj1.dtype.name, arrow_type)
+
+    def test_int_handling(self):
+        # Default integral types
+        int_schema = Schema({"_id": ObjectIdType(), "int64": int64()})
+        int64_arr = [(i if (i % 2 == 0) else None) for i in range(len(self.oids))]
+        self.coll.insert_many(
+            [{"_id": self.oids[i], "int64": int64_arr[i]} for i in range(len(self.oids))]
+        )
+
+        table = self.find_fn(self.coll, {}, schema=int_schema)
+
+        # Resulting datatype should be float64 according to the spec for numpy
+        # and pandas.
+        atype = self.pytype_tab_map[int]
+        self.assertType(table["int64"], atype)
+
+        # Does it contain NAs where we expect?
+        self.assertTrue(np.all(np.equal(isna(int64_arr), isna(table["int64"]))))
+
+        # Write
+        self.coll.drop()
+        table_write = self.table_from_dict({"int64": int64_arr})
+
+        write(self.coll, table_write)
+        res_table = self.find_fn(self.coll, {}, schema=int_schema)
+
+        self.equal_fn(res_table["int64"], table_write["int64"])
+        self.assert_in_idx(res_table, "_id")
+        self.assertType(res_table["int64"], atype)
+
+    def test_all_types(self):
+        for t in self.pytype_tab_map.keys():
+            self.assertTrue(_in_type_map(_TYPE_NORMALIZER_FACTORY[t](0)))
+
+    def test_other_handling(self):
+        # Tests other types, which are treated differently in
+        # arrow/pandas/numpy.
+        for gen in [str, float, datetime.datetime, ObjectId, Decimal128]:
+            con_type = self.pytype_tab_map[gen]  # Arrow/Pandas/Numpy
+            pytype = TestNullsBase.pytype_tab_map[gen]  # Arrow type specifically
+
+            other_schema = Schema({"_id": ObjectIdType(), "other": pytype})
+            others = [
+                self.pytype_cons_map[gen](i) if (i % 2 == 0) else None
+                for i in range(len(self.oids))
+            ]
+
+            self.setUp()
+            self.coll.insert_many(
+                [{"_id": self.oids[i], "other": others[i]} for i in range(len(self.oids))]
+            )
+            table = self.find_fn(self.coll, {}, schema=other_schema)
+
+            # Resulting datatype should be str in this case
+
+            self.assertType(table["other"], con_type)
+            self.assertEqual(
+                self.na_safe(con_type), np.all(np.equal(isna(others), isna(table["other"])))
+            )
+
+            def writeback():
+                # Write
+                self.coll.drop()
+                table_write_schema = Schema({"other": pytype})
+                table_write_schema_arrow = (
+                    pyarrow.schema([("other", pytype)])
+                    if (gen in [str, float, datetime.datetime])
+                    else None
+                )
+
+                table_write = self.table_from_dict(
+                    {"other": others}, schema=table_write_schema_arrow
+                )
+
+                write(self.coll, table_write)
+                res_table = self.find_fn(self.coll, {}, schema=table_write_schema)
+                self.equal_fn(res_table, table_write)
+                self.assertType(res_table["other"], con_type)
+
+            # Do we expect an exception to be raised?
+            if self.pytype_writeback_exc_map[gen] is not None:
+                expected_exc = self.pytype_writeback_exc_map[gen]
+                with self.assertRaises(expected_exc):
+                    writeback()
+            else:
+                writeback()
+
+    def test_bool_handling(self):
+        atype = self.pytype_tab_map[bool]
+        bool_schema = Schema({"_id": ObjectIdType(), "bool_": bool_()})
+        bools = [True if (i % 2 == 0) else None for i in range(len(self.oids))]
+
+        self.coll.insert_many(
+            [{"_id": self.oids[i], "bool_": bools[i]} for i in range(len(self.oids))]
+        )
+
+        table = self.find_fn(self.coll, {}, schema=bool_schema)
+
+        # Resulting datatype should be object
+        self.assertType(table["bool_"], atype)
+
+        # Does it contain Nones where expected?
+        self.assertTrue(np.all(np.equal(isna(bools), isna(table["bool_"]))))