ARROW-127 Time Zones are Not Properly Handled by Auto Schema (#106)

* ARROW-127 Time Zones are Not Properly Handled by Auto Schema

* clean up handling of tz info

* undo change to test

Author: blink1073

.github/workflows/test-python.yml

@@ -80,7 +80,7 @@ jobs:
         run: |
           # Install the test deps
           LIBBSON_INSTALL_DIR=$(pwd)/libbson python -m pip install -vv -e ".[test]"
-          LD_LIBRARY_PATH=$(pwd)/libbson/lib python -m unittest discover test -v
+          PYTHONWARNINGS=error LD_LIBRARY_PATH=$(pwd)/libbson/lib python -m unittest discover test -v
       - name: Check the manifest
         run: |
           pip install check-manifest

bindings/python/pymongoarrow/context.py

@@ -65,7 +65,7 @@ def from_schema(cls, schema, codec_options=DEFAULT_CODEC_OPTIONS):
             :class:`~bson.codec_options.CodecOptions`.
         """
         if schema is None:
-            return cls(schema, {})
+            return cls(schema, {}, codec_options)
 
         builder_map = {}
         str_type_map = _get_internal_typemap(schema.typemap)

bindings/python/pymongoarrow/lib.pyx

@@ -111,7 +111,7 @@ def process_bson_stream(bson_stream, context):
 
                     builder_type = _builder_type_map[ftype]
                     if builder_type == DatetimeBuilder and context.tzinfo is not None:
-                        arrow_type = timestamp(arrow_type.unit, tz=context.tzinfo)
+                        arrow_type = timestamp('ms', tz=context.tzinfo)
                         builder_map[key] = builder_type(dtype=arrow_type)
                     else:
                         builder_map[key] = builder_type()

bindings/python/pymongoarrow/types.py

@@ -75,7 +75,10 @@ def _is_decimal128_str(obj):
     Int64: lambda _: int64(),
     float: lambda _: float64(),
     int: lambda _: int64(),
-    datetime: lambda _: timestamp("ms"),  # TODO: add tzinfo support
+    # Note: we cannot infer a timezone form a raw datetime class,
+    # if a timezone is preferred then a timestamp with tz information
+    # must be used directly.
+    datetime: lambda _: timestamp("ms"),
     ObjectId: lambda _: ObjectIdType(),
     Decimal128: lambda _: Decimal128StringType(),
     str: lambda _: string(),
@@ -88,7 +91,7 @@ def _is_decimal128_str(obj):
     _atypes.is_int64: np.int64,
     _atypes.is_float64: np.float64,
     _atypes.is_timestamp: "datetime64[ms]",
-    _is_objectid: np.object,
+    _is_objectid: object,
     _atypes.is_string: np.str_,
     _atypes.is_boolean: np.bool_,
 }

bindings/python/test/test_arrow.py

@@ -373,7 +373,7 @@ def test_auto_schema_tz(self):
             ArrowSchema(
                 {
                     "bool": bool_(),
-                    "dt": timestamp("ms"),
+                    "dt": timestamp("ms", "US/Eastern"),
                     "string": string(),
                 }
             ),
@@ -388,6 +388,7 @@ def test_auto_schema_tz(self):
                 self.coll.with_options(codec_options=codec_options),
                 {} if func == find_arrow_all else [],
             ).drop(["_id"])
+            self.assertEqual(out["dt"].type.tz, "US/Eastern")
             self.assertEqual(data, out)
 
 

bindings/python/test/test_numpy.py

@@ -254,9 +254,7 @@ def test_auto_schema_tz(self):
         data = {
             "string": [str(i) for i in range(3)],
             "bool": [True for _ in range(3)],
-            "dt": [
-                datetime.datetime(1970 + i, 1, 1, tzinfo=timezone("US/Eastern")) for i in range(3)
-            ],
+            "dt": [datetime.datetime(1970 + i, 1, 1) for i in range(3)],
         }
         data = self.schemafied_ndarray_dict(data, schema)
         self.coll.drop()
@@ -270,6 +268,10 @@ def test_auto_schema_tz(self):
                     {} if func == find_numpy_all else [],
                 )
                 del out["_id"]
+                # Note: Numpy does not handle timezones.
+                assert out["dt"].dtype == np.dtype("<M8[ns]")
+                assert data["dt"].dtype == np.dtype("<M8[ms]")
+                out.pop("dt")
                 self.assert_numpy_equal(data, out)
 
 

bindings/python/test/test_pandas.py

@@ -146,7 +146,7 @@ def test_write_schema_validation(self):
             Schema(arrow_schema),
         )
 
-        schema = {"_id": "int32", "data": np.ubyte()}
+        schema = {"_id": "int32", "data": bytes}
         data = pd.DataFrame(
             data={"_id": [i for i in range(2)], "data": [i for i in range(2)]}
         ).astype(schema)
@@ -223,7 +223,7 @@ def test_auto_schema_heterogeneous(self):
     def test_auto_schema_tz(self):
         schema = {
             "bool": "bool",
-            "dt": "datetime64[ns]",
+            "dt": "datetime64[ns, US/Eastern]",
             "string": "str",
         }
         data = pd.DataFrame(