I'm trying to partition my df into a few group of ...
# general
k
I'm trying to partition my df into a few group of parquet files based on their size and run an operation on each group individually. For example, from a dataset of 1TB with original parquets of varying sizes, I want to run an operation which collects the results from groups of 10GB each, and write out 100 results into 100 files. Which way would be best? (1) divide the original df using into_partitions, then run a udf on each micropartition from iter_partitions (2) glob and do a rolling sum on 'size' column then divide by partition size and do a groupby and agg the list of parquet files and read_parquet on the list of paths then run the operation on each row (3) glob and do a rolling sum on 'size' column then divide by partition size and run a for loop to filter on the partition number Or any better ideas? I'm thinking this could have been easier if there's a wrapper for iter_partitions which gives an iterator of DFs instead of micropartitions?
j
Hmm, would you perhaps be able to do this with:
Copy code
df = daft.from_glob_path(...)
df = df.where("size < 10000000000")
filepaths = df.to_pydict()["paths"]

df_files_under_10gb = daft.read_parquet(filepaths)
df_files_under_10gb = df_files_under_10gb.into_partitions(100)
df.write_parquet(...)
I also might be misunderstanding your ask
k
I'm trying to coalesce partitions not by number but up to a certain threshold size and then run a UDF on them
j
I see… Yeah that’s tricky right now since in order to know the size of the data we have to materialize the entire data first which is quite expensive. What is the overall goal here? Are you trying to get better-sized files?
k
Yes trying to get evenly sized files which can be loaded in as even sized batches
j
If you need really really evenly sized files, you might need to do a full shuffle with something like
.repartition()
but note that this is quite expensive and requires a full shuffle. Otherwise,
.into_partitions()
is your best bet. You can control the file sizes by using
daft.set_execution_config(_*parquet_target_filesize*_=…)
z
Self proclaimed daft noob, but hoping this helps… Think the technical term for what you’re describing is binpacking https://en.m.wikipedia.org/wiki/Bin_packing_problem There’s a binpacking library here https://pypi.org/project/binpacking/ Could do something like… • list the files and their sizes • Binpack files, assigning a bin_id to each file • Group by bin_id • map_groups, running your UDF on each bin This is pretty much your option #2 Think it would look something like…
Copy code
import daft
import binpacking

# Read files and their sizes
df = daft.from_glob_path(…)
files = df.to_pylist()
uri_to_size = {f[“uri”]: f[“size”] for f in files}

# Binpack
max_bin_size = # Specify max bin size 
bins = binpacking.to_constant_volume(uri_to_size, max_bin_size)
binned_files = [
    {
        “bin_id”: bin_id,
        “uri”: uri,
        “size”: size
    }
    for bin_id, bin in enumerate(bins)
    for uri, size in bin.items()
]

# Group by bin, map groups
@daft.udf(return_type=daft.DataTypes.string())
def my_udf(group):
    “””Aggregates the parquet files in the group, writing a new parquet file and returning its path”””
    # think you have a lot of options here for what you actually use to process the parquet files
    pass

df = daft.from_pylist(binned_files)
df = df.groupby(“bin_id”).map_groups(my_udf)
df.collect()
But I’m curious what do you actually want to do in your udf? In the example above, daft is basically just acting as a fancy ray/local dispatcher, depending on which runner you use. You could accomplish the same thing with a thread pool executor and the ray python client. It’d just be a little lower level. If you modeled the udf in more native daft expressions, think you wouldn’t have to worry about binpacking. Also curious if by even sized files, are you saying you want each group worker to have an even sized input? Or do you want all group workers to produce an even sized output?
k
Thank you! I'll have to try them out later on! I'm trying to do this because I want to have even sized files as outputs for another system to ingest and they have strict requirements on how they want their data to be formatted and how much data should be in each file.