rust engine consume a lot of memory compared to pyarrow

[djouallah]

Environment

Delta-rs version:
0.21.0

Binding:

Environment:

• OS:
  Linux

---

Bug

switching from pyarrow engine to rust increase memory usage by nearly 3X, the job used to works fine, but now, getting OOM errors.

I added a reproducible example with only 60 input files to demo the issue

https://colab.research.google.com/drive/1fahlV0FgKSAS8sQvRMu47s3bDP1ekLbb#scrollTo=333a177b-f075-412e-8ca1-32d44f8c07eb

[rtyler]

@djouallah 👋 in the attached notebook, which write_deltalake call is resulting in memory pressure? I'm less familiar with duckdb, but I assume the df objects that it is producing are pyarrow.DataSet? or are they another type?

[djouallah]

it is an Arrow RecordBatchReader I think

[ion-elgreco]

Rust engine materalizes everything to memory prior to starting the whole process. Pyarrow probably writes batch by batch when you pass a reader

[djouallah]

just for my own understanding, is this something that can be fixed by datafusion ?

[ion-elgreco]

@djouallah there is a PR to address this but the contributor didn't have time to finish it yet: #2289

[rtyler]

I have been analyzing the performance around this issue, even with today's latest HEAD of main

Rust engine

Pyarrow engine

The root cause of this behavior is, in my opinion, the collection of Vec<RecordBatch> by the Rust writer. I have been experimenting with some alternative approaches that allow the writer to use either an iterator or a stream of RecordBatch to perform incremental read/write behavior rather than full-read, then full-write behaviors as it is doing today.

There is some trickiness around Send +Sync here because of how the WriteBuilder works with its IntoFuture implementation.

Nothing required from anybody, I just wanted to share the analysis and progress here. 🤔 🏃

[djouallah]

using 0.22, I just use the pyarrow engine, yes, processing 2300 csv files and it works great !!!