[FSTORE-1651] add dynamic queries to external feature groups

python/hsfs/core/arrow_flight_client.py

@@ -46,10 +46,8 @@
 if HAS_POLARS:
     import polars as pl
 
-
 _logger = logging.getLogger(__name__)
 
-
 _arrow_flight_instance = None
 
 
@@ -91,17 +89,15 @@ def _is_no_commits_found_error(exception):
 
 
 def _should_retry_healthcheck(exception):
-    return (
-        isinstance(exception, pyarrow._flight.FlightUnavailableError)
-        or isinstance(exception, pyarrow._flight.FlightTimedOutError)
+    return isinstance(exception, pyarrow._flight.FlightUnavailableError) or isinstance(
+        exception, pyarrow._flight.FlightTimedOutError
     )
 
 
 def _should_retry_certificate_registration(exception):
-    return (
-        _should_retry_healthcheck(exception)
-        or _is_feature_query_service_queue_full_error(exception)
-    )
+    return _should_retry_healthcheck(
+        exception
+    ) or _is_feature_query_service_queue_full_error(exception)
 
 
 # Avoid unnecessary client init
@@ -207,7 +203,9 @@ def __init__(self, disabled_for_session: bool = False):
 
         try:
             self._health_check()
-            if "get-version" in [action.type for action in self._connection.list_actions()]:
+            if "get-version" in [
+                action.type for action in self._connection.list_actions()
+            ]:
                 self._server_version = self._get_server_version()
             else:
                 self._server_version = None
@@ -522,6 +520,7 @@ def create_training_dataset(
             print("Error calling action:", e)
 
     def create_query_object(self, query, query_str, on_demand_fg_aliases=None):
+        selected_features = [feat.name for feat in query.features]
         if on_demand_fg_aliases is None:
             on_demand_fg_aliases = []
         features = {}
@@ -531,7 +530,9 @@ def create_query_object(self, query, query_str, on_demand_fg_aliases=None):
             fg_connector = _serialize_featuregroup_connector(
                 fg, query, on_demand_fg_aliases
             )
-            features[fg_name] = [feat.name for feat in fg.features]
+            features[fg_name] = [
+                feat.name for feat in fg.features if feat.name in selected_features
+            ]
             connectors[fg_name] = fg_connector
         filters = _serialize_filter_expression(query.filters, query)
 
@@ -592,11 +593,22 @@ def _serialize_featuregroup_connector(fg, query, on_demand_fg_aliases):
         connector["time_travel_type"] = None
         connector["type"] = fg.storage_connector.type
         connector["options"] = fg.storage_connector.connector_options()
-        connector["query"] = fg.query[:-1] if fg.query.endswith(";") else fg.query
+        connector["query"] = fg.query
         for on_demand_fg_alias in on_demand_fg_aliases:
+            # backend attaches dynamic query to on_demand_fg_alias.on_demand_feature_group.query if any
             if on_demand_fg_alias.on_demand_feature_group.name == fg.name:
+                connector["query"] = (
+                    on_demand_fg_alias.on_demand_feature_group.query
+                    if fg.query is None
+                    else fg.query
+                )
                 connector["alias"] = on_demand_fg_alias.alias
                 break
+        connector["query"] = (
+            connector["query"][:-1]
+            if connector["query"].endswith(";")
+            else connector["query"]
+        )
         if query._left_feature_group == fg:
             connector["filters"] = _serialize_filter_expression(
                 query._filter, query, True

python/hsfs/feature_group.py

@@ -442,9 +442,12 @@ def select_features(
         # Returns
             `Query`. A query object with all features of the feature group.
         """
-        query = self.select_except(
-            self.primary_key + self.partition_key + self.foreign_key + [self.event_time]
-        )
+        select_features = self.primary_key + self.foreign_key + [self.event_time]
+        if not isinstance(self, ExternalFeatureGroup):
+            select_features = select_features + self.partition_key
+
+        query = self.select_except(select_features)
+
         _logger.info(
             f"Using {[f.name for f in query.features]} as features for the query."
             "To include primary key and event time use `select_all`."

python/tests/core/test_arrow_flight_client.py

@@ -351,7 +351,7 @@ def test_construct_query_object_without_fs_excluded(self, mocker, backend_fixtur
         test_fg1 = feature_group.FeatureGroup.from_response_json(json1)
         mocker.patch("hsfs.constructor.query.Query.to_string", return_value="")
         mocker.patch("hsfs.constructor.query.Query._to_string", return_value="")
-        query = test_fg1.select_except(["intt"]).filter(Feature("intt") > 500)
+        query = test_fg1.filter(Feature("intt") > 500)
 
         # Act
         query_object = arrow_flight_client.get_instance().create_query_object(