Feature store public API

Author: lpenet

dataikuapi/dss/dataset.py

@@ -737,6 +737,26 @@ 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.
+        Changes of this property will be applied when calling :meth:`save` and require the "Manage Feature Store" permission.
+
+        :rtype: bool
+        """
+        return self.settings["featureGroup"]
+
+    def set_feature_group(self, status):
+        """
+        Indicates whether the Dataset is defined 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,14 @@
+class DSSFeatureStore(object):
+    def __init__(self, client):
+        """
+        A handle on the Feature Store. Do not call this methid directly, call """
+        self.client = client
+
+    def list_feature_groups(self):
+        """
+        Get a list of names of datasets defined as feature groups on the DSS instance.
+
+        :return: list of names of datasets defined as feature groups
+        :rtype: list of str
+        """
+        return self.client._perform_json("GET", "/feature-store/feature-groups/list")

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,14 @@ 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):
+        return DSSFeatureStore(self)
+
+
+
 class TemporaryImportHandle(object):
     def __init__(self, client, import_id):
         self.client = client