Repository: alteryx/evalml (python)

Issue:
Harden ARIMA implementation
There are a number of small issues with our ARIMA implementation that cause and hide bugs. Namely:
- When calculating the forecast horizon, we try to calculate the distance forward but fall back to incorrect values because we end up checking `isinstance(np.int64, int)`, which returns False when it should be considered True.
- If X is None, we set `self.last_X_index` to be the entire y index, rather than the last value of y's index

These cause errors in their own right, and if they're fixed, this code block should raise a `ValueError: X array dims (n_rows) != n_periods`:
```
X = pd.DataFrame({"date": pd.date_range("1/1/21", periods=100)})
y = pd.Series([1, 4, 2] * 20 + [3, 1] * 20)

problem_configuration = {
    "max_delay": 52,
    "gap": 0,
    "forecast_horizon": 5,
    "time_index": "date",
}
X_train, X_holdout, y_train, y_holdout = split_data(
    X, y, problem_type="time series regression"
)
automl = AutoMLSearch(
    X_train,
    y_train,
    problem_type="time series regression",
    problem_configuration=problem_configuration,
)
automl.search()
arima = automl.get_pipeline(11)
arima.fit(X_train, y_train)
y_forecast = arima.get_forecast_predictions(X=X, y=y)
```

Pull request description:
Fix small ARIMA issues
Closes #4185 

Patch:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 13e41eab..6ad38c84 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -3,6 +3,7 @@ Release Notes
 **Future Releases**
     * Enhancements
     * Fixes
+        * Fixed small errors in ``ARIMARegressor`` implementation :pr:`4186`
     * Changes
     * Documentation Changes
     * Testing Changes
diff --git a/evalml/pipelines/components/estimators/regressors/arima_regressor.py b/evalml/pipelines/components/estimators/regressors/arima_regressor.py
index fb829f98..a037f211 100644
--- a/evalml/pipelines/components/estimators/regressors/arima_regressor.py
+++ b/evalml/pipelines/components/estimators/regressors/arima_regressor.py
@@ -3,6 +3,7 @@ from typing import Dict, Hashable, List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
+from pandas.api.types import is_integer_dtype
 from skopt.space import Integer
 
 from evalml.model_family import ModelFamily
@@ -154,24 +155,25 @@ class ARIMARegressor(Estimator):
         from sktime.forecasting.base import ForecastingHorizon
 
         # we can only calculate the difference if the indices are of the same type
-        if isinstance(X.index[0], type(self.last_X_index)):
-            units_diff = X.index[0] - self.last_X_index
-            if isinstance(X.index, pd.DatetimeIndex):
-                dates_diff = pd.date_range(
-                    start=self.last_X_index,
-                    end=X.index[0],
-                    freq=X.index.freq,
-                )
-                units_diff = len(dates_diff) - 1
-            fh_ = ForecastingHorizon(
-                [units_diff + i for i in range(len(X))],
-                is_relative=True,
-            )
-        else:
-            fh_ = ForecastingHorizon(
-                [i + 1 for i in range(len(X))],
-                is_relative=True,
+        units_diff = 1
+        if isinstance(X.index[0], type(self.last_X_index)) and isinstance(
+            X.index,
+            pd.DatetimeIndex,
+        ):
+            dates_diff = pd.date_range(
+                start=self.last_X_index,
+                end=X.index[0],
+                freq=X.index.freq,
             )
+            units_diff = len(dates_diff) - 1
+        elif is_integer_dtype(type(X.index[0])) and is_integer_dtype(
+            type(self.last_X_index),
+        ):
+            units_diff = X.index[0] - self.last_X_index
+        fh_ = ForecastingHorizon(
+            [units_diff + i for i in range(len(X))],
+            is_relative=True,
+        )
         return fh_
 
     def _get_sp(self, X: pd.DataFrame) -> int:
@@ -213,7 +215,7 @@ class ARIMARegressor(Estimator):
 
         sp = self._get_sp(X)
         self._component_obj.sp = sp
-        self.last_X_index = X.index[-1] if X is not None else y.index
+        self.last_X_index = X.index[-1] if X is not None else y.index[-1]
 
         X = self._remove_datetime(X, features=True)
 
@@ -279,7 +281,12 @@ class ARIMARegressor(Estimator):
                 # pmdarima requires the number of covariate rows to equal the length of the total number of periods (X.shape[0] == fh_[-1]) if covariates are used.
                 # We circument this by adding arbitrary rows to the start of X since sktime discards these values when predicting.
                 num_rows_diff = fh_[-1] - X.shape[0]
-                X_ = pd.concat([X.head(num_rows_diff), X], ignore_index=True)
+                filler = pd.DataFrame(
+                    columns=X.columns,
+                    index=range(num_rows_diff),
+                ).fillna(0)
+                X_ = pd.concat([filler, X], ignore_index=True)
+                X_.ww.init(schema=X.ww.schema)
             else:
                 X_ = X
             y_pred_intervals = self._component_obj.predict_interval(


Test patch:
diff --git a/evalml/tests/component_tests/test_arima_regressor.py b/evalml/tests/component_tests/test_arima_regressor.py
index 29913bab..6cf71b3d 100644
--- a/evalml/tests/component_tests/test_arima_regressor.py
+++ b/evalml/tests/component_tests/test_arima_regressor.py
@@ -119,6 +119,7 @@ def test_match_indices(ts_data):
         datetime_feature=False,
         no_features=False,
         test_features_index_dt=False,
+        match_indices=False,
     )
 
     assert not X_train.index.equals(y_train.index)
@@ -271,6 +272,7 @@ def test_fit_predict_sk_failure(
         datetime_feature,
         no_features,
         test_features_index_dt,
+        match_indices=False,
     )
 
     a_clf = AutoARIMA(maxiter=10)
@@ -475,6 +477,38 @@ def test_arima_regressor_can_forecast_arbitrary_dates(use_covariates, ts_data):
     )
 
 
+@pytest.mark.parametrize("use_covariates", [True, False])
+def test_arima_regressor_can_forecast_arbitrary_dates_past_holdout(
+    use_covariates,
+    ts_data,
+):
+    X, _, y = ts_data(
+        train_features_index_dt=False,
+        train_target_index_dt=False,
+    )
+
+    # Create a training and testing set that are not continuous
+    X_train, X_test, y_train, _ = split_data(
+        X,
+        y,
+        problem_type="time series regression",
+        test_size=0.1,
+        random_seed=0,
+    )
+    X_train, _, y_train, _ = split_data(
+        X_train,
+        y_train,
+        problem_type="time series regression",
+        test_size=0.4,
+        random_seed=0,
+    )
+
+    arima = ARIMARegressor(use_covariates=use_covariates)
+    arima.fit(X_train, y_train)
+
+    arima.predict(X_test)
+
+
 @pytest.mark.parametrize(
     "nullable_ltype",
     ["IntegerNullable", "AgeNullable"],
diff --git a/evalml/tests/conftest.py b/evalml/tests/conftest.py
index 0c667a57..d1714abf 100644
--- a/evalml/tests/conftest.py
+++ b/evalml/tests/conftest.py
@@ -317,11 +317,14 @@ def ts_data():
         test_features_index_dt=True,
         freq="D",
         problem_type="time series regression",
+        match_indices=True,
     ):
         X = pd.DataFrame(index=[i + 1 for i in range(50)])
         dates = pd.date_range("1/1/21", periods=50, freq=freq)
         feature = pd.Series([1, 4, 2] * 10 + [3, 1] * 10, index=X.index)
         y = pd.Series([1, 2, 3, 4, 5, 6, 5, 4, 3, 2] * 5)
+        if match_indices:
+            y.index = X.index
 
         X_train = X.iloc[:40]
         X_test = X.iloc[40:]