chore: Fix test asserts for offline store write and improve some errors (feast-dev#2964)

* chore: Fix test asserts for offline store write and improve some error messages

Signed-off-by: Achal Shah <[email protected]>

* wait for write to finish

Signed-off-by: Achal Shah <[email protected]>

* wait for write to finish

Signed-off-by: Achal Shah <[email protected]>

* detailed error messages

Signed-off-by: Achal Shah <[email protected]>

* sort and reset index

Signed-off-by: Achal Shah <[email protected]>

* fix

Signed-off-by: Achal Shah <[email protected]>

Author: achals

sdk/python/feast/feature_store.py

@@ -1131,8 +1131,7 @@ def create_saved_dataset(
 
         if not from_.metadata:
             raise ValueError(
-                "RetrievalJob must contains metadata. "
-                "Use RetrievalJob produced by get_historical_features"
+                f"The RetrievalJob {type(from_)} must implement the metadata property."
             )
 
         dataset = SavedDataset(

sdk/python/feast/infra/offline_stores/bigquery.py

@@ -306,7 +306,7 @@ def write_logged_features(
                         file_obj=f,
                         destination=destination.table,
                         job_config=job_config,
-                    )
+                    ).result()
 
             return
 
@@ -319,7 +319,7 @@ def write_logged_features(
                 file_obj=parquet_temp_file,
                 destination=destination.table,
                 job_config=job_config,
-            )
+            ).result()
 
     @staticmethod
     def offline_write_batch(
@@ -373,7 +373,7 @@ def offline_write_batch(
                 file_obj=parquet_temp_file,
                 destination=feature_view.batch_source.table,
                 job_config=job_config,
-            )
+            ).result()
 
 
 class BigQueryRetrievalJob(RetrievalJob):

sdk/python/feast/usage.py

@@ -35,7 +35,7 @@
 USAGE_ENDPOINT = "https://usage.feast.dev"
 
 _logger = logging.getLogger(__name__)
-_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+_executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)
 
 _is_enabled = os.getenv(FEAST_USAGE, default=DEFAULT_FEAST_USAGE_VALUE) == "True"
 

sdk/python/tests/integration/offline_store/test_offline_write.py

@@ -126,11 +126,12 @@ def test_writing_consecutively_to_offline_store(environment, universal_data_sour
             "created": [ts, ts],
         },
     )
+    first_df = first_df.astype({"conv_rate": "float32", "acc_rate": "float32"})
     store.write_to_offline_store(
         driver_stats.name, first_df, allow_registry_cache=False
     )
 
-    after_write_df = store.get_historical_features(
+    after_write_df: pd.DataFrame = store.get_historical_features(
         entity_df=entity_df,
         features=[
             "driver_stats:conv_rate",
@@ -139,21 +140,26 @@ def test_writing_consecutively_to_offline_store(environment, universal_data_sour
         ],
         full_feature_names=False,
     ).to_df()
-
-    assert len(after_write_df) == len(first_df)
-    assert np.where(
-        after_write_df["conv_rate"].reset_index(drop=True)
-        == first_df["conv_rate"].reset_index(drop=True)
-    )
-    assert np.where(
-        after_write_df["acc_rate"].reset_index(drop=True)
-        == first_df["acc_rate"].reset_index(drop=True)
+    after_write_df = after_write_df.sort_values("event_timestamp").reset_index(
+        drop=True
     )
-    assert np.where(
-        after_write_df["avg_daily_trips"].reset_index(drop=True)
-        == first_df["avg_daily_trips"].reset_index(drop=True)
+
+    print(f"After: {after_write_df}\nFirst: {first_df}")
+    print(
+        f"After: {after_write_df['conv_rate'].reset_index(drop=True)}\nFirst: {first_df['conv_rate'].reset_index(drop=True)}"
     )
 
+    assert len(after_write_df) == len(first_df)
+    for field in ["conv_rate", "acc_rate", "avg_daily_trips"]:
+        assert np.equal(
+            after_write_df[field].reset_index(drop=True),
+            first_df[field].reset_index(drop=True),
+        ).all(), (
+            f"Field: {field}\n"
+            f"After: {after_write_df[field].reset_index(drop=True)}\n"
+            f"First: {first_df[field].reset_index(drop=True)}"
+        )
+
     second_df = pd.DataFrame.from_dict(
         {
             "event_timestamp": [ts + timedelta(hours=5), ts + timedelta(hours=6)],
@@ -164,6 +170,7 @@ def test_writing_consecutively_to_offline_store(environment, universal_data_sour
             "created": [ts, ts],
         },
     )
+    second_df = second_df.astype({"conv_rate": "float32", "acc_rate": "float32"})
 
     store.write_to_offline_store(
         driver_stats.name, second_df, allow_registry_cache=False
@@ -190,18 +197,17 @@ def test_writing_consecutively_to_offline_store(environment, universal_data_sour
         ],
         full_feature_names=False,
     ).to_df()
-
+    after_write_df = after_write_df.sort_values("event_timestamp").reset_index(
+        drop=True
+    )
     expected_df = pd.concat([first_df, second_df])
     assert len(after_write_df) == len(expected_df)
-    assert np.where(
-        after_write_df["conv_rate"].reset_index(drop=True)
-        == expected_df["conv_rate"].reset_index(drop=True)
-    )
-    assert np.where(
-        after_write_df["acc_rate"].reset_index(drop=True)
-        == expected_df["acc_rate"].reset_index(drop=True)
-    )
-    assert np.where(
-        after_write_df["avg_daily_trips"].reset_index(drop=True)
-        == expected_df["avg_daily_trips"].reset_index(drop=True)
-    )
+    for field in ["conv_rate", "acc_rate", "avg_daily_trips"]:
+        assert np.equal(
+            after_write_df[field].reset_index(drop=True),
+            expected_df[field].reset_index(drop=True),
+        ).all(), (
+            f"Field: {field}\n"
+            f"After: {after_write_df[field].reset_index(drop=True)}\n"
+            f"First: {expected_df[field].reset_index(drop=True)}"
+        )