Skip to content

Commit

Permalink
Only return parquet metadata if intending to write
Browse files Browse the repository at this point in the history
  • Loading branch information
martindurant committed Oct 22, 2024
1 parent fcae1a2 commit 5aa1a26
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion src/dask_awkward/lib/io/parquet.py
Original file line number Diff line number Diff line change
Expand Up @@ -483,6 +483,7 @@ def __init__(
npartitions: int,
prefix: str | None = None,
storage_options: dict | None = None,
write_metadata: bool = False,
**kwargs: Any,
):
self.fs = fs
Expand All @@ -496,16 +497,19 @@ def __init__(
if isinstance(self.fs.protocol, str)
else self.fs.protocol[0]
)
self.write_metadata = write_metadata
self.kwargs = kwargs

def __call__(self, data, block_index):
filename = f"part{str(block_index[0]).zfill(self.zfill)}.parquet"
if self.prefix is not None:
filename = f"{self.prefix}-{filename}"
filename = self.fs.unstrip_protocol(f"{self.path}{self.fs.sep}{filename}")
return ak.to_parquet(
out = ak.to_parquet(
data, filename, **self.kwargs, storage_options=self.storage_options
)
if self.write_metadata:
return out


def to_parquet(
Expand Down Expand Up @@ -667,6 +671,7 @@ def to_parquet(
parquet_old_int96_timestamps=parquet_old_int96_timestamps,
parquet_compliant_nested=parquet_compliant_nested,
parquet_extra_options=parquet_extra_options,
write_metadata=write_metadata,
),
array,
BlockIndex((array.npartitions,)),
Expand Down

0 comments on commit 5aa1a26

Please sign in to comment.