Speed up PandasDataset further

src/gluonts/dataset/pandas.py

@@ -15,7 +15,7 @@
 from dataclasses import dataclass, field
 from typing import (
     Any,
-    cast,
+    Collection,
     Dict,
     Iterator,
     Iterable,
@@ -26,10 +26,11 @@
 )
 
 import pandas as pd
+import numpy as np
 from pandas.core.indexes.datetimelike import DatetimeIndexOpsMixin
 from toolz import first
 
-from gluonts.dataset.common import DataEntry, ProcessDataEntry
+from gluonts.dataset.common import DataEntry, _as_period
 from gluonts.dataset.field_names import FieldName
 from gluonts.itertools import Map, SizedIterable
 
@@ -109,7 +110,6 @@ class PandasDataset:
     def __post_init__(self) -> None:
         if isinstance(self.target, list) and len(self.target) == 1:
             self.target = self.target[0]
-        self.one_dim_target = not isinstance(self.target, list)
 
         if isinstance(self.dataframes, dict):
             self._pairs = self.dataframes.items()
@@ -127,10 +127,6 @@ def __post_init__(self) -> None:
             ), "You need to provide `freq` along with `timestamp`"
             self.freq = infer_freq(first(self._pairs)[1].index)
 
-        self.process = ProcessDataEntry(
-            cast(str, self.freq), one_dim_target=self.one_dim_target
-        )
-
         self._data_entries = Map(self._pair_to_dataentry, self._pairs)
 
     def _pair_to_dataentry(
@@ -170,7 +166,8 @@ def _pair_to_dataentry(
         if item_id is not None:
             data_entry["item_id"] = item_id
 
-        return self.process(data_entry)
+        data_entry["start"] = _as_period(data_entry["start"], self.freq)
+        return data_entry
 
     def __iter__(self) -> Iterator[DataEntry]:
         for entry in self._data_entries:
@@ -231,6 +228,20 @@ def infer_freq(index: pd.Index):
     return pd.infer_freq(index)
 
 
+def extract_dynamic_array(
+    df: pd.DataFrame, col_names: Union[str, Collection[str]]
+) -> np.ndarray:
+    return df[col_names].values.transpose()
+
+
+def extract_static_array(
+    df: pd.DataFrame, col_names: Union[str, Collection[str]]
+) -> np.ndarray:
+    if isinstance(col_names, str):
+        return df[col_names].values[:1]
+    return df[col_names].values[0, :]
+
+
 def as_dataentry(
     data: pd.DataFrame,
     target: Union[str, List[str]],
@@ -277,21 +288,29 @@ def as_dataentry(
     start = data.loc[:, timestamp].iloc[0] if timestamp else data.index[0]
     dataentry = {FieldName.START: start}
 
-    def set_field(fieldname, col_names, f=lambda x: x):
-        if len(col_names) > 0:
-            dataentry[fieldname] = [
-                f(data.loc[:, n].to_list()) for n in col_names
-            ]
-
-    if isinstance(target, str):
-        dataentry[FieldName.TARGET] = data.loc[:, target].to_list()
-    else:
-        set_field(FieldName.TARGET, target)
-    set_field(FieldName.FEAT_DYNAMIC_REAL, feat_dynamic_real)
-    set_field(FieldName.FEAT_DYNAMIC_CAT, feat_dynamic_cat)
-    set_field(FieldName.FEAT_STATIC_REAL, feat_static_real, lambda x: x[0])
-    set_field(FieldName.FEAT_STATIC_CAT, feat_static_cat, lambda x: x[0])
-    set_field(FieldName.PAST_FEAT_DYNAMIC_REAL, past_feat_dynamic_real)
+    dataentry[FieldName.TARGET] = extract_dynamic_array(data, target)
+
+    if len(feat_dynamic_real) > 0:
+        dataentry[FieldName.FEAT_DYNAMIC_REAL] = extract_dynamic_array(
+            data, feat_dynamic_real
+        )
+    if len(feat_dynamic_cat) > 0:
+        dataentry[FieldName.FEAT_DYNAMIC_CAT] = extract_dynamic_array(
+            data, feat_dynamic_cat
+        )
+    if len(past_feat_dynamic_real) > 0:
+        dataentry[FieldName.PAST_FEAT_DYNAMIC_REAL] = extract_dynamic_array(
+            data, past_feat_dynamic_real
+        )
+    if len(feat_static_real) > 0:
+        dataentry[FieldName.FEAT_STATIC_REAL] = extract_static_array(
+            data, feat_static_real
+        )
+    if len(feat_static_cat) > 0:
+        dataentry[FieldName.FEAT_STATIC_CAT] = extract_static_array(
+            data, feat_static_cat
+        )
+
     return dataentry