use head instead of return_schema

dask · Nov 14, 2020 · 44f3a04 · 44f3a04
1 parent e57faa8
commit 44f3a04
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
@@ -977,7 +977,7 @@ def write_partition(
         compression=None,
         index_cols=None,
         schema=None,
-        return_schema=False,
+        head=False,
         **kwargs,
     ):
         _meta = None
@@ -1020,10 +1020,11 @@ def write_partition(
                 _meta.set_file_path(filename)
         # Return the schema needed to write the metadata
         if return_metadata:
-            if return_schema:
-                return [{"schema": t.schema, "meta": _meta}]
-            else:
-                return [{"meta": _meta}]
+            d = {"meta": _meta}
+            if head:
+                # Only return schema if this is the "head" partition
+                d["schema"] = t.schema
+            return [d]
         else:
             return []
 

diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
@@ -583,9 +583,7 @@ def to_parquet(
                 partition_on,
                 write_metadata_file,
             ],
-            toolz.merge(kwargs_pass, {"return_schema": True})
-            if d == 0
-            else kwargs_pass,
+            toolz.merge(kwargs_pass, {"head": True}) if d == 0 else kwargs_pass,
         )
         part_tasks.append((name, d))