Merge branch 'google:main' into main

google · Jul 4, 2023 · d0a65f8 · d0a65f8
2 parents 47544ab + ed62877
commit d0a65f8
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 11 deletions.
diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py
@@ -223,6 +223,9 @@ def extract_rows(self, uri: str, coordinates: t.List[t.Dict]) -> t.Iterator[t.Di
                 row = {n: to_json_serializable_type(ensure_us_time_resolution(v.values))
                        for n, v in row_ds.data_vars.items()}
 
+                # Serialize coordinates.
+                it = {k: to_json_serializable_type(v) for k, v in it.items()}
+
                 # Add indexed coordinates.
                 row.update(it)
                 # Add un-indexed coordinates.

diff --git a/weather_mv/loader_pipeline/util.py b/weather_mv/loader_pipeline/util.py
@@ -204,14 +204,11 @@ def get_coordinates(ds: xr.Dataset, uri: str = '') -> t.Iterator[t.Dict]:
     """Generates normalized coordinate dictionaries that can be used to index Datasets with `.loc[]`."""
     # Creates flattened iterator of all coordinate positions in the Dataset.
     #
-    # Coordinates have been pre-processed to remove NaNs and to format datetime objects
-    # to ISO format strings.
-    #
-    # Example: (-108.0, 49.0, '2018-01-02T22:00:00+00:00')
+    # Example: (-108.0, 49.0, datetime.datetime('2018-01-02T22:00:00+00:00'))
     coords = itertools.product(
         *(
             (
-                to_json_serializable_type(v)
+                v
                 for v in ensure_us_time_resolution(ds[c].variable.values).tolist()
             )
             for c in ds.coords.indexes
@@ -220,7 +217,7 @@ def get_coordinates(ds: xr.Dataset, uri: str = '') -> t.Iterator[t.Dict]:
     # Give dictionary keys to a coordinate index.
     #
     # Example:
-    #   {'longitude': -108.0, 'latitude': 49.0, 'time': '2018-01-02T23:00:00+00:00'}
+    #   {'longitude': -108.0, 'latitude': 49.0, 'time': datetime.datetime('2018-01-02T23:00:00+00:00')}
     idx = 0
     total_coords = math.prod(ds.coords.dims.values())
     for idx, it in enumerate(coords):

diff --git a/weather_mv/loader_pipeline/util_test.py b/weather_mv/loader_pipeline/util_test.py
@@ -38,7 +38,9 @@ def test_gets_indexed_coordinates(self):
         ds = xr.open_dataset(self.test_data_path)
         self.assertEqual(
             next(get_coordinates(ds)),
-            {'latitude': 49.0, 'longitude': -108.0, 'time': '2018-01-02T06:00:00+00:00'}
+            {'latitude': 49.0,
+             'longitude':-108.0,
+             'time': datetime.fromisoformat('2018-01-02T06:00:00+00:00').replace(tzinfo=None)}
         )
 
     def test_no_duplicate_coordinates(self):
@@ -89,12 +91,24 @@ def test_get_coordinates(self):
             actual,
             [
                 [
-                    {'longitude': -108.0, 'latitude': 49.0, 'time': '2018-01-02T06:00:00+00:00'},
-                    {'longitude': -108.0, 'latitude': 49.0, 'time': '2018-01-02T07:00:00+00:00'},
-                    {'longitude': -108.0, 'latitude': 49.0, 'time': '2018-01-02T08:00:00+00:00'},
+                    {'longitude': -108.0,
+                     'latitude': 49.0,
+                     'time': datetime.fromisoformat('2018-01-02T06:00:00+00:00').replace(tzinfo=None)
+                    },
+                    {'longitude': -108.0,
+                     'latitude': 49.0,
+                     'time': datetime.fromisoformat('2018-01-02T07:00:00+00:00').replace(tzinfo=None)
+                    },
+                    {'longitude': -108.0,
+                     'latitude': 49.0,
+                     'time': datetime.fromisoformat('2018-01-02T08:00:00+00:00').replace(tzinfo=None)
+                    },
                 ],
                 [
-                    {'longitude': -108.0, 'latitude': 49.0, 'time': '2018-01-02T09:00:00+00:00'}
+                    {'longitude': -108.0,
+                     'latitude': 49.0,
+                     'time': datetime.fromisoformat('2018-01-02T09:00:00+00:00').replace(tzinfo=None)
+                    }
                 ]
             ]
         )