fix: use itertuples to ingest pandas dataframe to FeatureStore

#2032

Author: Rui Wang Napieralski

src/sagemaker/feature_store/feature_group.py

@@ -179,10 +179,12 @@ def _ingest_single_batch(
             end_index (int): ending position to ingest in this batch.
         """
         logger.info("Started ingesting index %d to %d", start_index, end_index)
-        for _, row in data_frame[start_index:end_index].iterrows():
+        for row in data_frame[start_index:end_index].itertuples(index=False):
             record = [
-                FeatureValue(feature_name=name, value_as_string=str(value))
-                for name, value in row.items()
+                FeatureValue(
+                    feature_name=data_frame.columns[index], value_as_string=str(row[index])
+                )
+                for index in range(len(row))
             ]
             sagemaker_session.put_record(
                 feature_group_name=feature_group_name, record=[value.to_dict() for value in record]

tests/integ/test_feature_store.py

@@ -105,6 +105,17 @@ def pandas_data_frame():
     return df
 
 
+@pytest.fixture
+def pandas_data_frame_without_string():
+    df = pd.DataFrame(
+        {
+            "feature1": pd.Series(np.arange(10), dtype="int64"),
+            "feature2": pd.Series([time.time()] * 10, dtype="float64"),
+        }
+    )
+    return df
+
+
 @pytest.fixture
 def record():
     return [
@@ -188,6 +199,34 @@ def test_create_feature_store(
     assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
 
 
+def test_ingest_without_string_feature(
+    feature_store_session,
+    role,
+    feature_group_name,
+    offline_store_s3_uri,
+    pandas_data_frame_without_string,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame_without_string)
+
+    with cleanup_feature_group(feature_group):
+        output = feature_group.create(
+            s3_uri=offline_store_s3_uri,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature2",
+            role_arn=role,
+            enable_online_store=True,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        ingestion_manager = feature_group.ingest(
+            data_frame=pandas_data_frame_without_string, max_workers=3, wait=False
+        )
+        ingestion_manager.wait()
+
+    assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
+
+
 def _wait_for_feature_group_create(feature_group: FeatureGroup):
     status = feature_group.describe().get("FeatureGroupStatus")
     while status == "Creating":