fix export error when export_stats columns is null

Cathy0908 · Jan 20, 2025 · 7401646 · 7401646
1 parent 030e786
commit 7401646
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/data_juicer/core/exporter.py b/data_juicer/core/exporter.py
@@ -114,12 +114,13 @@ def _export_impl(self, dataset, export_path, suffix, export_stats=True):
                 export_columns.append(Fields.stats)
             if Fields.meta in dataset.features:
                 export_columns.append(Fields.meta)
-            ds_stats = dataset.select_columns(export_columns)
-            stats_file = export_path.replace('.' + suffix, '_stats.jsonl')
-            Exporter.to_jsonl(
-                ds_stats,
-                stats_file,
-                num_proc=self.num_proc if self.export_in_parallel else 1)
+            if len(export_columns):
+                ds_stats = dataset.select_columns(export_columns)
+                stats_file = export_path.replace('.' + suffix, '_stats.jsonl')
+                Exporter.to_jsonl(
+                    ds_stats,
+                    stats_file,
+                    num_proc=self.num_proc if self.export_in_parallel else 1)
 
         if self.export_ds:
             # fetch the corresponding export method according to the suffix