feat: Filter subset features in postgres #3174 (#3203)

* Filter subset features Signed-off-by: NHUAN.TRAN <[email protected]> Signed-off-by: NHUAN.TRAN <[email protected]> * Fix python lint Signed-off-by: NHUAN.TRAN <[email protected]> Signed-off-by: NHUAN.TRAN <[email protected]> Signed-off-by: NHUAN.TRAN <[email protected]> Co-authored-by: NHUAN.TRAN <[email protected]>
feast-dev · Sep 22, 2022 · b48d36b · b48d36b
1 parent cb8db84
commit b48d36b
Showing 1 changed file with 24 additions and 11 deletions.
diff --git 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()