Repository: mindsdb/lightwood

Issue:
Error if train TS predictor on data with single value in 'group_by'
I tryed train TS predictor on data:
```
df = pd.DataFrame({
            'group': ["A" for _ in range(100, 210)],
            'order': [x for x in range(100, 210)],
            'x1': [x for x in range(100, 210)],
            'x2': [x * 2 for x in range(100, 210)],
            'y': [x * 3 for x in range(100, 210)]
        })
```
I set `group_by` to `group`. I got error:
```
Traceback (most recent call last):
  File "/home/maxs/dev/mdb/venv38/sources/mindsdb/mindsdb/interfaces/model/learn_process.py", line 100, in run_learn
    run_fit(predictor_id, df)
  File "/home/maxs/dev/mdb/venv38/sources/mindsdb/mindsdb/utilities/functions.py", line 56, in wrapper
    return func(*args, **kwargs)
  File "/home/maxs/dev/mdb/venv38/sources/mindsdb/mindsdb/interfaces/model/learn_process.py", line 87, in run_fit
    raise e
  File "/home/maxs/dev/mdb/venv38/sources/mindsdb/mindsdb/interfaces/model/learn_process.py", line 67, in run_fit
    predictor.learn(df)
  File "/tmp/5ebf11ab5e9ba76011bd4d012f7779ef3ad842556160a13916371366856095033.py", line 397, in learn
    data = self.preprocess(data)
  File "/tmp/5ebf11ab5e9ba76011bd4d012f7779ef3ad842556160a13916371366856095033.py", line 137, in preprocess
    data = transform_timeseries(
  File "/home/maxs/dev/mdb/venv38/sources/lightwood/lightwood/data/timeseries_transform.py", line 105, in transform_timeseries
    for _, df in original_df.groupby(gb_arr):
  File "/home/maxs/dev/mdb/venv38/lib/python3.8/site-packages/pandas/core/frame.py", line 7631, in groupby
    return DataFrameGroupBy(
  File "/home/maxs/dev/mdb/venv38/lib/python3.8/site-packages/pandas/core/groupby/groupby.py", line 889, in __init__
    grouper, exclusions, obj = get_grouper(
  File "/home/maxs/dev/mdb/venv38/lib/python3.8/site-packages/pandas/core/groupby/grouper.py", line 862, in get_grouper
    raise KeyError(gpr)
KeyError: 'group'
```
As i see in debugger, `original_df` in `.../timeseries_transform.py", line 105` does not have 'group' column. I guess it was deleted because it have only single value. 
If it so, then i think we should not remove such column if user explicitly use it in 'group_by'


Discussion:

Resolving patch:
diff --git a/lightwood/api/json_ai.py b/lightwood/api/json_ai.py
--- a/lightwood/api/json_ai.py
+++ b/lightwood/api/json_ai.py
@@ -185,15 +185,18 @@ def generate_json_ai(
     exec(IMPORT_EXTERNAL_DIRS, globals())
     target = problem_definition.target
     input_cols = []
+    tss = problem_definition.timeseries_settings
+
     for col_name, col_dtype in type_information.dtypes.items():
         if (
-            col_name not in type_information.identifiers
-            and col_dtype not in (dtype.invalid, dtype.empty)
-            and col_name != target
+                (col_name not in type_information.identifiers
+                 and col_dtype not in (dtype.invalid, dtype.empty)
+                 and col_name != target)
+                or
+                (tss.group_by is not None and col_name in tss.group_by)
         ):
             input_cols.append(col_name)
 
-    tss = problem_definition.timeseries_settings
     is_target_predicting_encoder = False
     is_ts = problem_definition.timeseries_settings.is_timeseries
     # Single text column classification
diff --git a/lightwood/data/cleaner.py b/lightwood/data/cleaner.py
--- a/lightwood/data/cleaner.py
+++ b/lightwood/data/cleaner.py
@@ -280,8 +280,12 @@ def _remove_columns(data: pd.DataFrame, identifiers: Dict[str, object], target:
     data = deepcopy(data)
     to_drop = [*[x for x in identifiers.keys() if x != target],
                *[x for x in data.columns if x in dtype_dict and dtype_dict[x] == dtype.invalid]]
+
     exceptions = ["__mdb_make_predictions"]
-    to_drop = [x for x in to_drop if x in data.columns]
+    if timeseries_settings.group_by is not None:
+        exceptions += timeseries_settings.group_by
+
+    to_drop = [x for x in to_drop if x in data.columns and x not in exceptions]
     data = data.drop(columns=to_drop)
 
     if mode == "train":
@@ -298,7 +302,6 @@ def _remove_columns(data: pd.DataFrame, identifiers: Dict[str, object], target:
     for name in list(data.columns):
         if name not in dtype_dict and name not in exceptions:
             data = data.drop(columns=[name])
-
     return data
 
 
diff --git a/lightwood/mixer/sktime.py b/lightwood/mixer/sktime.py
--- a/lightwood/mixer/sktime.py
+++ b/lightwood/mixer/sktime.py
@@ -117,8 +117,8 @@ def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],
                 series = series.reset_index(drop=True)
 
                 for idx, _ in enumerate(series.iteritems()):
-                    ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(
-                        np.arange(idx, idx + self.n_ts_predictions)).tolist()
+                    preds = forecaster.predict(np.arange(idx, idx + self.n_ts_predictions)).squeeze()
+                    ydf['prediction'].iloc[series_idxs[idx]] = preds.tolist()
 
             if self.grouped_by == ['__default']:
                 break


Test patch: