Fix bindings helper function (#712)

* Fix bindings helper function After the pivoting was speed up in #705, the helper function for the spark and dask bindings also needs to change. Additionally, this allows to introduce a fix for #709. * Changelog
blue-yonder · Jun 12, 2020 · f1fcdd2 · f1fcdd2
1 parent 29cdae3
commit f1fcdd2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -19,6 +19,7 @@ Unreleased
     - Speed up the result pivoting (#705)
 - Bugfixes:
     - 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
@@ -24,10 +24,10 @@ def _feature_extraction_on_chunk_helper(df, column_id, column_kind,
     chunk = df[column_id].iloc[0], df[column_kind].iloc[0], df.sort_values(column_sort)[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,