feat: Add snowflake support for date & number with scale (#3148)

sfc-gh-madkins · web-flow · commit 50e8755d41ca · 2022-08-29T15:03:38.000-07:00
Signed-off-by: Miles Adkins &lt;miles.adkins@snowflake.com&gt;

Signed-off-by: Miles Adkins &lt;miles.adkins@snowflake.com&gt;
diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py
@@ -334,7 +334,11 @@ def generate_snowflake_materialization_query(
             )
 
             if feature_value_type_name == "UNIX_TIMESTAMP":
-                feature_sql = f'{feature_sql}(DATE_PART(EPOCH_NANOSECOND, "{feature.name}")) AS "{feature.name}"'
+                feature_sql = f'{feature_sql}(DATE_PART(EPOCH_NANOSECOND, "{feature.name}"::TIMESTAMP_LTZ)) AS "{feature.name}"'
+            elif feature_value_type_name == "DOUBLE":
+                feature_sql = (
+                    f'{feature_sql}("{feature.name}"::DOUBLE) AS "{feature.name}"'
+                )
             else:
                 feature_sql = f'{feature_sql}("{feature.name}") AS "{feature.name}"'
 
diff --git a/sdk/python/feast/infra/offline_stores/snowflake_source.py b/sdk/python/feast/infra/offline_stores/snowflake_source.py
@@ -264,18 +264,17 @@ def get_table_column_names_and_types(
                             ]
                         else:
                             raise NotImplementedError(
-                                "Numbers larger than INT64 are not supported"
+                                "NaNs or Numbers larger than INT64 are not supported"
                             )
                 else:
-                    raise NotImplementedError(
-                        "The following Snowflake Data Type is not supported: DECIMAL -- Convert to DOUBLE"
-                    )
-            elif row["type_code"] in [3, 5, 9, 10, 12]:
+                    row["snowflake_type"] = "NUMBERwSCALE"
+
+            elif row["type_code"] in [5, 9, 10, 12]:
                 error = snowflake_unsupported_map[row["type_code"]]
                 raise NotImplementedError(
                     f"The following Snowflake Data Type is not supported: {error}"
                 )
-            elif row["type_code"] in [1, 2, 4, 6, 7, 8, 11, 13]:
+            elif row["type_code"] in [1, 2, 3, 4, 6, 7, 8, 11, 13]:
                 row["snowflake_type"] = snowflake_type_code_map[row["type_code"]]
             else:
                 raise NotImplementedError(
@@ -291,6 +290,7 @@ def get_table_column_names_and_types(
     0: "NUMBER",
     1: "DOUBLE",
     2: "VARCHAR",
+    3: "DATE",
     4: "TIMESTAMP",
     6: "TIMESTAMP_LTZ",
     7: "TIMESTAMP_TZ",
@@ -300,7 +300,6 @@ def get_table_column_names_and_types(
 }
 
 snowflake_unsupported_map = {
-    3: "DATE -- Convert to TIMESTAMP",
     5: "VARIANT -- Try converting to VARCHAR",
     9: "OBJECT -- Try converting to VARCHAR",
     10: "ARRAY -- Try converting to VARCHAR",
diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py
@@ -625,8 +625,10 @@ def snowflake_type_to_feast_value_type(snowflake_type: str) -> ValueType:
         "VARCHAR": ValueType.STRING,
         "NUMBER32": ValueType.INT32,
         "NUMBER64": ValueType.INT64,
+        "NUMBERwSCALE": ValueType.DOUBLE,
         "DOUBLE": ValueType.DOUBLE,
         "BOOLEAN": ValueType.BOOL,
+        "DATE": ValueType.UNIX_TIMESTAMP,
         "TIMESTAMP": ValueType.UNIX_TIMESTAMP,
         "TIMESTAMP_TZ": ValueType.UNIX_TIMESTAMP,
         "TIMESTAMP_LTZ": ValueType.UNIX_TIMESTAMP,
