I'm running a group by query on the 1.8TB dataset I mentioned above. The ray job is failing because some RaySwordfishActors are running out of memory, either during the repartition or write parquet step (it looks like they're interleaved, so it's hard to tell). The log I'm seeing is:

(raylet) node_manager.cc:3193: 1 Workers (tasks / actors) killed due to memory pressure (OOM)

I've tried various k8s pod sizing - currently I have 80 worker pods with 8cpu / 32gb. Does anyone know why so much memory would be used, or how to tune / think about queries like this?

Hmm this is not ideal, but I wonder if it'll help to do this in two stages: 1. Write out partitioned data 2. Read partitioned data with agg list So something along the lines of

import daft
from daft import col

df = daft.read_parquet(
    "<abfs://my_container@my_service_account.dfs.core.windows.net/in_table_path/>"
)

df = df.select(col("group_col"), col("other_col"))

df.write_parquet(
    "<abfs://my_container@my_service_account.dfs.core.windows.net/tmp_grouped/>",
    partition_cols=["group_col"],
)

# Stage 2

df = daft.read_parquet(
    "<abfs://my_container@my_service_account.dfs.core.windows.net/tmp_grouped/>",
    hive_partitioning=True,
)

df = df.groupby("group_col").agg_list(col("other_col"))

df.write_parquet(
    "<abfs://my_container@my_service_account.dfs.core.windows.net/out_table_path/>"
)

I can give that a try, but does that mean it's necessary to break up all queries needing a shuffle of this much data?

I think it's more the nature of an agg_list. You would need to keep the entire dataset in-memory/spill to disk until all rows have been shuffled

I do see that some spilling to disk is happening, since I see logs like

(raylet, ip=10.224.20.85) Spilled 8674 MiB, 866 objects, write throughput 1621 MiB/s

, but not enough to keep the actor from getting killed

That's really odd... just to confirm, we have permissions to create new folders in abfs?

yeah we do - it wrote out the tmp dir files fine without

partition_cols

Ah cool, awesome debugging! Yup so if we broke it down fully: • With

agg_list

a specific

group_col

needs to be fully aggregated before it can be flushed from memory. With the initial dataset, we don't know when we've seen all values belonging to a specific

group_col

until we've shuffled all the data • By adding a partitioning step, all values belonging to a specific

group_col

now belong on the same worker on the same partition. So when we finally do the

agg_list

, on each worker, at most one partition needs to be materialized in-memory before we know that the `agg_list`s for belonging to that partition are completed. And whenever a batch of `agg_list`s is completed, it goes on to the next operator which is a write so it's removed from memory • Also, your idea to do an in-memory repartitioning was a great choice! I should've led with that instead of suggesting to write out the data 😅

Thanks for breaking that down - it makes sense!