fix: prepare_data — close double-split and missing sample_weight_all on non-shuffle splits (#1554)

Author: immu4989

flaml/automl/task/generic_task.py

@@ -880,6 +880,7 @@ def prepare_data(
                         y_train_all = np.concatenate([y_train_all, y_train_all[:n][rare_index]])
                     count += rare_count
                 logger.info(f"class {label} augmented from {rare_count} to {count}")
+        state.sample_weight_all = sample_weight_full
         SHUFFLE_SPLIT_TYPES = ["uniform", "stratified"]
         if is_spark_dataframe:
             # no need to shuffle pyspark dataframe
@@ -952,7 +953,7 @@ def prepare_data(
                     ) = self._split_pyspark(state, X_train_all, y_train_all, split_ratio)
                 else:
                     X_train, X_val, y_train, y_val = self._split_pyspark(state, X_train_all, y_train_all, split_ratio)
-            if split_type == "group":
+            elif split_type == "group":
                 gss = GroupShuffleSplit(n_splits=1, test_size=split_ratio, random_state=RANDOM_SEED)
                 for train_idx, val_idx in gss.split(X_train_all, y_train_all, state.groups_all):
                     if data_is_df:

test/automl/test_split.py

@@ -52,6 +52,28 @@ def test_time():
     _test(split_type="time")
 
 
+def test_time_split_with_sample_weight():
+    """Regression for #887 and #1553: split_type='time' + sample_weight + holdout."""
+    from sklearn.datasets import make_classification
+
+    X, y = make_classification(n_samples=200, n_features=10, n_informative=5, random_state=42)
+    sample_weight = np.ones(len(y))
+    automl = AutoML()
+    automl.fit(
+        X_train=X,
+        y_train=y,
+        sample_weight=sample_weight,
+        split_type="time",
+        retrain_full=True,
+        eval_method="holdout",
+        task="classification",
+        time_budget=-1,
+        max_iter=2,
+        estimator_list=["lgbm"],
+    )
+    assert automl.model is not None
+
+
 def test_groups_for_classification_task():
     from sklearn.externals._arff import ArffException