From 64daa07a65cd45f23dc32ea8ab67ea1f1cc98c1d Mon Sep 17 00:00:00 2001
From: Ashton Sidhu <ashton.sidhu@elastic.co>
Date: Tue, 4 Jan 2022 23:03:49 -0500
Subject: [PATCH] Using the 'date' field for datetime64+timezone columns

---
 eland/field_mappings.py              | 5 ++++-
 tests/dataframe/test_utils_pytest.py | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/eland/field_mappings.py b/eland/field_mappings.py
index c7817a7..92ac7f7 100644
--- a/eland/field_mappings.py
+++ b/eland/field_mappings.py
@@ -34,6 +34,7 @@ import numpy as np
 import pandas as pd  # type: ignore
 from pandas.core.dtypes.common import (  # type: ignore
     is_bool_dtype,
+    is_datetime64_any_dtype,
     is_datetime_or_timedelta_dtype,
     is_float_dtype,
     is_integer_dtype,
@@ -487,7 +488,7 @@ class FieldMappings:
         int64 int int_, int8, int16, int32, int64, uint8, uint16, uint32, uint64 Integer numbers
         float64 float float_, float16, float32, float64 Floating point numbers
         bool bool bool_ True/False values
-        datetime64 NA datetime64[ns] Date and time values
+        datetime64 NA datetime64[ns] datetime64[ns, TIMEZONE] Date and time values
         timedelta[ns] NA NA Differences between two datetimes
         category NA NA Finite list of text values
         ```
@@ -505,6 +506,8 @@ class FieldMappings:
             es_dtype = "keyword"
         elif is_datetime_or_timedelta_dtype(pd_dtype):
             es_dtype = "date"
+        elif is_datetime64_any_dtype(pd_dtype):
+            es_dtype = "date"
         else:
             warnings.warn(
                 f"No mapping for pd_dtype: [{pd_dtype}], using default mapping"
diff --git a/tests/dataframe/test_utils_pytest.py b/tests/dataframe/test_utils_pytest.py
index bc6b4ea..d14c332 100644
--- a/tests/dataframe/test_utils_pytest.py
+++ b/tests/dataframe/test_utils_pytest.py
@@ -37,6 +37,7 @@ class TestDataFrameUtils(TestData):
                 "E": [1.0, 2.0, 3.0],
                 "F": False,
                 "G": [1, 2, 3],
+                "H": pd.Timestamp("20190102", tz="UTC"),
             },
             index=["0", "1", "2"],
         )
@@ -51,6 +52,7 @@ class TestDataFrameUtils(TestData):
                     "E": {"type": "double"},
                     "F": {"type": "boolean"},
                     "G": {"type": "long"},
+                    "H": {"type": "date"},
                 }
             }
         }