Fix incorrect update_metadata_with_features calls in ArrowDataset (hu…

…ggingface#2258)
shirte · Apr 26, 2021 · 8e903b5 · 8e903b5
1 parent 5adf06a
commit 8e903b5
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -838,7 +838,7 @@ def flatten(self, new_fingerprint, max_depth=16) -> "Dataset":
             else:
                 break
         dataset.info.features = Features.from_arrow_schema(dataset._data.schema)
-        self._data = update_metadata_with_features(self._data, self.features)
+        dataset._data = update_metadata_with_features(dataset._data, dataset.features)
         logger.info(
             "Flattened dataset from depth {} to depth {}.".format(depth, 1 if depth + 1 < max_depth else "unknown")
         )
@@ -1018,7 +1018,7 @@ def remove_columns(self, column_names: Union[str, List[str]], new_fingerprint) -
             del dataset._info.features[column_name]
 
         dataset._data = dataset._data.drop(column_names)
-        dataset._data = update_metadata_with_features(dataset._data, self.features)
+        dataset._data = update_metadata_with_features(dataset._data, dataset.features)
         dataset._fingerprint = new_fingerprint
         return dataset
 
@@ -1109,7 +1109,7 @@ def rename(columns):
         )
 
         dataset._data = dataset._data.rename_columns(new_column_names)
-        dataset._data = update_metadata_with_features(dataset._data, self.features)
+        dataset._data = update_metadata_with_features(dataset._data, dataset.features)
         dataset._fingerprint = new_fingerprint
         return dataset