PySpark XGBoost integration

python-package/xgboost/spark/__init__.py

@@ -0,0 +1,16 @@
+"""XGBoost: eXtreme Gradient Boosting library.
+
+Contributors: https://github.com/dmlc/xgboost/blob/master/CONTRIBUTORS.md
+"""
+
+try:
+    import pyspark
+except ImportError:
+    raise RuntimeError("xgboost spark python API requires pyspark package installed.")
+
+from .estimator import (XgboostClassifier, XgboostClassifierModel,
+                        XgboostRegressor, XgboostRegressorModel)
+
+__all__ = ['XgboostClassifier', 'XgboostClassifierModel',
+           'XgboostRegressor', 'XgboostRegressorModel']
+

python-package/xgboost/spark/data.py

@@ -0,0 +1,262 @@
+import os
+from typing import Iterator
+import numpy as np
+import pandas as pd
+from scipy.sparse import csr_matrix
+from xgboost import DMatrix
+
+
+# Since sklearn's SVM converter doesn't address weights, this one does address weights:
+def _dump_libsvm(features, labels, weights=None, external_storage_precision=5):
+    esp = external_storage_precision
+    lines = []
+
+    def gen_label_str(row_idx):
+        if weights is not None:
+            return "{label:.{esp}g}:{weight:.{esp}g}".format(
+                label=labels[row_idx], esp=esp, weight=weights[row_idx])
+        else:
+            return "{label:.{esp}g}".format(label=labels[row_idx], esp=esp)
+
+    def gen_feature_value_str(feature_idx, feature_val):
+        return "{idx:.{esp}g}:{value:.{esp}g}".format(
+            idx=feature_idx, esp=esp, value=feature_val
+        )
+
+    is_csr_matrix = isinstance(features, csr_matrix)
+
+    for i in range(len(labels)):
+        current = [gen_label_str(i)]
+        if is_csr_matrix:
+            idx_start = features.indptr[i]
+            idx_end = features.indptr[i + 1]
+            for idx in range(idx_start, idx_end):
+                j = features.indices[idx]
+                val = features.data[idx]
+                current.append(gen_feature_value_str(j, val))
+        else:
+            for j, val in enumerate(features[i]):
+                current.append(gen_feature_value_str(j, val))
+        lines.append(" ".join(current) + "\n")
+    return lines
+
+
+# This is the updated version that handles weights
+def _stream_train_val_data(features, labels, weights, main_file,
+                           external_storage_precision):
+    lines = _dump_libsvm(features, labels, weights, external_storage_precision)
+    main_file.writelines(lines)
+
+
+def _stream_data_into_libsvm_file(data_iterator, has_weight,
+                                  has_validation, file_prefix,
+                                  external_storage_precision):
+    # getting the file names for storage
+    train_file_name = file_prefix + "/data.txt.train"
+    train_file = open(train_file_name, "w")
+    if has_validation:
+        validation_file_name = file_prefix + "/data.txt.val"
+        validation_file = open(validation_file_name, "w")
+
+    train_val_data = _process_data_iter(data_iterator,
+                                        train=True,
+                                        has_weight=has_weight,
+                                        has_validation=has_validation)
+    if has_validation:
+        train_X, train_y, train_w, _, val_X, val_y, val_w, _ = train_val_data
+        _stream_train_val_data(train_X, train_y, train_w, train_file,
+                               external_storage_precision)
+        _stream_train_val_data(val_X, val_y, val_w, validation_file,
+                               external_storage_precision)
+    else:
+        train_X, train_y, train_w, _ = train_val_data
+        _stream_train_val_data(train_X, train_y, train_w, train_file,
+                               external_storage_precision)
+
+    if has_validation:
+        train_file.close()
+        validation_file.close()
+        return train_file_name, validation_file_name
+    else:
+        train_file.close()
+        return train_file_name
+
+
+def _create_dmatrix_from_file(file_name, cache_name):
+    if os.path.exists(cache_name):
+        os.remove(cache_name)
+    if os.path.exists(cache_name + ".row.page"):
+        os.remove(cache_name + ".row.page")
+    if os.path.exists(cache_name + ".sorted.col.page"):
+        os.remove(cache_name + ".sorted.col.page")
+    return DMatrix(file_name + "#" + cache_name)
+
+
+def prepare_train_val_data(data_iterator,
+                           has_weight,
+                           has_validation,
+                           has_fit_base_margin=False):
+    def gen_data_pdf():
+        for pdf in data_iterator:
+            yield pdf
+
+    return _process_data_iter(gen_data_pdf(),
+                              train=True,
+                              has_weight=has_weight,
+                              has_validation=has_validation,
+                              has_fit_base_margin=has_fit_base_margin,
+                              has_predict_base_margin=False)
+
+
+def prepare_predict_data(data_iterator, has_predict_base_margin):
+    return _process_data_iter(data_iterator,
+                              train=False,
+                              has_weight=False,
+                              has_validation=False,
+                              has_fit_base_margin=False,
+                              has_predict_base_margin=has_predict_base_margin)
+
+
+def _check_feature_dims(num_dims, expected_dims):
+    """
+    Check all feature vectors has the same dimension
+    """
+    if expected_dims is None:
+        return num_dims
+    if num_dims != expected_dims:
+        raise ValueError("Rows contain different feature dimensions: "
+                         "Expecting {}, got {}.".format(
+                             expected_dims, num_dims))
+    return expected_dims
+
+
+def _row_tuple_list_to_feature_matrix_y_w(data_iterator, train, has_weight,
+                                          has_fit_base_margin,
+                                          has_predict_base_margin,
+                                          has_validation: bool = False):
+    """
+    Construct a feature matrix in ndarray format, label array y and weight array w
+    from the row_tuple_list.
+    If train == False, y and w will be None.
+    If has_weight == False, w will be None.
+    If has_base_margin == False, b_m will be None.
+    Note: the row_tuple_list will be cleared during
+    executing for reducing peak memory consumption
+    """
+    expected_feature_dims = None
+    label_list, weight_list, base_margin_list = [], [], []
+    label_val_list, weight_val_list, base_margin_val_list = [], [], []
+    values_list, values_val_list = [], []
+
+    # Process rows
+    for pdf in data_iterator:
+        if type(pdf) == tuple:
+            pdf = pd.concat(list(pdf), axis=1, names=["values", "baseMargin"])
+
+        if len(pdf) == 0:
+            continue
+        if train and has_validation:
+            pdf_val = pdf.loc[pdf["validationIndicator"], :]
+            pdf = pdf.loc[~pdf["validationIndicator"], :]
+
+        num_feature_dims = len(pdf["values"].values[0])
+
+        expected_feature_dims = _check_feature_dims(num_feature_dims,
+                                                    expected_feature_dims)
+
+        values_list.append(pdf["values"].to_list())
+        if train:
+            label_list.append(pdf["label"].to_list())
+        if has_weight:
+            weight_list.append(pdf["weight"].to_list())
+        if has_fit_base_margin or has_predict_base_margin:
+            base_margin_list.append(pdf.iloc[:, -1].to_list())
+        if has_validation:
+            values_val_list.append(pdf_val["values"].to_list())
+            if train:
+                label_val_list.append(pdf_val["label"].to_list())
+            if has_weight:
+                weight_val_list.append(pdf_val["weight"].to_list())
+            if has_fit_base_margin or has_predict_base_margin:
+                base_margin_val_list.append(pdf_val.iloc[:, -1].to_list())
+
+    # Construct feature_matrix
+    if expected_feature_dims is None:
+        return [], [], [], []
+
+    # Construct feature_matrix, y and w
+    feature_matrix = np.concatenate(values_list)
+    y = np.concatenate(label_list) if train else None
+    w = np.concatenate(weight_list) if has_weight else None
+    b_m = np.concatenate(base_margin_list) if (
+            has_fit_base_margin or has_predict_base_margin) else None
+    if has_validation:
+        feature_matrix_val = np.concatenate(values_val_list)
+        y_val = np.concatenate(label_val_list) if train else None
+        w_val = np.concatenate(weight_val_list) if has_weight else None
+        b_m_val = np.concatenate(base_margin_val_list) if (
+                has_fit_base_margin or has_predict_base_margin) else None
+        return feature_matrix, y, w, b_m, feature_matrix_val, y_val, w_val, b_m_val
+    return feature_matrix, y, w, b_m
+
+
+def _process_data_iter(data_iterator: Iterator[pd.DataFrame],
+                       train: bool,
+                       has_weight: bool,
+                       has_validation: bool,
+                       has_fit_base_margin: bool = False,
+                       has_predict_base_margin: bool = False):
+    """
+    If input is for train and has_validation=True, it will split the train data into train dataset
+    and validation dataset, and return (train_X, train_y, train_w, train_b_m <-
+    train base margin, val_X, val_y, val_w, val_b_m <- validation base margin)
+    otherwise return (X, y, w, b_m <- base margin)
+    """
+    if train and has_validation:
+        train_X, train_y, train_w, train_b_m, val_X, val_y, val_w, val_b_m = \
+            _row_tuple_list_to_feature_matrix_y_w(
+                data_iterator, train, has_weight, has_fit_base_margin,
+                has_predict_base_margin, has_validation)
+        return train_X, train_y, train_w, train_b_m, val_X, val_y, val_w, val_b_m
+    else:
+        return _row_tuple_list_to_feature_matrix_y_w(data_iterator, train, has_weight,
+                                                     has_fit_base_margin, has_predict_base_margin,
+                                                     has_validation)
+
+
+def convert_partition_data_to_dmatrix(partition_data_iter,
+                                      has_weight,
+                                      has_validation,
+                                      use_external_storage=False,
+                                      file_prefix=None,
+                                      external_storage_precision=5):
+    # if we are using external storage, we use a different approach for making the dmatrix
+    if use_external_storage:
+        if has_validation:
+            train_file, validation_file = _stream_data_into_libsvm_file(
+                partition_data_iter, has_weight,
+                has_validation, file_prefix, external_storage_precision)
+            training_dmatrix = _create_dmatrix_from_file(
+                train_file, "{}/train.cache".format(file_prefix))
+            val_dmatrix = _create_dmatrix_from_file(
+                validation_file, "{}/val.cache".format(file_prefix))
+            return training_dmatrix, val_dmatrix
+        else:
+            train_file = _stream_data_into_libsvm_file(
+                partition_data_iter, has_weight,
+                has_validation, file_prefix, external_storage_precision)
+            training_dmatrix = _create_dmatrix_from_file(
+                train_file, "{}/train.cache".format(file_prefix))
+            return training_dmatrix
+
+    # if we are not using external storage, we use the standard method of parsing data.
+    train_val_data = prepare_train_val_data(partition_data_iter, has_weight, has_validation)
+    if has_validation:
+        train_X, train_y, train_w, _, val_X, val_y, val_w, _ = train_val_data
+        training_dmatrix = DMatrix(data=train_X, label=train_y, weight=train_w)
+        val_dmatrix = DMatrix(data=val_X, label=val_y, weight=val_w)
+        return training_dmatrix, val_dmatrix
+    else:
+        train_X, train_y, train_w, _ = train_val_data
+        training_dmatrix = DMatrix(data=train_X, label=train_y, weight=train_w)
+        return training_dmatrix