fix: Support pandas mode in feature builder and fix dask column extraction

cutoutsy · cutoutsy · commit be89314e2d0a · 2026-04-16T15:43:10.000+08:00
Signed-off-by: cutoutsy &lt;cutoutsy@gmail.com&gt;
diff --git a/sdk/python/feast/infra/compute_engines/feature_builder.py b/sdk/python/feast/infra/compute_engines/feature_builder.py
@@ -158,7 +158,7 @@ def get_column_info(
         # we need to read ALL source columns, not just the output feature columns.
         # This is specifically for transformations that create new columns or need raw data.
         mode = getattr(getattr(view, "feature_transformation", None), "mode", None)
-        if mode == "ray" or getattr(mode, "value", None) == "ray":
+        if mode in ("ray", "pandas") or getattr(mode, "value", None) in ("ray", "pandas"):
             # Signal to read all columns by passing empty list for feature_cols
             # The transformation will produce the output columns defined in the schema
             feature_cols = []
diff --git a/sdk/python/feast/infra/offline_stores/dask.py b/sdk/python/feast/infra/offline_stores/dask.py
@@ -377,7 +377,10 @@ def evaluate_func():
                 df[DUMMY_ENTITY_ID] = DUMMY_ENTITY_VAL
                 columns_to_extract.add(DUMMY_ENTITY_ID)
 
-            return df[list(columns_to_extract)].persist()
+            if feature_name_columns:
+                df = df[list(columns_to_extract)]
+
+            return df.persist()
 
         # When materializing a single feature view, we don't need full feature names. On demand transforms aren't materialized
         return DaskRetrievalJob(
@@ -483,9 +486,10 @@ def evaluate_func():
                 columns_to_extract.add(DUMMY_ENTITY_ID)
             # TODO: Decides if we want to field mapping for pull_latest_from_table_or_query
             # This is default for other offline store.
-            df = df[list(columns_to_extract)]
-            df.persist()
-            return df
+            if feature_name_columns:
+                df = df[list(columns_to_extract)]
+                
+            return df.persist()
 
         # When materializing a single feature view, we don't need full feature names. On demand transforms aren't materialized
         return DaskRetrievalJob(
