add filter parameter for explicit splitting

Author: FChataigner

dataikuapi/dss/ml.py

@@ -4,7 +4,7 @@
 import json
 import time
 from .metrics import ComputedMetrics
-from .utils import DSSDatasetSelectionBuilder
+from .utils import DSSDatasetSelectionBuilder, DSSFilterBuilder
 
 class PredictionSplitParamsHandler(object):
     """Object to modify the train/test splitting params."""
@@ -57,14 +57,16 @@ def set_split_kfold(self, n_folds = 5, selection = None, dataset_name=None):
         if dataset_name is not None:
             sp["ssdDatasetSmartName"] = dataset_name
 
-    def set_split_explicit(self, train_selection, test_selection, dataset_name=None, test_dataset_name=None):
+    def set_split_explicit(self, train_selection, test_selection, dataset_name=None, test_dataset_name=None, train_filter=None, test_filter=None):
         """
         Sets the train/test split to explicit extract of one or two dataset
 
         :param object train_selection: A :class:`DSSDatasetSelectionBuilder` to build the settings of the extract of the train dataset. May be None (won't be changed)
         :param object test_selection: A :class:`DSSDatasetSelectionBuilder` to build the settings of the extract of the test dataset. May be None (won't be changed)
         :param str dataset_name: Name of dataset to use for the extracts. If None, the main dataset used to create the ML Task will be used.
         :param str test_dataset_name: Name of a second dataset to use for the test data extract. If None, both extracts are done from dataset_name
+        :param object train_filter: A :class:`DSSFilterBuilder` to build the settings of the filter of the train dataset. May be None (won't be changed)
+        :param object test_filter: A :class:`DSSFilterBuilder` to build the settings of the filter of the test dataset. May be None (won't be changed)
         """
         sp = self.mltask_settings["splitParams"]
         if dataset_name is None:
@@ -94,6 +96,17 @@ def set_split_explicit(self, train_selection, test_selection, dataset_name=None,
             else:
                 test_split["selection"] = test_selection
 
+        if train_filter is not None:
+            if isinstance(train_filter, DSSFilterBuilder):
+                train_split["filter"] = train_filter.build()
+            else:
+                train_split["filter"] = train_filter
+        if test_filter is not None:
+            if isinstance(test_filter, DSSFilterBuilder):
+                test_split["filter"] = test_filter.build()
+            else:
+                test_split["filter"] = test_filter
+
 
 class DSSMLTaskSettings(object):
     """

dataikuapi/dss/utils.py

@@ -42,3 +42,25 @@ def with_selected_partitions(self, ids):
         return self
 
 
+class DSSFilterBuilder(object):
+    """
+    Builder for a "filter". In DSS, a filter is used to define a subset of rows for processing.
+    """
+    def __init__(self):
+        self.filter = {"enabled":False, "distinct":False, "expression":None, "uiData":{"mode":"CUSTOM"}}
+
+    def build(self):
+        """Returns the built filter dict"""
+        return self.filter
+
+    def with_distinct(self):
+        """Sets the filter to deduplicate"""
+        self.filter["distinct"] = True
+        return self
+
+    def with_formula(self, expression):
+        """Sets the filter to deduplicate"""
+        self.filter["enabled"] = True
+        self.filter["expression"] = expression
+        self.filter["uiData"]["mode"] = "CUSTOM"
+        return self