Fix missing timestamp data type in Timestream (#1881)

Co-authored-by: kukushking <[email protected]>

Author: jaidisido

awswrangler/_data_types.py

@@ -263,6 +263,12 @@ def pyarrow2timestream(dtype: pa.DataType) -> str:  # pylint: disable=too-many-b
         return "BOOLEAN"
     if pa.types.is_string(dtype):
         return "VARCHAR"
+    if pa.types.is_date(dtype):
+        return "DATE"
+    if pa.types.is_time(dtype):
+        return "TIME"
+    if pa.types.is_timestamp(dtype):
+        return "TIMESTAMP"
     raise exceptions.UnsupportedType(f"Unsupported Amazon Timestream measure type: {dtype}")
 
 

awswrangler/timestream.py

@@ -27,6 +27,14 @@ def _df2list(df: pd.DataFrame) -> List[List[Any]]:
     return parameters
 
 
+def _format_measure(measure_name: str, measure_value: Any, measure_type: str) -> Dict[str, str]:
+    return {
+        "Name": measure_name,
+        "Value": str(round(measure_value.timestamp() * 1_000) if measure_type == "TIMESTAMP" else measure_value),
+        "Type": measure_type,
+    }
+
+
 def _write_batch(
     database: str,
     table: str,
@@ -66,7 +74,7 @@ def _write_batch(
                 record["MeasureName"] = measure_cols_names[0]
                 record["MeasureValueType"] = "MULTI"
                 record["MeasureValues"] = [
-                    {"Name": measure_name, "Value": str(measure_value), "Type": measure_value_type}
+                    _format_measure(measure_name, measure_value, measure_value_type)
                     for measure_name, measure_value, measure_value_type in zip(
                         measure_cols_names, rec[measure_cols_loc:dimensions_cols_loc], measure_types
                     )

tests/test_timestream.py

@@ -1,3 +1,4 @@
+import datetime as dt
 import logging
 from datetime import datetime
 
@@ -312,3 +313,34 @@ def test_create_table_additional_kwargs(timestream_database_and_table, timestrea
     wr.timestream.delete_table(database=timestream_database_and_table, table=f"{timestream_database_and_table}_3")
     tables_in_db = wr.timestream.list_tables(database=timestream_database_and_table)
     assert f"{timestream_database_and_table}_3" not in tables_in_db
+
+
+def test_timestamp_measure_column(timestream_database_and_table):
+    df = pd.DataFrame(
+        {
+            "time": [datetime.now()] * 3,
+            "dim0": ["foo", "boo", "bar"],
+            "dim1": [1, 2, 3],
+            "measure_f": [1.1, 1.2, 1.3],
+            "measure_t": [datetime.now(dt.timezone.utc)] * 3,
+        }
+    )
+
+    rejected_records = wr.timestream.write(
+        df=df,
+        database=timestream_database_and_table,
+        table=timestream_database_and_table,
+        time_col="time",
+        measure_col=["measure_f", "measure_t"],
+        dimensions_cols=["dim0", "dim1"],
+    )
+    assert len(rejected_records) == 0
+
+    df = wr.timestream.query(
+        f"""
+        SELECT
+            *
+        FROM "{timestream_database_and_table}"."{timestream_database_and_table}"
+        """,
+    )
+    assert df["measure_t"].dtype == "datetime64[ns]"