Merge pull request #220 from dataiku/feature/sc-81874-add-a-method-to-un-set-a-dataset-as-a-feature

Feature store public API

Author: lpenet

dataikuapi/dss/dataset.py

@@ -737,6 +737,25 @@ def add_time_partitioning_dimension(self, dim_name, period="DAY"):
     def add_raw_schema_column(self, column):
         self.settings["schema"]["columns"].append(column)
 
+    @property
+    def is_feature_group(self):
+        """
+        Indicates whether the Dataset is defined as a Feature Group, available in the Feature Store.
+
+        :rtype: bool
+        """
+        return self.settings["featureGroup"]
+
+    def set_feature_group(self, status):
+        """
+        (Un)sets the dataset as a Feature Group, available in the Feature Store.
+        Changes of this property will be applied when calling :meth:`save` and require the "Manage Feature Store" permission.
+
+        :param status: whether the dataset should be defined as a feature group
+        :type status: bool
+        """
+        self.settings["featureGroup"] = status
+
     def save(self):
         self.dataset.client._perform_empty(
                 "PUT", "/projects/%s/datasets/%s" % (self.dataset.project_key, self.dataset.dataset_name),

dataikuapi/dss/feature_store.py

@@ -0,0 +1,40 @@
+from dataikuapi.dss.dataset import DSSDataset
+
+
+class DSSFeatureGroupListItem(object):
+    def __init__(self, client, project_key, name):
+        self.client = client
+        self.project_key = project_key
+        self.name = name
+
+    @property
+    def id(self):
+        return self.project_key + "." + self.name
+
+    def get_as_dataset(self):
+        """
+        Gets the feature group as a dataset
+
+        :return: a handle on the dataset
+        :rtype: :class:`dataikuapi.dss.dataset.DSSDataset`
+        """
+        return DSSDataset(self.client, self.project_key, self.name)
+
+
+class DSSFeatureStore(object):
+    def __init__(self, client):
+        """
+        A handle on the Feature Store.
+        Do not create this class directly, use :meth:`DSSClient.get_feature_store`
+        """
+        self.client = client
+
+    def list_feature_groups(self):
+        """
+        Get a list of feature groups on which the user has at least read permissions
+
+        :return: list of feature groups
+        :rtype: list of :class:`dataikuapi.feature_store.DSSFeatureGroupListItem`
+        """
+        items = self.client._perform_json("GET", "/feature-store/feature-groups")
+        return [DSSFeatureGroupListItem(self.client, item["projectKey"], item["name"]) for item in items]

dataikuapi/dssclient.py

@@ -4,6 +4,7 @@
 from requests import exceptions
 from requests.auth import HTTPBasicAuth
 
+from .dss.feature_store import DSSFeatureStore
 from .dss.notebook import DSSNotebook
 from .dss.future import DSSFuture
 from .dss.projectfolder import DSSProjectFolder
@@ -1118,6 +1119,19 @@ def get_object_discussions(self, project_key, object_type, object_id):
         """
         return DSSObjectDiscussions(self, project_key, object_type, object_id)
 
+    ########################################################
+    # Feature Store
+    ########################################################
+    def get_feature_store(self):
+        """
+        Get a handle to interact with the Feature Store.
+
+        :return: a handle on the feature store
+        :rtype: :class:`dataikuapi.feature_store.DSSFeatureStore`
+        """
+        return DSSFeatureStore(self)
+
+
 class TemporaryImportHandle(object):
     def __init__(self, client, import_id):
         self.client = client