chore: Update document test to add other feature data types and fix f…

…eature logging Signed-off-by: Francisco Javier Arceo <[email protected]>
feast-dev · Dec 22, 2024 · e3bb522 · e3bb522
1 parent 7df287e
commit e3bb522
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/sdk/python/tests/data/data_creator.py b/sdk/python/tests/data/data_creator.py
@@ -84,6 +84,8 @@ def get_feature_values_for_dtype(
 def create_document_dataset() -> pd.DataFrame:
     data = {
         "item_id": [1, 2, 3],
+        "string_feature": ["a", "b", "c"],
+        "float_feature": [1.0, 2.0, 3.0],
         "embedding_float": [[4.0, 5.0], [1.0, 2.0], [3.0, 4.0]],
         "embedding_double": [[4.0, 5.0], [1.0, 2.0], [3.0, 4.0]],
         "ts": [

diff --git a/sdk/python/tests/integration/feature_repos/universal/feature_views.py b/sdk/python/tests/integration/feature_repos/universal/feature_views.py
@@ -17,7 +17,7 @@
 from feast.data_source import DataSource, RequestSource
 from feast.feature_view_projection import FeatureViewProjection
 from feast.on_demand_feature_view import PandasTransformation, SubstraitTransformation
-from feast.types import Array, FeastType, Float32, Float64, Int32, Int64
+from feast.types import Array, FeastType, Float32, Float64, Int32, Int64, String
 from tests.integration.feature_repos.universal.entities import (
     customer,
     driver,
@@ -160,8 +160,20 @@ def create_item_embeddings_feature_view(source, infer_features: bool = False):
         schema=None
         if infer_features
         else [
-            Field(name="embedding_double", dtype=Array(Float64)),
-            Field(name="embedding_float", dtype=Array(Float32)),
+            Field(
+                name="embedding_double",
+                dtype=Array(Float64),
+                vector_index=True,
+                vector_search_metric="L2",
+            ),
+            Field(
+                name="embedding_float",
+                dtype=Array(Float32),
+                vector_index=True,
+                vector_search_metric="L2",
+            ),
+            Field(name="string_feature", dtype=String),
+            Field(name="float_feature", dtype=Float32),
         ],
         source=source,
         ttl=timedelta(hours=2),

diff --git a/sdk/python/tests/integration/offline_store/test_feature_logging.py b/sdk/python/tests/integration/offline_store/test_feature_logging.py
@@ -106,7 +106,15 @@ def retrieve():
     )
 
     persisted_logs = persisted_logs[expected_columns]
+
     logs_df = logs_df[expected_columns]
+
+    # Convert timezone-aware datetime values to naive datetime values
+    logs_df[LOG_TIMESTAMP_FIELD] = logs_df[LOG_TIMESTAMP_FIELD].dt.tz_localize(None)
+    persisted_logs[LOG_TIMESTAMP_FIELD] = persisted_logs[
+        LOG_TIMESTAMP_FIELD
+    ].dt.tz_localize(None)
+
     pd.testing.assert_frame_equal(
         logs_df.sort_values(REQUEST_ID_FIELD).reset_index(drop=True),
         persisted_logs.sort_values(REQUEST_ID_FIELD).reset_index(drop=True),