write all NaN and NaT Dataframe created values as null

materialsproject · Feb 29, 2024 · 33806c1 · 33806c1
1 parent 767a024
commit 33806c1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/maggma/stores/open_data.py b/src/maggma/stores/open_data.py
@@ -4,6 +4,7 @@
 from typing import Dict, Generator, List, Optional, Tuple, Union
 
 import jsonlines
+import numpy as np
 import pandas as pd
 from boto3 import client as boto_client
 from botocore import UNSIGNED
@@ -625,6 +626,8 @@ def _write_doc_and_update_index(self, items: pd.DataFrame, index: pd.DataFrame)
         self.index.update(index)
 
     def _write_doc_to_s3(self, doc: pd.DataFrame, index: pd.DataFrame) -> None:
+        doc = doc.replace({pd.NaT: None}).replace({"NaT": None}).replace({np.NaN: None})
+
         string_io = StringIO()
         with jsonlines.Writer(string_io, dumps=json_util.dumps) as writer:
             for _, row in doc.iterrows():