During write_parquet if Ray worker crashes invalid partial partition may be persisted

[jpedrick-numeus]

Describe the bug

I'm running Daft + Ray on an instance with limited memory. When calling daft.DataFrame.write_parquet if workers crash for any reason, the partially written parquet files aren't cleaned up which causes problems in downstream processes. I, at minimum, need the entire write_parquet operation to return some kind of failed state. Ideally, however, would be that Daft would cleanup any potentially corrupted parquet files and retry the task.

To Reproduce

Run daft.DataFrame.write_parquet with too many threads for the memory available resulting in OOM crashes for Ray workers.

Expected behavior

Either being informed that the write operation is potentially corrupted from the return value(or thrown exception) or ensure that the write_parquet operation cleans up and potentially corrupted parts when a ray worker crashes.

Component(s)

Parquet

Additional context

No response

[andrewgazelka]

Thanks for reporting this. Want to take a look @jaychia, @colin-ho ?

[colin-ho]

Hey @jpedrick-numeus, this should be possible to implement. We can implement a cleanup mechanism to detect duplicate/corrupted files written by daft upon a successful write operation.

Setting write_mode='overwrite' already does this, of course, it will also clean up all files that previously existed in the directory. We can implement a similar logic for write_mode='append' that only cleans up files written by daft.

[jpedrick-numeus]

Hi @colin-ho , thanks for the reply. I'm not quite sure what you're saying here. If I say "write_mode='overwrite'" Daft will cleanup partial writes from a crashed ray worker?

My intuition is that currently DataFrame.write_parquet() call would still have the corrupted parquet files, and I would need to call write_parquet(write_mode='overwrite') to clean them up?

[colin-ho]

write_mode='overwrite' is designed to overwrite the directory, i.e. any existing files before the write will be removed.

The way write_mode='overwrite' works is:

• Daft writes the new parquet files directly to the target directory
• Once the files are successfully written, it removes all other files.
  This cleanup ensures only the successfully written files remain in the directory, which means any corrupted parquet files will also be removed.

You don't need to call df.write_parquet and then df.write_parquet(write_mode='overwrite'), just a single df.write_parquet(write_mode='overwrite') will work.

[jpedrick-numeus]

Ok, interesting, I'll try that then.