Fix scalability limitations of current implementation

[andygrove]

Query execution works by building up a tree of futures to execute each partition in each query stage.

The root node of each query stage is a RayShuffleWriterExec. It works by executing its child plan and fetching all of the results into memory (this is already not scalable because there could be millions or billions of rows). It then concatenates all of the batches into one large batch, which is returned. This large batch is then stored in Ray's object store and will be fetched by the next query stage.

The original disk-based shuffle mechanism (that was removed in #19) did not suffer from any of these issues because query results were streamed to disk in the writer and then streamed back out in the reader. However, this approach assumes that all workers have access to the same local file system.

[andygrove]

One approach would be to bring back the original shuffle code and then add a mechanism for reading data from another node by implementing a gRPC based service (such as Arrow Flight protocol) in each worker.

[robtandy]

Would it be possible to use a ray generator for the shuffle writer and reader execution plans?

Rather than serializing all of the data to disk in the writer, and reading it in the reader, we yield record batches in the Task. Then ray sends those in the in memory object store (spilling if required).

Then we can pipeline the execution of the writer and readers so that stages need not wait for completion of dependent stages to start processing.

Would this work?

Willing to work on this if you think its a valid approach.

[robtandy]

It would be nice if ray facilitated streaming via network between Actors in addition to exchange of messages via the object store. If its only to exchange streams of batches would Arrow Flight be necessary or could something simpler be used, assuming the Actors exist only ephemerally.

EDIT: Is the serialization of Arrow to/from the object store zero copy? If so, then that might be good enough to start. That wasn't clear to me when reading the ray docs which mention zero copy for numpy data structures.

EDIT 2: It seems that the pickle protocol used has out of band support for zero copy serialization for arrow: https://peps.python.org/pep-0574/

[robtandy]

I was able to refactor the code such that batches are exchanged in a streaming fashion through the ray object store.

Its too slow though and struggles to scale when run on larger TPCH scale factors. The heavy interaction between all stages and exchanging object refs pointing to batches is too much for Ray. I would observe a typical TPCH query creating 10000+ ray tasks during its execution and the bottle neck appeared to the be the object store.

I'll refactor again trying to exchange data with Arrow Flight as you have suggested @andygrove

[andygrove]

Thanks for testing this out @robtandy!

We know from experience in the Ballista and Comet projects that an Arrow Flight approach should work well, so let's hope it isn't too challenging to implement this here.

[andygrove]

We know from experience in the Ballista and Comet projects that an Arrow Flight approach should work well, so let's hope it isn't too challenging to implement this here.

Minor correction: Comet does not use Flight, but the approach is similar. We stream batches to disk using Arrow IPC, and then the shuffle files (actually blocks from the files) are transferred to the reader.

[robtandy]

The current implementation i have streams object references out of each partition through an intermediate Actor to hold the queues. Then subsequent stages connect to the intermediate actor and read streaming object references from the queue. They then have to ray.get() on the references to materialize the IPC encoded batch.

I was really hoping for more performance and scale out of this approach, and I particularly like that it reuses the RepartitionExec's as implemented in DataFusion. But alas.

My work isn't documented well yet but if you have time to discuss it sync, I'd appreciate your eyes on it and how we can improve it. Let me know if you have time to pair program for a bit so I can demo.

[edmondop]

Interested to join too

[robtandy]

Great! Im in eastern time zone US. @edmondop ? @andygrove ? This will be useful as I want to discuss in particular the approach to reusing the RepartitionExec's that I took and get feedback.

[edmondop]

Great! Im in eastern time zone US. @edmondop ? @andygrove ? This will be useful as I want to discuss in particular the approach to reusing the RepartitionExec's that I took and get feedback.

Can you join the Discord channel where we can sync more easierly maybe?

[robtandy]

Yep. Send me a link at rob.tandy@gmail.com

[edmondop]

rob.tandy@gmail.com

You can find all the links to the channels here https://datafusion.apache.org/contributor-guide/communication.html. We are in #datafusion-ray on Discord

[robtandy]

Joined!