diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py index 18f3b189e7..144b242a1d 100644 --- a/sdk/python/feast/infra/online_stores/contrib/postgres.py +++ b/sdk/python/feast/infra/online_stores/contrib/postgres.py @@ -114,17 +114,30 @@ def online_read( ) ) - cur.execute( - sql.SQL( - """ - SELECT entity_key, feature_name, value, event_ts - FROM {} WHERE entity_key = ANY(%s); - """ - ).format( - sql.Identifier(_table_id(project, table)), - ), - (keys,), - ) + if not requested_features: + cur.execute( + sql.SQL( + """ + SELECT entity_key, feature_name, value, event_ts + FROM {} WHERE entity_key = ANY(%s); + """ + ).format( + sql.Identifier(_table_id(project, table)), + ), + (keys,), + ) + else: + cur.execute( + sql.SQL( + """ + SELECT entity_key, feature_name, value, event_ts + FROM {} WHERE entity_key = ANY(%s) and feature_name = ANY(%s); + """ + ).format( + sql.Identifier(_table_id(project, table)), + ), + (keys, requested_features), + ) rows = cur.fetchall()