From 9b168fb48daf326bec74a2dcf4b368a073ce038b Mon Sep 17 00:00:00 2001
From: davitbzh <davit.bzhalava@gmail..com>
Date: Tue, 3 Dec 2024 21:19:37 +0100
Subject: [PATCH] format

---
 python/hsfs/feature_group.py       | 38 +++++++++++-------------------
 python/hsfs/feature_store.py       |  6 +++++
 python/tests/test_feature_group.py | 14 +++++++----
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/python/hsfs/feature_group.py b/python/hsfs/feature_group.py
index d426bcf21..8562924a8 100644
--- a/python/hsfs/feature_group.py
+++ b/python/hsfs/feature_group.py
@@ -320,36 +320,26 @@ def select_all(
         # Returns
             `Query`. A query object with all features of the feature group.
         """
-        if include_event_time and include_primary_key and include_foreign_key and include_partition_key:
+        removed_keys = []
+
+        if not include_event_time:
+            removed_keys += [self.event_time]
+        if not include_primary_key:
+            removed_keys += self.primary_key
+        if not include_foreign_key:
+            removed_keys += self.foreign_key
+        if not include_partition_key:
+            removed_keys += self.partition_key
+
+        if removed_keys:
+            return self.select_except(removed_keys)
+        else:
             return query.Query(
                 left_feature_group=self,
                 left_features=self._features,
                 feature_store_name=self._feature_store_name,
                 feature_store_id=self._feature_store_id,
             )
-        elif include_event_time:
-            return self.select_except(
-                self.primary_key + self.foreign_key + self.partition_key
-            )
-        elif include_primary_key:
-            return self.select_except(
-                self.foreign_key + self.partition_key + [self.event_time]
-            )
-        elif include_foreign_key:
-            return self.select_except(
-                self.primary_key + self.partition_key + [self.event_time]
-            )
-        elif include_partition_key:
-            return self.select_except(
-                self.primary_key + self.foreign_key + [self.event_time]
-            )
-        else:
-            return self.select_except(
-                self.primary_key
-                + self.partition_key
-                + self.foreign_key
-                + [self.event_time]
-            )
 
     def select_features(
         self,
diff --git a/python/hsfs/feature_store.py b/python/hsfs/feature_store.py
index 825dcf7c5..5a36e85a4 100644
--- a/python/hsfs/feature_store.py
+++ b/python/hsfs/feature_store.py
@@ -588,6 +588,9 @@ def plus_two(value):
                 feature group. This primary key can be a composite key of multiple
                 features and will be used as joining key, if not specified otherwise.
                 Defaults to empty list `[]`, and the feature group won't have any primary key.
+            foreign_key: A list of feature names to be used as foreign key for the feature group.
+                Foreign key is referencing the primary key of another feature group and can be used as joining key.
+                Defaults to empty list `[]`, and the feature group won't have any foreign key.
             embedding_index: [`EmbeddingIndex`](./embedding_index_api.md). If an embedding index is provided,
                 vector database is used as online feature store. This enables similarity search by
                 using [`find_neighbors`](./feature_group_api.md#find_neighbors).
@@ -1238,6 +1241,9 @@ def get_or_create_spine_group(
                 spine group. This primary key can be a composite key of multiple
                 features and will be used as joining key, if not specified otherwise.
                 Defaults to empty list `[]`, and the spine group won't have any primary key.
+            foreign_key: A list of feature names to be used as foreign key for the feature group.
+                Foreign key is referencing the primary key of another feature group and can be used as joining key.
+                Defaults to empty list `[]`, and the feature group won't have any foreign key.
             event_time: Optionally, provide the name of the feature containing the event
                 time for the features in this spine group. If event_time is set
                 the spine group can be used for point-in-time joins. Defaults to `None`.
diff --git a/python/tests/test_feature_group.py b/python/tests/test_feature_group.py
index fca9953e8..cdd0a6129 100644
--- a/python/tests/test_feature_group.py
+++ b/python/tests/test_feature_group.py
@@ -350,14 +350,20 @@ def test_select_all(self):
     def test_select_all_exclude_pk(self):
         query = test_feature_group.select_all(include_primary_key=False)
         features = query.features
-        assert len(features) == 3
-        assert set([f.name for f in features]) == {"ts", "f1", "f2"}
+        assert len(features) == 4
+        assert set([f.name for f in features]) == {"ts", "fk", "f1", "f2"}
+
+    def test_select_all_exclude_fk(self):
+        query = test_feature_group.select_all(include_foreign_key=False)
+        features = query.features
+        assert len(features) == 4
+        assert set([f.name for f in features]) == {"f1", "f2", "pk", "ts"}
 
     def test_select_all_exclude_ts(self):
         query = test_feature_group.select_all(include_event_time=False)
         features = query.features
-        assert len(features) == 3
-        assert set([f.name for f in features]) == {"pk", "f1", "f2"}
+        assert len(features) == 4
+        assert set([f.name for f in features]) == {"pk", "fk", "f1", "f2"}
 
     def test_select_all_exclude_pk_ts(self):
         query = test_feature_group.select_all(