Repository: dask/dask
File: dask/dataframe/io/parquet/arrow.py (Python)

Code before:

    md_list = []
    for keys, subgroup in data_df.groupby(partition_keys):
        if not isinstance(keys, tuple):
            keys = (keys,)
        subdir = "/".join(
            [
                "{colname}={value}".format(colname=name, value=val)
                for name, val in zip(partition_cols, keys)
            ]
        )
        subtable = pa.Table.from_pandas(
            subgroup, preserve_index=False, schema=subschema, safe=False
        )
        prefix = "/".join([root_path, subdir])
        pq._mkdir_if_not_exists(fs, prefix)
        outfile = guid() + ".parquet"
        full_path = "/".join([prefix, outfile])
        with fs.open(full_path, "wb") as f:
            pq.write_table(subtable, f, metadata_collector=md_list, **kwargs)
        md_list[-1].set_file_path("/".join([subdir, outfile]))

    return md_list


class ArrowEngine(Engine):
    @staticmethod
    def read_metadata(
        fs,
        paths,
        categories=None,
        index=None,
        gather_statistics=None,
        filters=None,
        split_row_groups=True,
        **kwargs,
    ):
        # Define the dataset object to use for metadata,
        # Also, initialize `parts`.  If `parts` is populated here,
        # then each part will correspond to a file.  Otherwise, each part will
        # correspond to a row group (populated below)
        parts, dataset = _determine_dataset_parts(
            fs, paths, gather_statistics, filters, kwargs.get("dataset", {})
        )
        col_chunk_paths = False
        if dataset.metadata:
            col_chunk_paths = all(
                dataset.metadata.row_group(i).column(0).file_path is not None
                for i in range(dataset.metadata.num_row_groups)
            )


Review comment (suggestion):
I don't know to what extent you want to keep the copy of the code from pyarrow as a pure copy, but, if including it here, there are some potential efficiency improvements to be made.

For example, just before calling this method, there is `t = pa.Table.from_pandas(df, preserve_index=preserve_index, schema=schema)` that converts the pandas DataFrame to an arrow table. You can also pass the dataframe directly, avoiding the conversion from arrow table back to a pandas dataframe here.

Code after:
                )
            else:
                dataset = pq.ParquetDataset(
                    allpaths[0], filesystem=fs, **dataset_kwargs
                )
            parts = [base + fs.sep + fn for fn in fns]
    else:
        # There is only one file to read
        dataset = pq.ParquetDataset(paths, filesystem=fs, **dataset_kwargs)
    return parts, dataset


def _write_partitioned(
    table, root_path, partition_cols, fs, preserve_index=True, **kwargs
):
    """ Write table to a partitioned dataset with pyarrow.

        Logic copied from pyarrow.parquet.
        (arrow/python/pyarrow/parquet.py::write_to_dataset)

        TODO: Remove this in favor of pyarrow's `write_to_dataset`
              once ARROW-8244 is addressed.
    """
    fs.mkdirs(root_path, exist_ok=True)

    df = table.to_pandas(ignore_metadata=True)
    partition_keys = [df[col] for col in partition_cols]
    data_df = df.drop(partition_cols, axis="columns")
    data_cols = df.columns.drop(partition_cols)
    if len(data_cols) == 0 and not preserve_index:
        raise ValueError("No data left to save outside partition columns")

    subschema = table.schema
    for col in table.schema.names:
        if col in partition_cols:
            subschema = subschema.remove(subschema.get_field_index(col))

    md_list = []
    for keys, subgroup in data_df.groupby(partition_keys):
        if not isinstance(keys, tuple):
            keys = (keys,)
        subdir = fs.sep.join(
            [
                "{colname}={value}".format(colname=name, value=val)
                for name, val in zip(partition_cols, keys)
            ]
        )
        subtable = pa.Table.from_pandas(
            subgroup, preserve_index=False, schema=subschema, safe=False
        )
        prefix = fs.sep.join([root_path, subdir])