Repository: blue-yonder/tsfresh

Issue:
High memory requirements during feature extraction
1. Your operating system
Linux

2. The version of *tsfresh* that you are using
0.16.0

3. The data on which the problem occurred (please do not upload 1000s of time series but try to boil the problem down to a small group or even a singular one)
```
<class 'pandas.core.frame.DataFrame'>
Int64Index: 6024604 entries, 0 to 1299438
Columns: 77 entries, index to machine_id
dtypes: float32(75), int64(1), object(1)
memory usage: 1.8+ GB
```

4. A minimal code snippet which reproduces the problem/bug
See below

5. Any reported errors or traceback
OOM

I was experimenting with tsfresh on the above mentioned dataset using a scikit learn pipeline:
`RelevantFeatureAugmenter(column_id='machine_id', column_sort='index', default_fc_parameters = MinimalFCParameters(), n_jobs=4, chunksize=10)
`
The code was executed from a Jupyter notebook which has 20GB of RAM available. The kernel died from out of memory errors, no matter which parameters I chose for `chunksize` or `n_jobs`. This seemed a bit strange cosidering the medium sized input and the tasks I was imagining tsfresh to do. My first idea was to fit (select features) only on a sample of the train data. This worked well, but the feature extraction during the transform step of the ~70 relevant features was still causing the same problem.
The next idea was scaling out. But before trying to distribute the work using dask, I started to trace at which step the memory starts to inflate and found that it the problem is the [sorting](https://github.com/blue-yonder/tsfresh/blob/a0039838a8264cbe8fb9f512849b778dfc3ccc8c/tsfresh/utilities/dataframe_functions.py#L359) of the [melted](https://github.com/blue-yonder/tsfresh/blob/a0039838a8264cbe8fb9f512849b778dfc3ccc8c/tsfresh/utilities/dataframe_functions.py#L338) dataframe while transforming the data into the normalized format. This doesn't seem too strange if you consider that it is actually sorting 6.024.604 * 75 = 451.845.300 rows.

Because this data transformation is done before sending any data to the distributor, scaling out won't be a solution here.

So I forked tsfresh and started to shift the heavy transformations from the "driver" to the workers and ended with skipping the melt transformation completely. A more straight forward approach was to simply iterate the wide format.

You can find my first draft here:
https://github.com/blue-yonder/tsfresh/compare/v0.16.0...hoesler:distribute-melt?expand=1

Using my modified code, the kernel barely exceeded 8GB during the fit step on the complete train data set. It seemed a bit slower, though I didn't do any benchmarks. But that is totally acceptable for me, especially because this could be easy improved by scaling out.

I would love to finalize that into a pull request, but I'd like to hear your opinion first.

Cheers
Christoph


Discussion:
coveralls: 
[![Coverage Status](https://coveralls.io/builds/31119546/badge)](https://coveralls.io/builds/31119546)

Coverage decreased (-1.4%) to 95.706% when pulling **d647441ac757cad9d409a542e3f8c07bcf5ecce3 on hoesler:distribute-melt** into **29cdae3079fdd8f6dc50ad7eebfbf6e33f0a2b46 on blue-yonder:master**.

nils-braun: @hoesler So first of all thank you very very much for this very good PR!
I just had a very quick over the code, and it looks quite well written. As @MaxBenChrist added himself as a reviewer, I will let him comment first but it already looks quite promising. 

Just to understand: the speed improvement came from the facts that you both do not need to do any pre-processing now and that you do need less memory, or? So the first part will especially be seen in "wide-column" (as you named it) data, but probably less in already melted data - whereas the second part will be seen always, correct?


hoesler: Hi @nils-braun.

Thank you very much. My primary goal was to improve memory consumption. This is archived by directly iterating the input data and skipping costly transformations. I tried to implement this iteration as performant as possible, but did not yet compare it to the speed of the original implementation. Under the new data model, however, it seems that accessing the "wide" dataframe in slices can speed things up, even though I am iterating over the groups multiple times using `itertools.isclice`. I didn't spent time in improving narrow and dict format yet. This still needs to be done.
TKlerx: Nice work @hoesler .
Any comment from @MaxBenChrist ?
I guess the coverage needs to be increased?!
MaxBenChrist: > Nice work @hoesler .
> Any comment from @MaxBenChrist ?
> I guess the coverage needs to be increased?!

sorry, was busy lately. very interesting work @hoesler, will try to understand everything this weekend.

We should aim for high Coverage, yes. but, If we know why the coverage drops and there are good reasons to not test it, then I am also fine with merging prs that reduce coverage.
hoesler: No problem. I am happy to improve coverage, just wanted to wait with the cumbersome work until I get some feedback on the overall design.
nils-braun: Thanks for answering all my questions @hoesler !

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 1ae9735..2082521 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -17,8 +17,11 @@ Unreleased
     - Implemented the benford distribution feature (#689)
     - Reworked the notebooks (#701, #704)
     - Speed up the result pivoting (#705)
+    - Add a test for the dask bindings (#719)
 - Bugfixes:
+    - Fixed a bug in the selection, that caused all regression tasks with un-ordered index to be wrong (#715)
     - Fixed readthedocs (#695, #696)
+    - Fix spark and dask after #705 and for non-id named id columns (#712)
 
 Version 0.16.0
 ==============
diff --git a/tsfresh/convenience/bindings.py b/tsfresh/convenience/bindings.py
index 817484b..47f8756 100644
--- a/tsfresh/convenience/bindings.py
+++ b/tsfresh/convenience/bindings.py
@@ -21,17 +21,20 @@ def _feature_extraction_on_chunk_helper(df, column_id, column_kind,
     elif default_fc_parameters is None and kind_to_fc_parameters is not None:
         default_fc_parameters = {}
 
-    chunk = df[column_id].iloc[0], df[column_kind].iloc[0], df.sort_values(column_sort)[column_value]
+    if column_sort is not None:
+        df = df.sort_values(column_sort)
+
+    chunk = df[column_id].iloc[0], df[column_kind].iloc[0], df[column_value]
     features = _do_extraction_on_chunk(chunk, default_fc_parameters=default_fc_parameters,
                                        kind_to_fc_parameters=kind_to_fc_parameters)
-    features = pd.DataFrame(features)
+    features = pd.DataFrame(features, columns=[column_id, "variable", "value"])
     features["value"] = features["value"].astype("double")
 
-    return features[[column_id, "variable", "value"]]
+    return features
 
 
 def dask_feature_extraction_on_chunk(df, column_id, column_kind,
-                                     column_sort, column_value,
+                                     column_value, column_sort=None,
                                      default_fc_parameters=None, kind_to_fc_parameters=None):
     """
     Extract features on a grouped dask dataframe given the column names and the extraction settings.
@@ -99,7 +102,7 @@ def dask_feature_extraction_on_chunk(df, column_id, column_kind,
     :type column_id: str
 
     :param column_sort: The name of the sort column.
-    :type column_sort: str
+    :type column_sort: str or None
 
     :param column_kind: The name of the column keeping record on the kind of the value.
     :type column_kind: str
@@ -121,8 +124,8 @@ def dask_feature_extraction_on_chunk(df, column_id, column_kind,
 
 
 def spark_feature_extraction_on_chunk(df, column_id, column_kind,
-                                      column_sort, column_value,
-                                      default_fc_parameters, kind_to_fc_parameters=None):
+                                      column_value, column_sort=None,
+                                      default_fc_parameters=None, kind_to_fc_parameters=None):
     """
     Extract features on a grouped spark dataframe given the column names and the extraction settings.
     This wrapper function should only be used if you have a spark dataframe as input.
@@ -184,7 +187,7 @@ def spark_feature_extraction_on_chunk(df, column_id, column_kind,
     :type column_id: str
 
     :param column_sort: The name of the sort column.
-    :type column_sort: str
+    :type column_sort: str or None
 
     :param column_kind: The name of the column keeping record on the kind of the value.
     :type column_kind: str
diff --git a/tsfresh/feature_extraction/data.py b/tsfresh/feature_extraction/data.py
new file mode 100644
index 0000000..17282b1
--- /dev/null
+++ b/tsfresh/feature_extraction/data.py
@@ -0,0 +1,351 @@
+import itertools
+from collections import namedtuple
+from typing import Generator, Iterable, Sized
+
+import pandas as pd
+
+
+class Timeseries(namedtuple('Timeseries', ['id', 'kind', 'data'])):
+    """
+    Timeseries tuple used for feature extraction.
+
+    Make sure `kind` is of type `str` to allow inference
+    of feature settings in `feature_extraction.settings.from_columns`.
+    """
+
+
+class TsData(Iterable[Timeseries], Sized):
+    """
+    TsData provides access to time series data for internal usage.
+
+    Implementations must implement `__iter__` and `__len__`.
+    Other methods should be overwritten if a more efficient solution exists for the underlying data store.
+    """
+
+    def __iter__(self):
+        raise NotImplementedError
+
+    def __len__(self):
+        raise NotImplementedError
+
+    def partition(self, chunk_size):
+        """
+        Split the data into iterable chunks of given `chunk_size`.
+
+        :param chunk_size: the size of the chunks
+        :type chunk_size: int
+
+        :return: chunks with at most `chunk_size` items
+        :rtype: Generator[Iterable[Timeseries]]
+        """
+        iterable = iter(self)
+        while True:
+            next_chunk = list(itertools.islice(iterable, chunk_size))
+            if not next_chunk:
+                return
+
+            yield next_chunk
+
+
+class SliceableTsData(TsData):
+    """
+    SliceableTsData uses a slice strategy to implement data partitioning.
+    Implementations of `TsData` may extend this class if they can iterate their elements in deterministic order and
+    if an efficient strategy exists to slice the data.
+
+    Because `__iter__` defaults to `slice(0)`, implementations must only implement `slice`.
+    """
+
+    def slice(self, offset, length=None):
+        """
+        Get a subset of the timeseries elements
+
+        :param offset: the offset in the sequence of elements
+        :type offset: int
+
+        :param length: if not `None`, the maximum number of elements the slice will yield
+        :type length: int
+
+        :return: a slice of the data
+        :rtype: Iterable[Timeseries]
+        """
+        raise NotImplementedError
+
+    def __iter__(self):
+        return self.slice(0)
+
+    def partition(self, chunk_size):
+        div, mod = divmod(len(self), chunk_size)
+        chunk_info = [(i * chunk_size, chunk_size) for i in range(div)]
+        if mod > 0:
+            chunk_info += [(div * chunk_size, mod)]
+
+        for offset, length in chunk_info:
+            yield _Slice(self, offset, length)
+
+
+class _Slice(Iterable[Timeseries]):
+    def __init__(self, data, offset, length):
+        """
+        Wraps the `slice` generator function as an iterable object which can be pickled and passed to the distributor
+        backend.
+
+        :type data: SliceableTsData
+        :type offset: int
+        :type length: int
+        """
+        self.data = data
+        self.offset = offset
+        self.length = length
+
+    def __iter__(self):
+        return self.data.slice(self.offset, self.length)
+
+    def __len__(self):
+        return self.length
+
+
+def _check_colname(*columns):
+    """
+    Check if given column names conflict with `settings.from_columns` (ends with '_' or contains '__').
+
+    :param columns: the column names to check
+    :type columns: str
+
+    :return: None
+    :rtype: None
+    :raise: ``ValueError`` if column names are invalid.
+    """
+
+    for col in columns:
+        if str(col).endswith("_"):
+            raise ValueError("Dict keys are not allowed to end with '_': {}".format(col))
+
+        if "__" in str(col):
+            raise ValueError("Dict keys are not allowed to contain '__': {}".format(col))
+
+
+def _check_nan(df, *columns):
+    """
+    Raise a ``ValueError`` if one of the columns does not exist or contains NaNs.
+
+    :param df: the pandas DataFrame to test for NaNs
+    :type df: pandas.DataFrame
+    :param columns: a list of columns to test for NaNs. If left empty, all columns of the DataFrame will be tested.
+    :type columns: str
+
+    :return: None
+    :rtype: None
+    :raise: ``ValueError`` if ``NaNs`` are found in the DataFrame.
+    """
+
+    for col in columns:
+        if col not in df.columns:
+            raise ValueError("Column not found: {}".format(col))
+
+        if df[col].isnull().any():
+            raise ValueError("Column must not contain NaN values: {}".format(col))
+
+
+class WideTsFrameAdapter(SliceableTsData):
+
+    def __init__(self, df, column_id, column_sort=None, value_columns=None):
+        """
+        Adapter for Pandas DataFrames in wide format, where multiple columns contain different time series for
+        the same id.
+
+        :param df: the data frame
+        :type df: pd.DataFrame
+
+        :param column_id: the name of the column containing time series group ids
+        :type column_id: str
+
+        :param column_sort: the name of the column to sort on
+        :type column_sort: str|None
+
+        :param value_columns: list of column names to treat as time series values.
+            If `None`, all columns except `column_id` and `column_sort` will be used.
+        :type value_columns: list[str]|None
+        """
+
+        _check_nan(df, column_id)
+
+        if value_columns is None:
+            value_columns = [col for col in df.columns if col not in [column_id, column_sort]]
+
+        if len(value_columns) == 0:
+            raise ValueError("You must provide at least one value column")
+
+        _check_nan(df, *value_columns)
+        _check_colname(*value_columns)
+
+        self.value_columns = value_columns
+
+        if column_sort is not None:
+            _check_nan(df, column_sort)
+            self.df_grouped = df.sort_values([column_sort]).groupby([column_id])
+        else:
+            self.df_grouped = df.groupby([column_id])
+
+    def __len__(self):
+        return self.df_grouped.ngroups * len(self.value_columns)
+
+    def slice(self, offset, length=None):
+        if 0 < offset >= len(self):
+            raise ValueError("offset out of range: {}".format(offset))
+
+        group_offset, column_offset = divmod(offset, len(self.value_columns))
+        kinds = self.value_columns[column_offset:]
+        i = 0
+
+        for group_name, group in itertools.islice(self.df_grouped, group_offset, None):
+            for kind in kinds:
+                i += 1
+                if length is not None and i > length:
+                    return
+                yield Timeseries(group_name, kind, group[kind])
+            kinds = self.value_columns
+
+
+class LongTsFrameAdapter(TsData):
+
+    def __init__(self, df, column_id, column_kind, column_value, column_sort=None):
+        """
+        Adapter for Pandas DataFrames in long format, where different time series for the same id are
+        labeled by column `column_kind`.
+
+        :param df: the data frame
+        :type df: pd.DataFrame
+
+        :param column_id: the name of the column containing time series group ids
+        :type column_id: str
+
+        :param column_kind: the name of the column containing time series kinds for each id
+        :type column_kind: str
+
+        :param column_value: the name of the column containing time series values
+        :type column_value: str
+
+        :param column_sort: the name of the column to sort on
+        :type column_sort: str|None
+        """
+
+        _check_nan(df, column_id, column_kind, column_value)
+        _check_colname(column_kind)
+
+        self.column_value = column_value
+
+        if column_sort is not None:
+            _check_nan(df, column_sort)
+            self.df_grouped = df.sort_values([column_sort]).groupby([column_id, column_kind])
+        else:
+            self.df_grouped = df.groupby([column_id, column_kind])
+
+    def __len__(self):
+        return len(self.df_grouped)
+
+    def __iter__(self):
+        for group_key, group in self.df_grouped:
+            yield Timeseries(group_key[0], str(group_key[1]), group[self.column_value])
+
+
+class TsDictAdapter(TsData):
+    def __init__(self, ts_dict, column_id, column_value, column_sort=None):
+        """
+        Adapter for a dict, which maps different time series kinds to Pandas DataFrames.
+
+        :param ts_dict: a dict of data frames
+        :type ts_dict: dict[str, pd.DataFrame]
+
+        :param column_id: the name of the column containing time series group ids
+        :type column_id: str
+
+        :param column_value: the name of the column containing time series values
+        :type column_value: str
+
+        :param column_sort: the name of the column to sort on
+        :type column_sort: str|None
+        """
+
+        _check_colname(*list(ts_dict.keys()))
+        for df in ts_dict.values():
+            _check_nan(df, column_id, column_value)
+
+        self.column_value = column_value
+
+        if column_sort is not None:
+            for key, df in ts_dict.items():
+                _check_nan(df, column_sort)
+
+            self.grouped_dict = {key: df.sort_values([column_sort]).groupby(column_id)
+                                 for key, df in ts_dict.items()}
+        else:
+            self.grouped_dict = {key: df.groupby(column_id) for key, df in ts_dict.items()}
+
+    def __iter__(self):
+        for kind, grouped_df in self.grouped_dict.items():
+            for ts_id, group in grouped_df:
+                yield Timeseries(ts_id, str(kind), group[self.column_value])
+
+    def __len__(self):
+        return sum(grouped_df.ngroups for grouped_df in self.grouped_dict.values())
+
+
+def to_tsdata(df, column_id=None, column_kind=None, column_value=None, column_sort=None):
+    """
+    Wrap supported data formats as a TsData object, i.e. an iterable of individual time series.
+
+    E.g. the DataFrame
+
+        ====  ======  =========
+          id  kind          val
+        ====  ======  =========
+           1  a       -0.21761
+           1  a       -0.613667
+           1  a       -2.07339
+           2  b       -0.576254
+           2  b       -1.21924
+        ====  ======  =========
+
+    into
+
+        Iterable((1, 'a', pd.Series([-0.217610, -0.613667, -2.073386]),
+         (2, 'b', pd.Series([-0.576254, -1.219238]))
+
+    :param df: one of the supported input formats
+    :type df: pd.DataFrame|dict|TsData
+
+    :param column_id: The name of the id column to group by.
+    :type column_id: str|None
+
+    :param column_kind: The name of the column keeping record on the kind of the value.
+    :type column_kind: str|None
+
+    :param column_value: The name for the column keeping the value itself.
+    :type column_value: str|None
+
+    :param column_sort: The name for the column to sort on.
+    :type column_sort: str|None
+
+    :return: a data adapter
+    :rtype: TsData
+    """
+
+    if isinstance(df, TsData):
+        return df
+
+    elif isinstance(df, pd.DataFrame):
+        if column_value is not None:
+            if column_kind is not None:
+                return LongTsFrameAdapter(df, column_id, column_kind, column_value, column_sort)
+            else:
+                return WideTsFrameAdapter(df, column_id, column_sort, [column_value])
+        else:
+            return WideTsFrameAdapter(df, column_id, column_sort)
+
+    elif isinstance(df, dict):
+        return TsDictAdapter(df, column_id, column_value, column_sort)
+
+    else:
+        raise ValueError("df must be a DataFrame or a dict of DataFrames. "
+                         "See https://tsfresh.readthedocs.io/en/latest/text/data_formats.html")
diff --git a/tsfresh/feature_extraction/extraction.py b/tsfresh/feature_extraction/extraction.py
index 272d9fb..0aafd01 100644
--- a/tsfresh/feature_extraction/extraction.py
+++ b/tsfresh/feature_extraction/extraction.py
@@ -12,8 +12,9 @@ import pandas as pd
 
 from tsfresh import defaults
 from tsfresh.feature_extraction import feature_calculators
+from tsfresh.feature_extraction.data import to_tsdata
 from tsfresh.feature_extraction.settings import ComprehensiveFCParameters
-from tsfresh.utilities import dataframe_functions, profiling
+from tsfresh.utilities import profiling
 from tsfresh.utilities.distribution import MapDistributor, MultiprocessingDistributor, \
     DistributorBaseClass
 from tsfresh.utilities.string_manipulation import convert_to_output_format
@@ -130,12 +131,7 @@ def extract_features(timeseries_container, default_fc_parameters=None,
 
     # Always use the standardized way of storing the data.
     # See the function normalize_input_to_internal_representation for more information.
-    df_melt, column_id, column_kind, column_value = \
-        dataframe_functions._normalize_input_to_internal_representation(
-            timeseries_container=timeseries_container,
-            column_id=column_id, column_kind=column_kind,
-            column_sort=column_sort,
-            column_value=column_value)
+
     # Use the standard setting if the user did not supply ones himself.
     if default_fc_parameters is None and kind_to_fc_parameters is None:
         default_fc_parameters = ComprehensiveFCParameters()
@@ -152,9 +148,10 @@ def extract_features(timeseries_container, default_fc_parameters=None,
         else:
             warnings.simplefilter("default")
 
-        result = _do_extraction(df=df_melt,
+        result = _do_extraction(df=timeseries_container,
                                 column_id=column_id, column_value=column_value,
                                 column_kind=column_kind,
+                                column_sort=column_sort,
                                 n_jobs=n_jobs, chunk_size=chunksize,
                                 disable_progressbar=disable_progressbar,
                                 show_warnings=show_warnings,
@@ -174,69 +171,7 @@ def extract_features(timeseries_container, default_fc_parameters=None,
     return result
 
 
-def generate_data_chunk_format(df, column_id, column_kind, column_value):
-    """Converts the dataframe df in into a list of individual time seriess.
-
-    E.g. the DataFrame
-
-        ====  ======  =========
-          id  kind          val
-        ====  ======  =========
-           1  a       -0.21761
-           1  a       -0.613667
-           1  a       -2.07339
-           2  b       -0.576254
-           2  b       -1.21924
-        ====  ======  =========
-
-    into
-
-        [(1, 'a', pd.Series([-0.217610, -0.613667, -2.073386]),
-         (2, 'b', pd.Series([-0.576254, -1.219238])]
-
-
-    The data is separated out into those single time series and the _do_extraction_on_chunk is
-    called on each of them. The results are then combined into a single pandas DataFrame.
-
-    The call is either happening in parallel or not and is showing a progress bar or not depending
-    on the given flags.
-
-    :param df: The dataframe in the normalized format which is used for extraction.
-    :type df: pd.DataFrame
-
-    :param column_id: The name of the id column to group by.
-    :type column_id: str
-
-    :param column_kind: The name of the column keeping record on the kind of the value.
-    :type column_kind: str
-
-    :param column_value: The name for the column keeping the value itself.
-    :type column_value: str
-
-    :return: the data in chunks
-    :rtype: list
-    """
-    MAX_VALUES_GROUPBY = 2147483647
-
-    if df[[column_id, column_kind]].nunique().prod() >= MAX_VALUES_GROUPBY:
-        _logger.error(
-            "The time series container has {} different ids and {} different kind of time series, "
-            "in total {} possible combinations. "
-            "Due to a limitation in pandas we can only process a maximum of {} id/kind combinations. "
-            "Please reduce your time series container and restart "
-            "the calculation".format(
-                df[column_id].nunique(), df[column_kind].nunique(),
-                df[[column_id, column_kind]].nunique().prod(), MAX_VALUES_GROUPBY)
-        )
-        raise ValueError(
-            "Number of ids/kinds are too high. Please reduce your data size and run feature extraction again.")
-    data_in_chunks = [x + (y,) for x, y in
-                      df.groupby([column_id, column_kind], as_index=True)[column_value]]
-
-    return data_in_chunks
-
-
-def _do_extraction(df, column_id, column_value, column_kind,
+def _do_extraction(df, column_id, column_value, column_kind, column_sort,
                    default_fc_parameters, kind_to_fc_parameters,
                    n_jobs, chunk_size, disable_progressbar, show_warnings, distributor):
     """
@@ -291,7 +226,7 @@ def _do_extraction(df, column_id, column_value, column_kind,
     :rtype: pd.DataFrame
     """
 
-    data_in_chunks = generate_data_chunk_format(df, column_id, column_kind, column_value)
+    data = to_tsdata(df, column_id, column_kind, column_value, column_sort)
 
     if distributor is None:
         if n_jobs == 0:
@@ -309,7 +244,7 @@ def _do_extraction(df, column_id, column_value, column_kind,
     kwargs = dict(default_fc_parameters=default_fc_parameters,
                   kind_to_fc_parameters=kind_to_fc_parameters)
 
-    result = distributor.map_reduce(_do_extraction_on_chunk, data=data_in_chunks,
+    result = distributor.map_reduce(_do_extraction_on_chunk, data=data,
                                     chunk_size=chunk_size,
                                     function_kwargs=kwargs)
     distributor.close()
diff --git a/tsfresh/feature_selection/relevance.py b/tsfresh/feature_selection/relevance.py
index 7697352..b6d1e20 100644
--- a/tsfresh/feature_selection/relevance.py
+++ b/tsfresh/feature_selection/relevance.py
@@ -131,6 +131,13 @@ def calculate_relevance_table(X, y, ml_task='auto', n_jobs=defaults.N_PROCESSES,
              not relevant] for this feature)
     :rtype: pandas.DataFrame
     """
+
+    # Make sure X and y both have the exact same indices
+    y = y.sort_index()
+    X = X.sort_index()
+
+    assert list(y.index) == list(X.index), "The index of X and y need to be the same"
+
     if ml_task not in ['auto', 'classification', 'regression']:
         raise ValueError('ml_task must be one of: \'auto\', \'classification\', \'regression\'')
     elif ml_task == 'auto':
diff --git a/tsfresh/utilities/dataframe_functions.py b/tsfresh/utilities/dataframe_functions.py
index d18c546..9d20e3a 100644
--- a/tsfresh/utilities/dataframe_functions.py
+++ b/tsfresh/utilities/dataframe_functions.py
@@ -5,16 +5,15 @@
 Utility functions for handling the DataFrame conversions to the internal normalized format
 (see ``normalize_input_to_internal_representation``) or on how to handle ``NaN`` and ``inf`` in the DataFrames.
 """
-import gc
 import warnings
 from collections import defaultdict
 
-from tsfresh import defaults
-from tsfresh.utilities.distribution import MapDistributor, MultiprocessingDistributor, DistributorBaseClass
-
 import numpy as np
 import pandas as pd
 
+from tsfresh import defaults
+from tsfresh.utilities.distribution import MapDistributor, MultiprocessingDistributor, DistributorBaseClass
+
 
 def check_for_nans_in_columns(df, columns=None):
     """
@@ -240,149 +239,6 @@ def get_ids(df_or_dict, column_id):
         raise TypeError("df_or_dict should be of type dict or pandas.DataFrame")
 
 
-# todo: add more testcases
-# todo: rewrite in a more straightforward way
-def _normalize_input_to_internal_representation(timeseries_container, column_id, column_sort,
-                                                column_kind, column_value):
-    """
-    Try to transform any given input to the internal representation of time series, which is a flat DataFrame
-    (the first format from see :ref:`data-formats-label`).
-
-    This function can transform pandas DataFrames in different formats or dictionaries into the internal format
-    that we use. It should not be called by the user.
-
-    :param timeseries_container: a pandas DataFrame or a dictionary. The required shape/form of the object depends on
-        the rest of the passed arguments.
-    :type timeseries_container: pandas.DataFrame or dict
-    :param column_id: it must be present in the pandas DataFrame or in all DataFrames in the dictionary.
-        It is not allowed to have NaN values in this column.
-    :type column_id: basestring
-    :param column_sort: if not None, sort the rows by this column. It is not allowed to
-        have NaN values in this column.
-    :type column_sort: basestring or None
-    :param column_kind: It can only be used when passing a pandas DataFrame (the dictionary is already assumed to be
-        grouped by the kind). Is must be present in the DataFrame and no NaN values are allowed. The DataFrame
-        will be grouped by the values in the kind column and each group will be one entry in the resulting
-        mapping.
-        If the kind column is not passed, it is assumed that each column in the pandas DataFrame (except the id or
-        sort column) is a possible kind and the DataFrame is split up into as many DataFrames as there are columns.
-        It is not allowed to have a value column then.
-    :type column_kind: basestring or None
-    :param column_value: If it is given, it must be present and not-NaN on the pandas DataFrames (or all pandas
-        DataFrames in the dictionaries). If it is None, the kind column must also be none.
-    :type column_value: basestring or None
-
-    :return: A tuple of 4 elements: the normalized DataFrame, the name of the id column, the name of the value column
-             and the name of the value column
-    :rtype: (pd.DataFrame, basestring, basestring, basestring)
-    :raise: ``ValueError`` when the passed combination of parameters is wrong or does not fit to the input DataFrame
-            or dict.
-    """
-    # Also make it possible to have a dict as an input
-    if isinstance(timeseries_container, dict):
-        if column_kind is not None:
-            raise ValueError("You passed in a dictionary and gave a column name for the kind. Both are not possible.")
-
-        column_kind = "_variables"
-
-        timeseries_container = {key: df.copy() for key, df in timeseries_container.items()}
-
-        for kind, df in timeseries_container.items():
-            df[column_kind] = kind
-
-        try:
-            timeseries_container = pd.concat(timeseries_container.values(), sort=True)
-        except TypeError:  # pandas < 0.23.0
-            timeseries_container = pd.concat(timeseries_container.values())
-        gc.collect()
-
-    # Check ID column
-    if column_id is None:
-        raise ValueError("You have to set the column_id which contains the ids of the different time series")
-
-    if column_id not in timeseries_container.columns:
-        raise AttributeError("The given column for the id is not present in the data.")
-
-    if timeseries_container[column_id].isnull().any():
-        raise ValueError("You have NaN values in your id column.")
-
-    # Check sort column
-    if column_sort is not None:
-        if timeseries_container[column_sort].isnull().any():
-            raise ValueError("You have NaN values in your sort column.")
-
-    # Check that either kind and value is None or both not None.
-    if column_kind is None and column_value is not None:
-        column_kind = "_variables"
-        timeseries_container = timeseries_container.copy()
-        timeseries_container[column_kind] = column_value
-    if column_kind is not None and column_value is None:
-        raise ValueError("If passing the kind, you also have to pass the value.")
-
-    if column_kind is None and column_value is None:
-        if column_sort is not None:
-            sort = timeseries_container[column_sort].values
-            timeseries_container = timeseries_container.drop(column_sort, axis=1)
-        else:
-            sort = range(len(timeseries_container))
-            column_sort = "_sort"
-
-        column_kind = "_variables"
-        column_value = "_values"
-
-        if not set(timeseries_container.columns) - {column_id}:
-            raise ValueError("There is no column with values in your data!")
-
-        # We need to preserve the index. However, pandas has hard times to parse the columns if they
-        # have different types, so we need to preserve them.
-        # At least until https://github.com/pandas-dev/pandas/pull/28859 is merged.
-        if isinstance(column_id, int) or isinstance(column_id, float):
-            # some arbitrary number
-            index_name = column_id + 999
-        else:
-            index_name = "_temporary_index_column"
-
-        timeseries_container.index.name = index_name
-        timeseries_container = pd.melt(timeseries_container.reset_index(),
-                                       id_vars=[index_name, column_id],
-                                       value_name=column_value, var_name=column_kind)
-        timeseries_container = timeseries_container.set_index(index_name)
-        timeseries_container[column_sort] = np.tile(sort, (len(timeseries_container) // len(sort)))
-
-    # Check kind column
-    if column_kind not in timeseries_container.columns:
-        raise AttributeError("The given column for the kind is not present in the data.")
-
-    if timeseries_container[column_kind].isnull().any():
-        raise ValueError("You have NaN values in your kind column.")
-
-    # Check value column
-    if column_value not in timeseries_container.columns:
-        raise ValueError("The given column for the value is not present in the data.")
-
-    if timeseries_container[column_value].isnull().any():
-        raise ValueError("You have NaN values in your value column.")
-
-    if column_sort:
-        timeseries_container = timeseries_container.sort_values([column_id, column_kind, column_sort])
-        timeseries_container = timeseries_container.drop(column_sort, axis=1)
-    else:
-        timeseries_container = timeseries_container.sort_values([column_id, column_kind])
-
-    # The kind columns should always be of type "str" to make the inference of feature settings later in `from_columns`
-    # work
-    timeseries_container[column_kind] = timeseries_container[column_kind].astype(str)
-
-    # Make sure we have only parsable names
-    for kind in timeseries_container[column_kind].unique():
-        if kind.endswith("_"):
-            raise ValueError("The kind {kind} is not allowed to end with '_'".format(kind=kind))
-        if "__" in kind:
-            raise ValueError("The kind {kind} is not allowed to contain '__'".format(kind=kind))
-
-    return timeseries_container, column_id, column_kind, column_value
-
-
 def _roll_out_time_series(timeshift, grouped_data, rolling_direction, max_timeshift, min_timeshift,
                           column_sort, column_id):
     """
diff --git a/tsfresh/utilities/distribution.py b/tsfresh/utilities/distribution.py
index 0490b8a..815cfd0 100644
--- a/tsfresh/utilities/distribution.py
+++ b/tsfresh/utilities/distribution.py
@@ -8,14 +8,17 @@ Essentially, a Distributor organizes the application of feature calculators to d
 Design of this module by Nils Braun
 """
 
-import math
 import itertools
+import math
 import warnings
-from collections import Iterable
+from collections import Generator, Iterable
 from functools import partial
 from multiprocessing import Pool
+
 from tqdm import tqdm
 
+from tsfresh.feature_extraction.data import TsData
+
 
 def _function_with_partly_reduce(chunk_list, map_function, kwargs):
     """
@@ -74,25 +77,30 @@ class DistributorBaseClass:
     @staticmethod
     def partition(data, chunk_size):
         """
-        This generator chunks a list of data into slices of length chunk_size. If the chunk_size is not a divider of the
-        data length, the last slice will be shorter than chunk_size.
+        This generator partitions an iterable into slices of length `chunk_size`.
+        If the chunk size is not a divider of the data length, the last slice will be shorter.
 
-        :param data: The data to chunk.
-        :type data: list
-        :param chunk_size: Each chunks size. The last chunk may be smaller.
+        :param data: The data to partition.
+        :type data: Iterable
+        :param chunk_size: The chunk size. The last chunk might be smaller.
         :type chunk_size: int
 
         :return: A generator producing the chunks of data.
-        :rtype: generator
+        :rtype: Generator[Iterable]
         """
 
-        iterable = iter(data)
-        while True:
-            next_chunk = list(itertools.islice(iterable, chunk_size))
-            if not next_chunk:
-                return
-
-            yield next_chunk
+        if isinstance(data, TsData):
+            return data.partition(chunk_size)
+        else:
+            def partition_iterable():
+                iterable = iter(data)
+                while True:
+                    next_chunk = list(itertools.islice(iterable, chunk_size))
+                    if not next_chunk:
+                        return
+                    yield next_chunk
+
+            return partition_iterable()
 
     def __init__(self):
         """


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
index cba7937..d647178 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -11,6 +11,8 @@ import pandas as pd
 # todo: add test cases for float data
 # todo: add test cases for nans, -infs, infs
 # todo: add test cases with time series of length one
+
+
 class DataTestCase(TestCase):
     def create_test_data_sample(self):
         cid = np.repeat([10, 500], 40)
@@ -30,6 +32,53 @@ class DataTestCase(TestCase):
         df.index.name = None
         return df
 
+    def create_test_data_sample_wide(self):
+        rec = np.rec.array([
+            (0, 10, 0, 11, 50),
+            (1, 10, 1, 9, 40),
+            (2, 10, 2, 67, 39),
+            (3, 10, 3, 45, 7),
+            (4, 10, 4, 30, 53),
+            (5, 10, 5, 58, 23),
+            (6, 10, 6, 62, 16),
+            (7, 10, 7, 19, 37),
+            (8, 10, 8, 56, 66),
+            (9, 10, 9, 29, 38),
+            (10, 10, 10, 0, 6),
+            (11, 10, 11, 27, 47),
+            (12, 10, 12, 36, 3),
+            (13, 10, 13, 43, 61),
+            (14, 10, 14, 33, 44),
+            (15, 10, 15, 2, 42),
+            (16, 10, 16, 24, 78),
+            (17, 10, 17, 71, 31),
+            (18, 10, 18, 41, 21),
+            (19, 10, 19, 28, 55),
+            (20, 500, 0, 15, 4),
+            (21, 500, 1, 35, 14),
+            (22, 500, 2, 25, 68),
+            (23, 500, 3, 32, 22),
+            (24, 500, 4, 69, 18),
+            (25, 500, 5, 65, 52),
+            (26, 500, 6, 70, 54),
+            (27, 500, 7, 64, 60),
+            (28, 500, 8, 51, 79),
+            (29, 500, 9, 46, 12),
+            (30, 500, 10, 5, 49),
+            (31, 500, 11, 77, 63),
+            (32, 500, 12, 26, 8),
+            (33, 500, 13, 73, 59),
+            (34, 500, 14, 76, 1),
+            (35, 500, 15, 75, 13),
+            (36, 500, 16, 72, 20),
+            (37, 500, 17, 74, 17),
+            (38, 500, 18, 10, 48),
+            (39, 500, 19, 57, 34)],
+            dtype=[('index', '<i8'), ('id', '<i8'), ('sort', '<i8'), ('a', '<i8'), ('b', '<i8')])
+        df = pd.DataFrame.from_records(rec)
+        df = df.set_index("index", drop=True)
+        return df
+
     def create_test_data_sample_with_time_index(self):
         cid = np.repeat([10, 500], 40)
         ckind = np.repeat(["a", "b", "a", "b"], 20)
diff --git a/tests/integrations/examples/test_driftbif_simulation.py b/tests/integrations/examples/test_driftbif_simulation.py
index f8e0748..d07abeb 100644
--- a/tests/integrations/examples/test_driftbif_simulation.py
+++ b/tests/integrations/examples/test_driftbif_simulation.py
@@ -7,6 +7,7 @@ import unittest
 import pandas as pd
 
 from tsfresh.examples.driftbif_simulation import velocity, load_driftbif, sample_tau
+from tsfresh import extract_relevant_features
 
 
 class DriftBifSimlationTestCase(unittest.TestCase):
@@ -63,6 +64,18 @@ class DriftBifSimlationTestCase(unittest.TestCase):
         self.assertEqual(v.shape, (Nt, 3),
                          'The returned vector should reflect the dimension of the initial condition.')
 
+    def test_relevant_feature_extraction(self):
+        df, y = load_driftbif(100, 10, classification=False, seed=42)
+
+        df['id'] = df['id'].astype('str')
+        y.index = y.index.astype('str')
+
+        X = extract_relevant_features(df, y,
+                                      column_id="id", column_sort="time",
+                                      column_kind="dimension", column_value="value")
+
+        self.assertGreater(len(X.columns), 10)
+
 
 class SampleTauTestCase(unittest.TestCase):
     def test_range(self):
diff --git a/tests/integrations/test_bindings.py b/tests/integrations/test_bindings.py
new file mode 100644
index 0000000..832808b
--- /dev/null
+++ b/tests/integrations/test_bindings.py
@@ -0,0 +1,33 @@
+from unittest import TestCase
+
+from tsfresh.convenience.bindings import dask_feature_extraction_on_chunk
+from tsfresh.feature_extraction.settings import MinimalFCParameters
+
+from dask import dataframe as dd
+import pandas as pd
+
+
+class DaskBindingsTestCase(TestCase):
+    def test_feature_extraction(self):
+        df = pd.DataFrame({"my_id": [1, 1, 1, 2, 2, 2], "my_kind": ["a"]*6,
+                           "my_value": [1, 2, 3, 4, 5, 6]})
+
+        df = dd.from_pandas(df, chunksize=3)
+
+        df_grouped = df.groupby(["my_id", "my_kind"])
+
+        features = dask_feature_extraction_on_chunk(df_grouped, column_id="my_id",
+                                                    column_kind="my_kind",
+                                                    column_value="my_value",
+                                                    column_sort=None,
+                                                    default_fc_parameters=MinimalFCParameters())
+
+        features = features.categorize(columns=["variable"])
+        features = features.reset_index(drop=True)
+
+        feature_table = features.pivot_table(index="my_id", columns="variable", values="value", aggfunc="sum")
+
+        feature_table = feature_table.compute()
+
+        self.assertEqual(len(feature_table.columns), len(MinimalFCParameters()))
+        self.assertEqual(len(feature_table), 2)
diff --git a/tests/units/feature_extraction/test_data.py b/tests/units/feature_extraction/test_data.py
new file mode 100644
index 0000000..b31b9a9
--- /dev/null
+++ b/tests/units/feature_extraction/test_data.py
@@ -0,0 +1,173 @@
+import numpy as np
+import pandas as pd
+
+from tests.fixtures import DataTestCase
+from tsfresh.feature_extraction.data import to_tsdata, LongTsFrameAdapter, WideTsFrameAdapter, TsDictAdapter
+from tsfresh.utilities.distribution import MultiprocessingDistributor
+
+TEST_DATA_EXPECTED_TUPLES = \
+    [(10, 'a', pd.Series([36, 71, 27, 62, 56, 58, 67, 11, 2, 24, 45, 30, 0,
+                          9, 41, 28, 33, 19, 29, 43],
+                         index=[10] * 20)),
+     (10, 'b', pd.Series([78, 37, 23, 44, 6, 3, 21, 61, 39, 31, 53, 16, 66,
+                          50, 40, 47, 7, 42, 38, 55],
+                         index=[10] * 20)),
+     (500, 'a', pd.Series([76, 72, 74, 75, 32, 64, 46, 35, 15, 70, 57, 65,
+                           51, 26, 5, 25, 10, 69, 73, 77],
+                          index=[500] * 20)),
+     (500, 'b', pd.Series([8, 60, 12, 68, 22, 17, 18, 63, 49, 34, 20, 52,
+                           48, 14, 79, 4, 1, 59, 54, 13],
+                          index=[500] * 20))]
+
+WIDE_TEST_DATA_EXPECTED_TUPLES = \
+    [(10, 'a', pd.Series([11, 9, 67, 45, 30, 58, 62, 19, 56, 29, 0, 27, 36,
+                          43, 33, 2, 24, 71, 41, 28],
+                         index=list(range(20)))),
+     (10, 'b', pd.Series([50, 40, 39, 7, 53, 23, 16, 37, 66, 38, 6, 47, 3,
+                          61, 44, 42, 78, 31, 21, 55],
+                         index=list(range(20)))),
+     (500, 'a', pd.Series([15, 35, 25, 32, 69, 65, 70, 64, 51, 46, 5, 77,
+                           26, 73, 76, 75, 72, 74, 10, 57],
+                          index=list(range(20, 40)))),
+     (500, 'b', pd.Series([4, 14, 68, 22, 18, 52, 54, 60, 79, 12, 49, 63,
+                           8, 59, 1, 13, 20, 17, 48, 34],
+                          index=list(range(20, 40))))]
+
+
+class DataAdapterTestCase(DataTestCase):
+
+    def test_long_tsframe(self):
+        df = self.create_test_data_sample()
+        data = LongTsFrameAdapter(df, "id", "kind", "val", "sort")
+
+        self.assert_tsdata(data, TEST_DATA_EXPECTED_TUPLES)
+
+    def test_wide_tsframe(self):
+        df = self.create_test_data_sample_wide()
+        data = WideTsFrameAdapter(df, "id", "sort")
+
+        self.assert_tsdata(data, WIDE_TEST_DATA_EXPECTED_TUPLES)
+
+    def test_dict_tsframe(self):
+        df = {key: df for key, df in self.create_test_data_sample().groupby(["kind"])}
+        data = TsDictAdapter(df, "id", "val", "sort")
+
+        self.assert_tsdata(data, TEST_DATA_EXPECTED_TUPLES)
+
+    def assert_tsdata(self, data, expected):
+        self.assertEqual(len(data), len(expected))
+        self.assertEqual(sum(1 for _ in data), len(data))
+        self.assertEqual(sum(1 for _ in data.partition(1)), len(expected))
+        self.assertEqual((sum(sum(1 for _ in g) for g in data.partition(1))), len(data))
+        self.assert_data_chunk_object_equal(data, expected)
+
+    def assert_data_chunk_object_equal(self, result, expected):
+        dic_result = {str(x[0]) + "_" + str(x[1]): x[2] for x in result}
+        dic_expected = {str(x[0]) + "_" + str(x[1]): x[2] for x in expected}
+        for k in dic_result.keys():
+            pd.testing.assert_series_equal(dic_result[k], dic_expected[k], check_names=False)
+
+    def test_simple_data_sample_two_timeseries(self):
+        df = pd.DataFrame({"id": [10] * 4, "kind": ["a"] * 2 + ["b"] * 2, "val": [36, 71, 78, 37]})
+        df.set_index("id", drop=False, inplace=True)
+        df.index.name = None
+
+        result = to_tsdata(df, "id", "kind", "val")
+        expected = [(10, 'a', pd.Series([36, 71], index=[10] * 2, name="val")),
+                    (10, 'b', pd.Series([78, 37], index=[10] * 2, name="val"))]
+        self.assert_data_chunk_object_equal(result, expected)
+
+    def test_simple_data_sample_four_timeseries(self):
+        df = self.create_test_data_sample()
+        df.index.name = None
+        df.sort_values(by=["id", "kind", "sort"], inplace=True)
+
+        result = to_tsdata(df, "id", "kind", "val", "sort")
+        expected = TEST_DATA_EXPECTED_TUPLES
+
+        self.assert_data_chunk_object_equal(result, expected)
+
+    def test_with_dictionaries_two_rows(self):
+        test_df = pd.DataFrame([{"value": 2, "sort": 2, "id": "id_1"},
+                                {"value": 1, "sort": 1, "id": "id_1"}])
+        test_dict = {"a": test_df, "b": test_df}
+
+        result = to_tsdata(test_dict, column_id="id", column_value="value", column_sort="sort")
+        expected = [("id_1", 'a', pd.Series([1, 2], index=[1, 0], name="value")),
+                    ("id_1", 'b', pd.Series([1, 2], index=[1, 0], name="value"))]
+        self.assert_data_chunk_object_equal(result, expected)
+
+    def test_wide_dataframe_order_preserved_with_sort_column(self):
+        """ verifies that the order of the sort column from a wide time series container is preserved
+        """
+
+        test_df = pd.DataFrame({'id': ["a", "a", "b"],
+                                'v1': [3, 2, 1],
+                                'v2': [13, 12, 11],
+                                'sort': [103, 102, 101]})
+
+        result = to_tsdata(test_df, column_id="id", column_sort="sort")
+        expected = [("a", 'v1', pd.Series([2, 3], index=[1, 0], name="v1")),
+                    ("a", 'v2', pd.Series([12, 13], index=[1, 0], name="v2")),
+                    ("b", 'v1', pd.Series([1], index=[2], name="v1")),
+                    ("b", 'v2', pd.Series([11], index=[2], name="v2"))]
+        self.assert_data_chunk_object_equal(result, expected)
+
+    def test_with_wrong_input(self):
+        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": np.NaN}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", "kind", "value", "sort")
+
+        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "strange_id", "kind", "value", "sort")
+
+        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", "strange_kind", "value", "sort")
+
+        test_df = pd.DataFrame([{"id": np.NaN, "kind": "a", "value": 3, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", "kind", "value", "sort")
+
+        test_df = pd.DataFrame([{"id": 0, "kind": np.NaN, "value": 3, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", "kind", "value", "sort")
+
+        test_df = pd.DataFrame([{"id": 2}, {"id": 1}])
+        test_dict = {"a": test_df, "b": test_df}
+
+        # If there are more than one column, the algorithm can not choose the correct column
+        self.assertRaises(ValueError, to_tsdata, test_dict,
+                          "id", None, None, None)
+
+        test_dict = {"a": pd.DataFrame([{"id": 2, "value_a": 3}, {"id": 1, "value_a": 4}]),
+                     "b": pd.DataFrame([{"id": 2}, {"id": 1}])}
+
+        # If there are more than one column, the algorithm can not choose the correct column
+        self.assertRaises(ValueError, to_tsdata, test_dict,
+                          "id", None, None, None)
+
+        test_df = pd.DataFrame([{"id": 0, "value": np.NaN}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", None, "value", None)
+
+        test_df = pd.DataFrame([{"id": 0, "value": np.NaN}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          None, None, "value", None)
+
+        test_df = pd.DataFrame([{"id": 0, "a_": 3, "b": 5, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", None, None, "sort")
+
+        test_df = pd.DataFrame([{"id": 0, "a__c": 3, "b": 5, "sort": 1}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", None, None, "sort")
+
+        test_df = pd.DataFrame([{"id": 0}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", None, None, None)
+
+        test_df = pd.DataFrame([{"id": 0, "sort": 0}])
+        self.assertRaises(ValueError, to_tsdata, test_df,
+                          "id", None, None, "sort")
diff --git a/tests/units/feature_extraction/test_extraction.py b/tests/units/feature_extraction/test_extraction.py
index 11c4f5a..ab01786 100644
--- a/tests/units/feature_extraction/test_extraction.py
+++ b/tests/units/feature_extraction/test_extraction.py
@@ -10,7 +10,7 @@ import pandas as pd
 from mock import Mock
 
 from tests.fixtures import DataTestCase
-from tsfresh.feature_extraction.extraction import extract_features, generate_data_chunk_format
+from tsfresh.feature_extraction.extraction import extract_features
 from tsfresh.feature_extraction.settings import ComprehensiveFCParameters
 from tsfresh.utilities.distribution import DistributorBaseClass
 
@@ -267,44 +267,3 @@ class DistributorUsageTestCase(DataTestCase):
 
         self.assertTrue(mock.close.called)
         self.assertTrue(mock.map_reduce.called)
-
-
-class GenerateDataChunkTestCase(DataTestCase):
-
-    def assert_data_chunk_object_equal(self, result, expected):
-        dic_result = {str(x[0]) + "_" + str(x[1]): x[2] for x in result}
-        dic_expected = {str(x[0]) + "_" + str(x[1]): x[2] for x in expected}
-        for k in dic_result.keys():
-            pd.testing.assert_series_equal(dic_result[k], dic_expected[k])
-
-    def test_simple_data_sample_two_timeseries(self):
-        df = pd.DataFrame({"id": [10] * 4, "kind": ["a"] * 2 + ["b"] * 2, "val": [36, 71, 78, 37]})
-        df.set_index("id", drop=False, inplace=True)
-        df.index.name = None
-
-        result = generate_data_chunk_format(df, "id", "kind", "val")
-        expected = [(10, 'a', pd.Series([36, 71], index=[10] * 2, name="val")),
-                    (10, 'b', pd.Series([78, 37], index=[10] * 2, name="val"))]
-        self.assert_data_chunk_object_equal(result, expected)
-
-    def test_simple_data_sample_four_timeseries(self):
-        df = self.create_test_data_sample()
-        # todo: investigate the names that are given
-        df.index.name = None
-        df.sort_values(by=["id", "kind", "sort"], inplace=True)
-
-        result = generate_data_chunk_format(df, "id", "kind", "val")
-        expected = [(10, 'a', pd.Series([36, 71, 27, 62, 56, 58, 67, 11, 2, 24, 45, 30, 0,
-                                        9, 41, 28, 33, 19, 29, 43],
-                                        index=[10] * 20, name="val")),
-                    (10, 'b', pd.Series([78, 37, 23, 44, 6, 3, 21, 61, 39, 31, 53, 16, 66,
-                                         50, 40, 47, 7, 42, 38, 55],
-                                        index=[10] * 20, name="val")),
-                    (500, 'a', pd.Series([76, 72, 74, 75, 32, 64, 46, 35, 15, 70, 57, 65,
-                                          51, 26, 5, 25, 10, 69, 73, 77],
-                                         index=[500] * 20, name="val")),
-                    (500, 'b', pd.Series([8, 60, 12, 68, 22, 17, 18, 63, 49, 34, 20, 52,
-                                          48, 14, 79, 4, 1, 59, 54, 13],
-                                         index=[500] * 20, name="val"))]
-
-        self.assert_data_chunk_object_equal(result, expected)
diff --git a/tests/units/feature_selection/test_relevance.py b/tests/units/feature_selection/test_relevance.py
index a7a8b71..40ee2ee 100644
--- a/tests/units/feature_selection/test_relevance.py
+++ b/tests/units/feature_selection/test_relevance.py
@@ -83,8 +83,14 @@ class TestCalculateRelevanceTable:
 
         assert 0.5 == relevance_table.loc['feature_binary'].p_value
         assert 0.7 == relevance_table.loc['feature_real'].p_value
-        significance_test_feature_binary_mock.assert_called_once_with(X['feature_binary'], y=y_real)
-        significance_test_feature_real_mock.assert_called_once_with(X['feature_real'], y=y_real)
+
+        assert significance_test_feature_binary_mock.call_count == 1
+        pd.testing.assert_series_equal(significance_test_feature_binary_mock.call_args[0][0], X["feature_binary"])
+        pd.testing.assert_series_equal(significance_test_feature_binary_mock.call_args[1]["y"], y_real)
+
+        assert significance_test_feature_real_mock.call_count == 1
+        pd.testing.assert_series_equal(significance_test_feature_real_mock.call_args[0][0], X["feature_real"])
+        pd.testing.assert_series_equal(significance_test_feature_real_mock.call_args[1]["y"], y_real)
 
     @mock.patch('tsfresh.feature_selection.relevance.target_real_feature_real_test')
     @mock.patch('tsfresh.feature_selection.relevance.target_real_feature_binary_test')
diff --git a/tests/units/utilities/test_dataframe_functions.py b/tests/units/utilities/test_dataframe_functions.py
index 39c0b75..77b9279 100644
--- a/tests/units/utilities/test_dataframe_functions.py
+++ b/tests/units/utilities/test_dataframe_functions.py
@@ -4,210 +4,11 @@
 import warnings
 from unittest import TestCase
 
-import pandas as pd
-from tsfresh.utilities import dataframe_functions
 import numpy as np
+import pandas as pd
 from pandas.testing import assert_frame_equal, assert_series_equal
 
-
-class NormalizeTestCase(TestCase):
-    def test_with_dictionaries_one_row(self):
-        test_df = pd.DataFrame([{"value": 1, "id": "id_1"}])
-        test_dict = {"a": test_df, "b": test_df}
-
-        # A kind is not allowed with dicts
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_dict,
-                          "id", None, "a kind", None)
-
-        # The value must be present
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_dict,
-                          "id", None, None, "something other")
-
-    def test_with_dictionaries_two_rows(self):
-        test_df = pd.DataFrame([{"value": 2, "sort": 2, "id": "id_1"},
-                                {"value": 1, "sort": 1, "id": "id_1"}])
-        test_dict = {"a": test_df, "b": test_df}
-
-        # If there are more than one column, the algorithm can not choose the correct column
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_dict,
-                          "id", None, None, None)
-
-        # Sorting should work
-        result_df, column_id, column_kind, column_value = \
-            dataframe_functions._normalize_input_to_internal_representation(test_dict, "id", "sort", None, "value")
-        self.assertEqual(column_value, "value")
-        self.assertEqual(column_id, "id")
-
-        # Assert sorted and without sort column
-        self.assertEqual(result_df[result_df[column_kind] == "a"].iloc[0].to_dict(),
-                         {"_variables": "a", "value": 1, "id": "id_1"})
-        self.assertEqual(result_df[result_df[column_kind] == "a"].iloc[1].to_dict(),
-                         {"_variables": "a", "value": 2, "id": "id_1"})
-
-    def test_with_dictionaries_two_rows_sorted(self):
-        test_df = pd.DataFrame([{"value": 2, "id": "id_1"},
-                                {"value": 1, "id": "id_1"}])
-        test_dict = {"a": test_df, "b": test_df}
-
-        # Pass the id
-        result_df, column_id, column_kind, column_value = \
-            dataframe_functions._normalize_input_to_internal_representation(test_dict, "id", None, None, "value")
-        self.assertEqual(column_value, "value")
-        self.assertEqual(column_id, "id")
-
-        self.assertEqual(result_df[result_df[column_kind] == "a"].iloc[0].to_dict(),
-                         {"_variables": "a", "value": 2, "id": "id_1"})
-
-    def test_with_df_1(self):
-        # give everyting
-        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": 1}])
-        result_df, column_id, column_kind, column_value = \
-            dataframe_functions._normalize_input_to_internal_representation(test_df, "id", "sort", "kind", "value")
-
-        self.assertEqual(column_id, "id")
-        self.assertEqual(column_value, "value")
-        self.assertEqual(column_kind, "kind")
-        self.assertIn("a", set(result_df[column_kind]))
-        self.assertCountEqual(list(result_df.columns), ["id", "value", "kind"])
-        self.assertEqual(list(result_df[result_df[column_kind] == "a"]["value"]), [3])
-        self.assertEqual(list(result_df[result_df[column_kind] == "a"]["id"]), [0])
-
-    def test_with_df_2(self):
-        # give no kind
-        test_df = pd.DataFrame([{"id": 0, "value": 3, "sort": 1}])
-        result_df, column_id, column_kind, column_value = \
-            dataframe_functions._normalize_input_to_internal_representation(test_df, "id", "sort", None, "value")
-
-        self.assertEqual(column_id, "id")
-        self.assertEqual(column_value, "value")
-        self.assertEqual(column_kind, "_variables")
-        self.assertIn("value", set(result_df[column_kind]))
-        self.assertCountEqual(list(result_df.columns), ["id", "value", "_variables"])
-        self.assertEqual(list(result_df[result_df[column_kind] == "value"]["value"]), [3])
-        self.assertEqual(list(result_df[result_df[column_kind] == "value"]["id"]), [0])
-
-    def test_with_df_3(self):
-        # Let the function find the values
-        test_df = pd.DataFrame([{"id": 0, "a": 3, "b": 5, "sort": 1}])
-        result_df, column_id, column_kind, column_value = \
-            dataframe_functions._normalize_input_to_internal_representation(test_df, "id", "sort", None, None)
-
-        self.assertEqual(column_id, "id")
-        self.assertEqual(column_value, "_values")
-        self.assertEqual(column_kind, "_variables")
-        self.assertIn("a", set(result_df[column_kind]))
-        self.assertIn("b", set(result_df[column_kind]))
-        self.assertCountEqual(list(result_df.columns), ["_values", "_variables", "id"])
-        self.assertEqual(list(result_df[result_df[column_kind] == "a"]["_values"]), [3])
-        self.assertEqual(list(result_df[result_df[column_kind] == "a"]["id"]), [0])
-        self.assertEqual(list(result_df[result_df[column_kind] == "b"]["_values"]), [5])
-        self.assertEqual(list(result_df[result_df[column_kind] == "b"]["id"]), [0])
-
-    def test_with_wrong_input(self):
-        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": np.NaN}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", "sort", "kind", "value")
-
-        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": 1}])
-        self.assertRaises(AttributeError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "strange_id", "sort", "kind", "value")
-
-        test_df = pd.DataFrame([{"id": 0, "kind": "a", "value": 3, "sort": 1}])
-        self.assertRaises(AttributeError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", "sort", "strange_kind", "value")
-
-        test_df = pd.DataFrame([{"id": np.NaN, "kind": "a", "value": 3, "sort": 1}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", "sort", "kind", "value")
-
-        test_df = pd.DataFrame([{"id": 0, "kind": np.NaN, "value": 3, "sort": 1}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", "sort", "kind", "value")
-
-        test_df = pd.DataFrame([{"id": 2}, {"id": 1}])
-        test_dict = {"a": test_df, "b": test_df}
-
-        # If there are more than one column, the algorithm can not choose the correct column
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_dict,
-                          "id", None, None, None)
-
-        test_dict = {"a": pd.DataFrame([{"id": 2, "value_a": 3}, {"id": 1, "value_a": 4}]),
-                     "b": pd.DataFrame([{"id": 2}, {"id": 1}])}
-
-        # If there are more than one column, the algorithm can not choose the correct column
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_dict,
-                          "id", None, None, None)
-
-        test_df = pd.DataFrame([{"id": 0, "value": np.NaN}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", None, None, "value")
-
-        test_df = pd.DataFrame([{"id": 0, "value": np.NaN}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          None, None, None, "value")
-
-        test_df = pd.DataFrame([{"id": 0, "a_": 3, "b": 5, "sort": 1}])
-        self.assertRaises(
-            ValueError,
-            dataframe_functions._normalize_input_to_internal_representation,
-            test_df,
-            "id",
-            "sort",
-            None,
-            None)
-
-        test_df = pd.DataFrame([{"id": 0, "a__c": 3, "b": 5, "sort": 1}])
-        self.assertRaises(
-            ValueError,
-            dataframe_functions._normalize_input_to_internal_representation,
-            test_df,
-            "id",
-            "sort",
-            None,
-            None)
-
-        test_df = pd.DataFrame([{"id": 0}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", None, None, None)
-
-        test_df = pd.DataFrame([{"id": 0, "sort": 0}])
-        self.assertRaises(ValueError, dataframe_functions._normalize_input_to_internal_representation, test_df,
-                          "id", "sort", None, None)
-
-    def test_wide_dataframe_order_preserved_with_sort_column(self):
-        """ verifies that the order of the sort column from a wide time series container is preserved
-        """
-
-        test_df = pd.DataFrame({'id': ["a", "a", "b"],
-                                'v1': [3, 2, 1],
-                                'v2': [13, 12, 11],
-                                'sort': [103, 102, 101]})
-
-        melt_df, _, _, _ = \
-            dataframe_functions._normalize_input_to_internal_representation(
-                test_df, column_id="id", column_sort="sort", column_kind=None, column_value=None)
-
-        assert (test_df.sort_values("sort").query("id=='a'")["v1"].values ==
-                melt_df.query("id=='a'").query("_variables=='v1'")["_values"].values).all()
-        assert (test_df.sort_values("sort").query("id=='a'")["v2"].values ==
-                melt_df.query("id=='a'").query("_variables=='v2'")["_values"].values).all()
-
-    def test_wide_dataframe_order_preserved(self):
-        """ verifies that the order of the time series inside a wide time series container are preserved
-        (columns_sort=None)
-        """
-        test_df = pd.DataFrame({'id': ["a", "a", "a", "b"],
-                                'v1': [4, 3, 2, 1],
-                                'v2': [14, 13, 12, 11]})
-
-        melt_df, _, _, _ = \
-            dataframe_functions._normalize_input_to_internal_representation(
-                test_df, column_id="id", column_sort=None, column_kind=None, column_value=None)
-
-        assert (test_df.query("id=='a'")["v1"].values ==
-                melt_df.query("id=='a'").query("_variables=='v1'")["_values"].values).all()
-        assert (test_df.query("id=='a'")["v2"].values ==
-                melt_df.query("id=='a'").query("_variables=='v2'")["_values"].values).all()
+from tsfresh.utilities import dataframe_functions
 
 
 class RollingTestCase(TestCase):
diff --git a/tsfresh/feature_selection/significance_tests.py b/tsfresh/feature_selection/significance_tests.py
index fe909a5..8dfb646 100644
--- a/tsfresh/feature_selection/significance_tests.py
+++ b/tsfresh/feature_selection/significance_tests.py
@@ -203,6 +203,8 @@ def __check_if_pandas_series(x, y):
         raise TypeError("x should be a pandas Series")
     if not isinstance(y, pd.Series):
         raise TypeError("y should be a pandas Series")
+    if not list(y.index) == list(x.index):
+        raise ValueError("X and y need to have the same index!")
 
 
 def __check_for_binary_target(y):