Hey folks, I'd like to integrate daft into a ray ...
# general
n
Hey folks, I'd like to integrate daft into a ray distributed training pipeline. I have already used the following approaches for this: 1. daft.DataFrame.to_ray___dataset 2. daft.DataFrame.write_parquet -> Load directly using ray.data.read_parquet 3. daft.DataFrame.to_torch_map_dataset 4. daft.DataFrame.to_torch_iter_dataset In all cases, the API collects the dataframe before iterating through it. Instead, I'd like to deal with partitions iteratively in a ray dataset. Is there a way to create a ray DataSource? Or use iter_partitions() to get some of this functionality? Having seen the info on the

new streaming execution engine▾

, will this be coming to Ray some time soon?
c
Hi @Neil Wadhvana, you can try
iter_partitions
, which will yield an iterator of partitions in the form of Ray ObjectRefs. Regarding the streaming engine, it is on our roadmap to integrate it with Ray, but our priority is to release it for local mode first.
j
Let us know how
iter_partitions
performs for you? It should be very simple (and cheap) to convert these objectrefs into the blocks that Ray Dataset expects. Here’s the code in Daft that does it: https://github.com/Eventual-Inc/Daft/blob/main/daft/runners/ray_runner.py#L138-L172 If you figure it out, we should maybe upstream your changes to change how Daft creates Ray datasets: https://github.com/Eventual-Inc/Daft/blob/main/daft/runners/ray_runner.py#L279-L290 (or we can maybe upstream a contribution into Ray Data to have a Daft Datasource)
g
the Daft Datasource at Ray Data would be awesome... Ray docs recommend using Spark or Dask as last mile preprocessing before ML training, but doing that with Daft would be faster/cheaper/simpler/etc etc 😁
❤️ 1
n
Hey thread, putting this here for visibility: https://github.com/Eventual-Inc/Daft/issues/3374
I did look at
iter_partitions
, but this requires an executed dataframe.
t
Hi, we started using daft today to read from Iceberg tables and prep data for an ML pipeline hosted on Ray. Started with ~10M rows and ran into the same OOM problem with
to_ray_dataset()
. Fortunately,
write_parquet()
->
read_parquet()
worked for us. Looking forward to a better fix for this because we'll have some 500M row datasets to process soon. 🙏
j
Thanks — can you give us a better idea of: 1. The plan you are running (
df.explain(True)
) 2. What machines are you running 3. What parquet files are you reading (types of columns, size of files, where the files are stored)
t
1. It's very simple so I can explain the entire procedure here. a. Read 2 snapshots from an Iceberg table: current and previous. (see note in 3 below about partitions). b. Perform an outer join between the two to generate a new DF. c. Calculate a CDC column and add it to the DF from 2. d. Try to create a Ray dataset which hits OOM. 2. Single node Ray cluster and Ray on K8s (3-5 nodes), both with 96GB+. 3. Num of files varies wildly by snapshot. Older snapshots have only 100+ files. Newer snapshots have more due to some really annoying behavior of the Dataframe writer in Spark where we can sometimes see 5000+ files for the same table. To mitigate this, we load the Iceberg table with
into_partitions()
on the 2 original DFs to something like 100-200 otherwise the join itself takes 10-20m before running OOM (before reaching the call to generate a Ray dataset). Column count is ~20, nothing complex so far like dense vectors (all text or number data, some lists). Iceberg table is stored in S3 buckets. Another thing I noticed is that using the native runner on my M2 Pro MBP is 2x faster than trying to use the Ray runner on a 32 core 96GB single node ray cluster. Furthermore, I didn't need do the repartitioning workaround with the native runner; it materialized the outer join in ~3 minutes.
👀 1
The intermediate Parquet file we're generating as part of the workaround is ~10 GB.
g
out of curiosity, have you compacted the iceberg tables before reading them? 100+ files for 10GB worth of parquet files sounds too much... after running compaction you should get around 20 files (default size is 0.5 GB). 5000+ small files is going to be too much
t
Yeah I understand. We write the table at least once a day, but currently I was thinking about running compaction once a week. I'm still trying to understand the Spark-Iceberg behavior for why so many parquet files are getting created when the actual underlying data has much fewer partitions based on the specified partition keys.
g
yeah I hear you. spark is going to create as many files as tasks it uses to write. the available compute in the cluster, the data itself and spark's AQE configuration affect that... if you write +5000 files every write, every day, you should run compaction every day to keep your table good for readers beyond that... try reading with ray/daft before and after compaction, to test the "its too many small files" hypothesis 😉
👍 1
j
Actually, Daft does a pretty good job of “coalescing reads”. So if you have many small Parquet files it will actually coalesce them into a single bulk-read, based on some logic of estimating their size once read into memory! This helps alleviate the “small file problem”. But yes if you have hundreds of thousands/millions of tiny files that will still cause performance issues.
g
oh that's interesting!
j
WRT to the join OOMing, this is something that @Colin Ho is looking at. I’m guessing that your join (without into_partitions) is OOMing the head node? We’ve noticed Ray taking an absurd amount of memory on the head node when performing these operations on O(thousands) of partitions, and are building some workarounds for that right now with a target of a January release.
t
Daft does a pretty good job of “coalescing reads”.
Yes Daft seems to handle this OK, especially in native mode. It's when using the Ray runner and generating the Ray data set that we see the OOMs.
Anyway, I updated the table props to explicitly set
write.distribution-mode
to
hash
, did another table update, and it looks like the new snapshot's num partitions on the Daft DF is much closer to the num partitions of the Spark DF during the write.
I’m guessing that your join (without into_partitions) is OOMing the head node?
Yes we run OOM on just the Daft DF outer join, but there's a separate issue which I think is what @Neil Wadhvana called out originally, that the
to_ray_dataset()
runs OOM. In that case, it's all nodes and not just the head.
j
Gotcha. Head node OOM during joins with many partitions is something we’re aware of and @Colin Ho has some new join algorithms that can help there.
to_ray_dataset()
OOM is likely because of the large materialization into Ray. We should chat more about enabling a better workflow here, would love to chat over a call to hash it out.
👍🏽 1