Add feature weights support (#265)

* added support for feature_weights parameter and associated tests

* added tests for the feature_weights parameter

* ran code formatter

* revert version

* address lint_test failures

* addressed further test_lint failures

* addressing pytest failure reduce resource request

* addressing pytest failure increase boosting rounds for convergence

* addressing pytest failure increase boosting rounds for convergence

Author: mozjay0619

.gitignore

@@ -90,4 +90,7 @@ project-id
 # Downloaded test data
 *.csv
 *.csv.gz
-*.parquet
+*.parquet
+
+# Byte-compiled files
+__pycache__/

xgboost_ray/main.py

@@ -328,6 +328,8 @@ def _get_dmatrix(data: RayDMatrix, param: Dict) -> xgb.DMatrix:
                 param["label"] = [param["label"]]
             if not isinstance(param["weight"], list):
                 param["weight"] = [param["weight"]]
+            if not isinstance(param["feature_weights"], list):
+                param["feature_weights"] = [param["feature_weights"]]
             if not isinstance(param["qid"], list):
                 param["qid"] = [param["qid"]]
             if not isinstance(param["data"], list):
@@ -354,6 +356,7 @@ def _get_dmatrix(data: RayDMatrix, param: Dict) -> xgb.DMatrix:
                 "data": concat_dataframes(param["data"]),
                 "label": concat_dataframes(param["label"]),
                 "weight": concat_dataframes(param["weight"]),
+                "feature_weights": concat_dataframes(param["feature_weights"]),
                 "qid": concat_dataframes(param["qid"]),
                 "base_margin": concat_dataframes(param["base_margin"]),
                 "label_lower_bound": concat_dataframes(
@@ -365,6 +368,7 @@ def _get_dmatrix(data: RayDMatrix, param: Dict) -> xgb.DMatrix:
 
         ll = param.pop("label_lower_bound", None)
         lu = param.pop("label_upper_bound", None)
+        fw = param.pop("feature_weights", None)
 
         if LEGACY_MATRIX:
             param.pop("base_margin", None)
@@ -378,7 +382,8 @@ def _get_dmatrix(data: RayDMatrix, param: Dict) -> xgb.DMatrix:
         matrix = xgb.DMatrix(**param)
 
         if not LEGACY_MATRIX:
-            matrix.set_info(label_lower_bound=ll, label_upper_bound=lu)
+            matrix.set_info(
+                label_lower_bound=ll, label_upper_bound=lu, feature_weights=fw)
 
     data.update_matrix_properties(matrix)
     return matrix

xgboost_ray/matrix.py

@@ -102,6 +102,7 @@ def __init__(
             label: List[Optional[Data]],
             missing: Optional[float],
             weight: List[Optional[Data]],
+            feature_weights: List[Optional[Data]],
             qid: List[Optional[Data]],
             base_margin: List[Optional[Data]],
             label_lower_bound: List[Optional[Data]],
@@ -118,6 +119,7 @@ def __init__(
         self._label = label
         self._missing = missing
         self._weight = weight
+        self._feature_weights = feature_weights
         self._qid = qid
         self._base_margin = base_margin
         self._label_lower_bound = label_lower_bound
@@ -151,6 +153,7 @@ def next(self, input_data: Callable):
             data=self._prop(self._data),
             label=self._prop(self._label),
             weight=self._prop(self._weight),
+            feature_weights=self._prop(self._feature_weights),
             qid=self._prop(self._qid),
             group=None,
             label_lower_bound=self._prop(self._label_lower_bound),
@@ -168,6 +171,7 @@ def __init__(self,
                  label: Optional[Data] = None,
                  missing: Optional[float] = None,
                  weight: Optional[Data] = None,
+                 feature_weights: Optional[Data] = None,
                  base_margin: Optional[Data] = None,
                  label_lower_bound: Optional[Data] = None,
                  label_upper_bound: Optional[Data] = None,
@@ -182,6 +186,7 @@ def __init__(self,
         self.label = label
         self.missing = missing
         self.weight = weight
+        self.feature_weights = feature_weights
         self.base_margin = base_margin
         self.label_lower_bound = label_lower_bound
         self.label_upper_bound = label_upper_bound
@@ -248,8 +253,8 @@ def _split_dataframe(
         """
         Split dataframe into
 
-        `features`, `labels`, `weight`, `base_margin`, `label_lower_bound`,
-        `label_upper_bound`
+        `features`, `labels`, `weight`, `feature_weights`, `base_margin`,
+        `label_lower_bound`, `label_upper_bound`
 
         """
         # sort dataframe by qid if exists (required by DMatrix)
@@ -268,6 +273,11 @@ def _split_dataframe(
         if exclude:
             exclude_cols.add(exclude)
 
+        feature_weights, exclude = data_source.get_column(
+            local_data, self.feature_weights)
+        if exclude:
+            exclude_cols.add(exclude)
+
         qid, exclude = data_source.get_column(local_data, self.qid)
         if exclude:
             exclude_cols.add(exclude)
@@ -291,8 +301,8 @@ def _split_dataframe(
         if exclude_cols:
             x = x[[col for col in x.columns if col not in exclude_cols]]
 
-        return x, label, weight, base_margin, label_lower_bound, \
-            label_upper_bound, qid
+        return x, label, weight, feature_weights, base_margin, \
+            label_lower_bound, label_upper_bound, qid
 
     def load_data(self,
                   num_actors: int,
@@ -380,7 +390,7 @@ def load_data(self,
         # yet. Instead, we'll be selecting the rows below.
         local_df = data_source.load_data(
             self.data, ignore=self.ignore, indices=None, **self.kwargs)
-        x, y, w, b, ll, lu, qid = self._split_dataframe(
+        x, y, w, fw, b, ll, lu, qid = self._split_dataframe(
             local_df, data_source=data_source)
 
         if isinstance(x, list):
@@ -396,6 +406,7 @@ def load_data(self,
                 "data": ray.put(x.iloc[indices]),
                 "label": ray.put(y.iloc[indices] if y is not None else None),
                 "weight": ray.put(w.iloc[indices] if w is not None else None),
+                "feature_weights": ray.put(fw),
                 "base_margin": ray.put(b.iloc[indices]
                                        if b is not None else None),
                 "label_lower_bound": ray.put(ll.iloc[indices]
@@ -545,7 +556,7 @@ def load_data(self,
                 indices=rank_shards,
                 ignore=self.ignore,
                 **self.kwargs)
-            x, y, w, b, ll, lu, qid = self._split_dataframe(
+            x, y, w, fw, b, ll, lu, qid = self._split_dataframe(
                 local_df, data_source=data_source)
 
             if isinstance(x, list):
@@ -557,16 +568,16 @@ def load_data(self,
             indices = _get_sharding_indices(sharding, rank, num_actors, n)
 
             if not indices:
-                x, y, w, b, ll, lu, qid = (None, None, None, None, None, None,
-                                           None)
+                x, y, w, fw, b, ll, lu, qid = (None, None, None, None, None,
+                                               None, None, None)
                 n = 0
             else:
                 local_df = data_source.load_data(
                     self.data,
                     ignore=self.ignore,
                     indices=indices,
                     **self.kwargs)
-                x, y, w, b, ll, lu, qid = self._split_dataframe(
+                x, y, w, fw, b, ll, lu, qid = self._split_dataframe(
                     local_df, data_source=data_source)
 
                 if isinstance(x, list):
@@ -579,6 +590,7 @@ def load_data(self,
                 "data": ray.put(x),
                 "label": ray.put(y),
                 "weight": ray.put(w),
+                "feature_weights": ray.put(fw),
                 "base_margin": ray.put(b),
                 "label_lower_bound": ray.put(ll),
                 "label_upper_bound": ray.put(lu),
@@ -684,6 +696,7 @@ def __init__(self,
                  data: Data,
                  label: Optional[Data] = None,
                  weight: Optional[Data] = None,
+                 feature_weights: Optional[Data] = None,
                  base_margin: Optional[Data] = None,
                  missing: Optional[float] = None,
                  label_lower_bound: Optional[Data] = None,
@@ -739,6 +752,7 @@ def __init__(self,
                 label=label,
                 missing=missing,
                 weight=weight,
+                feature_weights=feature_weights,
                 base_margin=base_margin,
                 label_lower_bound=label_lower_bound,
                 label_upper_bound=label_upper_bound,
@@ -755,6 +769,7 @@ def __init__(self,
                 label=label,
                 missing=missing,
                 weight=weight,
+                feature_weights=feature_weights,
                 base_margin=base_margin,
                 label_lower_bound=label_lower_bound,
                 label_upper_bound=label_upper_bound,

xgboost_ray/tests/test_end_to_end.py

@@ -375,6 +375,47 @@ def testRanking(self):
         auc_rec = evals_result["train"]["aucpr"]
         self.assertTrue((p <= q for p, q in zip(auc_rec, auc_rec[1:])))
 
+    @unittest.skipIf(xgb.__version__ < "1.3.0",
+                     f"not supported in xgb version {xgb.__version__}")
+    def testFeatureWeightsParam(self):
+        """Test the feature_weights parameter for xgb version >= 1.3.0.
+        Adapted from the official demo codes:
+        http://xgboost.readthedocs.io/en/stable/python/examples/
+        feature_weights.html"""
+
+        rng = np.random.RandomState(1994)
+
+        kRows = 1000
+        kCols = 10
+
+        X = rng.randn(kRows, kCols)
+        y = rng.randn(kRows)
+        fw = np.ones(shape=(kCols, ))
+        for i in range(kCols):
+            fw[i] *= float(i)
+        train_set = RayDMatrix(X, y, feature_weights=fw)
+
+        evals_result = {}
+        bst = train(
+            {
+                "objective": "reg:squarederror",
+                "eval_metric": ["rmse", "error"],
+                "colsample_bynode": 0.1,
+            },
+            train_set,
+            num_boost_round=250,
+            evals_result=evals_result,
+            evals=[(train_set, "train")],
+            verbose_eval=False,
+            ray_params=RayParams(
+                num_actors=2,  # Number of remote actors
+                cpus_per_actor=1))
+
+        feature_map = bst.get_fscore()
+        # feature zero has 0 weight
+        self.assertTrue(feature_map.get("f0", None) is None)
+        self.assertTrue(max(feature_map.values()) == feature_map.get("f9"))
+
 
 if __name__ == "__main__":
     import pytest

xgboost_ray/tests/test_matrix.py

@@ -368,6 +368,15 @@ def testLegacyParams(self):
             label_lower_bound=label_lower_bound,
             label_upper_bound=label_upper_bound)
 
+    @unittest.skipIf(xgb.__version__ < "1.3.0",
+                     f"not supported in xgb version {xgb.__version__}")
+    def testFeatureWeightsParam(self):
+        """Test the feature_weights parameter for xgb version >= 1.3.0"""
+        in_x = self.x
+        in_y = self.y
+        feature_weights = np.arange(len(in_y))
+        self._testMatrixCreation(in_x, in_y, feature_weights=feature_weights)
+
     @unittest.skipIf("qid" not in inspect.signature(xgb.DMatrix).parameters,
                      f"not supported in xgb version {xgb.__version__}")
     def testQidSortedBehaviorXGBoost(self):