Update type conversion from pandas to timestamp to support various the timestamp types (feast-dev#1603)

achals · web-flow · commit f29cb8cca5d8 · 2021-06-01T23:53:08.000-07:00
* Fix precommit to run black on commit

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Update type conversion from pandas to timestamp to support various the specifications

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Undo precommit changs

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* make format

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/data_source.py b/sdk/python/feast/data_source.py
@@ -598,8 +598,7 @@ def __init__(
         self._file_options = FileOptions(file_format=file_format, file_url=file_url)
 
         super().__init__(
-            event_timestamp_column
-            or self._infer_event_timestamp_column(r"timestamp\[\w\w\]"),
+            event_timestamp_column or self._infer_event_timestamp_column(r"^timestamp"),
             created_timestamp_column,
             field_mapping,
             date_partition_column,
@@ -744,7 +743,6 @@ def get_table_column_names_and_types(self) -> Iterable[Tuple[str, str]]:
         from google.cloud import bigquery
 
         client = bigquery.Client()
-        bq_columns_query = ""
         name_type_pairs = []
         if self.table_ref is not None:
             project_id, dataset_id, table_id = self.table_ref.split(".")
diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Union
 
@@ -438,8 +439,14 @@ def pa_to_value_type(pa_type: object):
 
 
 def pa_to_feast_value_type(value: Union[pa.lib.ChunkedArray, str]) -> ValueType:
+    value_type = (
+        value.type.__str__() if isinstance(value, pa.lib.ChunkedArray) else value
+    )
+
+    if re.match(r"^timestamp", value_type):
+        return ValueType.INT64
+
     type_map = {
-        "timestamp[ms]": ValueType.INT64,
         "int32": ValueType.INT32,
         "int64": ValueType.INT64,
         "double": ValueType.DOUBLE,
@@ -455,9 +462,7 @@ def pa_to_feast_value_type(value: Union[pa.lib.ChunkedArray, str]) -> ValueType:
         "list<item: binary>": ValueType.BYTES_LIST,
         "list<item: bool>": ValueType.BOOL_LIST,
     }
-    return type_map[
-        value.type.__str__() if isinstance(value, pa.lib.ChunkedArray) else value
-    ]
+    return type_map[value_type]
 
 
 def pa_column_to_timestamp_proto_column(column: pa.lib.ChunkedArray) -> List[Timestamp]:
