When staging entities dataset timestamps are being truncated to ms (#1207)

pyalex · web-flow · commit 8df96a1eea0e · 2020-12-02T15:49:59.000+08:00
* floor event timestamp to ms

Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;

* fix bq staging

Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py
@@ -979,6 +979,12 @@ def get_historical_features(
             feature_refs, self.project
         )
 
+        assert all(ft.batch_source.created_timestamp_column for ft in feature_tables), (
+            "All BatchSources attached to retrieved FeatureTables "
+            "must have specified `created_timestamp_column` to be used in "
+            "historical dataset generation."
+        )
+
         if output_location is None:
             output_location = os.path.join(
                 self._config.get(opt.HISTORICAL_FEATURE_OUTPUT_LOCATION),
diff --git a/sdk/python/feast/staging/entities.py b/sdk/python/feast/staging/entities.py
@@ -29,6 +29,11 @@ def stage_entities_to_fs(
     entity_staging_uri = urlparse(os.path.join(staging_location, str(uuid.uuid4())))
     staging_client = get_staging_client(entity_staging_uri.scheme, config)
     with tempfile.NamedTemporaryFile() as df_export_path:
+        # prevent casting ns -> ms exception inside pyarrow
+        entity_source["event_timestamp"] = entity_source["event_timestamp"].dt.floor(
+            "ms"
+        )
+
         entity_source.to_parquet(df_export_path.name)
         bucket = (
             None if entity_staging_uri.scheme == "file" else entity_staging_uri.netloc
@@ -70,6 +75,9 @@ def stage_entities_to_bq(
         f"_entities_{datetime.now():%Y%m%d%H%M%s}",
     )
 
+    # prevent casting ns -> ms exception inside pyarrow
+    entity_source["event_timestamp"] = entity_source["event_timestamp"].dt.floor("ms")
+
     load_job: bigquery.LoadJob = bq_client.load_table_from_dataframe(
         entity_source, destination
     )
diff --git a/tests/e2e/test_historical_features.py b/tests/e2e/test_historical_features.py
@@ -36,11 +36,7 @@ def read_parquet(uri):
 
 
 def generate_data():
-    retrieval_date = (
-        datetime.utcnow()
-        .replace(hour=0, minute=0, second=0, microsecond=0)
-        .replace(tzinfo=None)
-    )
+    retrieval_date = datetime.utcnow().replace(tzinfo=None)
     retrieval_outside_max_age_date = retrieval_date + timedelta(1)
     event_date = retrieval_date - timedelta(2)
     creation_date = retrieval_date - timedelta(1)
