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

NhuanTDBK · NHUAN.TRAN · web-flow · commit b48d36bee074 · 2022-09-22T13:38:38.000-07:00
* Filter subset features

Signed-off-by: NHUAN.TRAN &lt;nhuan.tran@onemount.com&gt;
Signed-off-by: NHUAN.TRAN &lt;NhuanTDBK@users.noreply.github.com&gt;

* Fix python lint

Signed-off-by: NHUAN.TRAN &lt;NhuanTDBK@users.noreply.github.com&gt;

Signed-off-by: NHUAN.TRAN &lt;nhuan.tran@onemount.com&gt;
Signed-off-by: NHUAN.TRAN &lt;NhuanTDBK@users.noreply.github.com&gt;
Co-authored-by: NHUAN.TRAN &lt;nhuan.tran@onemount.com&gt;
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()
 
