[Feature] Row streaming from a DataFrame

[jaychia]

Summary

Retrieval of rows to the driver from a DataFrame through a streaming interface

Please upvote this discussion and share your use-cases if this is a feature you would like to see implemented!

Use-cases

1. Retrieve rows to pipe to some proprietary/currently unsupported storage backend (e.g. dynamodb)
2. Retrieve rows for ML model training

Proposal

row_iterator = df.iterrows()
colnames = row_iterator.column_names()

for row in row_iterator:
    row: tuple[PyObj]
    ...

Extensions to this API should be made for:

1. Retrieval of rows in batches
2. Retrieval of rows based on some offset
3. [Stretch] Retrieval of rows by distributed clients that poll rows from the same DataFrame

[sandias42]

I've been looking for something like this and would be very excited if it existed. I have two usecases:

1. Feeding data to an ML model in batches (must be batches of rows with the same batch_size, eventually yielding the entire dataset without resampling any rows)
2. Submitting batches of data to an API (an embedding endpoint, in this case) which has specifications for the maximum number of rows per batch (and where getting exactly that number saves me money and time).

For 2) having some notion of order preserved would be nice but not required.

As far as I can tell the best current option for this is converting to a ray dataset and using get_batch or iter_batches?

It's also possible my use case would be achieved by a window function implementation (which I saw yall working on). It would also be possible to achieve my objectives if I could repartition with an exact number of rows per partition (or even partition boundaries as dask allows) plus a map_partition (I believe your udfs are effectively doing a map_partition even though the docs don't make promises that each time the function is called its with the data from one partition), or if I could assign an incrementing global index (which I know is expensive but would make it possible). I know you are trying to avoid some of these things.

Really excited about the project, by the way. Just moved from dask as you can probably tell and its quite refreshing to have a library which works as advertised.

[jaychia]

Really excited about the project, by the way. Just moved from dask as you can probably tell and its quite refreshing to have a library which works as advertised.

Welcome to Daft @sandias42 and thank you for the kind words!

I think we have to get some of our documentation in order to reflect new functionality that has been added since this discussion was put up, and probably write some tutorials after really ironing out the user-story, but here are our current recommendations.

• "I believe your udfs are effectively doing a map_partition" - correct 😛. In fact, you can also do your own naive batching inside of your UDF to achieve (1) "Feeding data to an ML model in batches" for a parallel model batch inference use-case and (2) "Submitting batches of data to an API".
• For ML model training, there actually is a way to run the dataframe and then retrieve data by partition: df.iter_partitions(). I realize now (with great dismay 😭) that this isn't on our documentation............. I am also working on a way to having the rate of consumption of this iterator provide backpressure to the rate of data production. When this is done, this will let you iterate on the dataframe on a per-partition basis (with backpressure to avoid the dataframe just running to full completion!) using: for partition in df.iter_partitions(buffer_size=1): ....

For the first point, I wonder if we should just expose a @daft.udf(max_batch_size=8) kwarg that will then feed data of a certain batch size (instead of the full partition) to the UDF, since this is a pretty common use-case. Let me know your thoughts here!

Also to answer your other questions:

1. We also support df.to_ray_dataset() for cases where you are fine just materializing your dataset as a Ray dataset and then continuing with the rest of your work from there. This works great if you are just using the Ray ecosystem (e.g. Ray train/Ray tune) for your workload, because they all work with Ray Datasets.
2. It's also possible my use case would be achieved by a window function implementation - similar, but I think this use-case would only require non-overlapping (or "tumbling") windows, and not require windows that end with the current partition (do not "bleed over" into the next partition). This simplifies the problem quite a bit in a distributed context.
3. It would also be possible to achieve my objectives if I could repartition with an exact number of rows per partition - yes, but we're actually trying to move towards a world where we increasingly try to abstract away details about partitioning from our users. More details to come soon 🙂
4. Check out also df.into_partitions(N) API for cheaply splitting/coalescing partitions to tune the number of partitions (and the size of those partitions).

cc @clarkzinzow @samster25 too

[sandias42]

Hey @jaychia thanks so much for the thoughtful and speedy reply!

The challenge I see with using a UDF / map_partition in my use case is handling the "ragged" last chunk per partition. For both my applications, I really need to have exactly the desired number of rows per batch. With a single smaller "ragged" batch out of a very large dataset, which in some cases I'd just drop entirely.

It's not a big deal to have one funky batch, but in my large dataset I have thousands of partitions. Because I can't control the size of a partition, each call of my "batchify" UDF on a partition will produce an incorrectly sized batch on the ragged end, and thousands of incorrectly sized batches is a serious issue for both cost and correctness. In this sense my use case does "bleed over" across partitions, because when partition_length % batch_size != 0 (which will be true for almost all partitions I can't control the size of) I need to look at the next partition to collect the desired number of rows.

Re the other options you mention

• iter_partitions has the same problem as UDF- if I can't control partition size, then the "batches" will be variable size, which is unacceptable in my use case. (btw you mention ML training- I think I'm not unusual in that my application cannot tolerate variable sized batches? Like, tensor cores really want matmuls on matrices of multiples of 8 for performance and variable sized batches violate the statistical assumptions of standard optimizers like Adam)
• df.to_ray_dataset Oof, I didn't realize this fully materialized. I guess at this point I might as well write to s3 and switch over to ray dataset entirely so I don't need to have cluster memory pressure from caching the whole thing.
• @daft.udf(max_batch_size=8) I don't think max_batch_size would solve my problem unless you were promising that all batches were going to be max_batch_size except the last one in the entire dataset (NOT the last in the partition), ideally with a drop_last behavior as the deep learning Dataset batch iterators typically provide.
• df.into_partitions(N) this sounds useful in general, and might help in cases where the partitions were too small relative to my batchsize or something, but without promises about what partition sizes in num_rows will result (which I understand you want to abstract away) won't help.

FWIW I really like the idea of daft abstracting away partitioning. But that aim feels in conflict with the interface for e.g. UDFs, iter_partitions, which force me to think about partitioning by handing me Series of mysterious length.

[jaychia]

This all makes sense! Thanks for taking the time to elaborate 😍

If you need perfectly batch-aligned partitions (except the last one), I can't think of a great solution that provides full parallelism over all partitions without some extra work. There are a few solutions I can think of, but none of them are fully satisfactory:

• 1-pass serial algorithm: Proceed serially from the first row and run in batches until the last row, processing batches across partition lines when necessary. Unfortunately we won't get any parallelism here...
• 1-pass algorithm by coalescing partitions together (df.into_partitions(N)): This is the "good enough" approach, where instead of creating O(thousand) of partitions, you can first perform a coalesce and reduce that to O(hundred) and then run in batches across each partition, reducing the amount of raggedness.
• 2-pass algorithm with materialization: Materialize the entire dataset, perform a right-sizing operation across the entire dataframe, then run in parallel across the newly right-sized partitions.

Is the last one (2-pass algorithm with materialization) something you're looking for? It's expensive, but I think if we want perfect boundaries aligned to your batch size it might be the only solution.

[sandias42]

Thanks! To make sure I understand what you mean, the lack of parallelism you are pointing out is that the 1-pass algorithms won't be able to yield multiple batches to some subsequent consumer at the same time?

If so, that does seem like a substantial downside. The 2-pass algorithm sounds preferable despite the materialization costs.

That said, it seems like there should be a parallel implementation possible here?

Smth like: each of num_workers parallel workers fetches a partition from the partition queue, and yields batches of the correct size until len(remaining_partition) < batch_size at which point it fetches another partition. If the queue is exhausted when a new partition is needed to complete the batch, the ragged remaining_partition is sent to (the same) single worker to be concatenated and yielded. (a more sophisticated implementation could send the ragged edges just to neighboring workers instead of a single central worker)

In this implementation you will only have at most num_workers ragged edges to be coalesced by the final node, and each ragged edge will be < batch_size which doesn't seem so bad in terms of data transfer and memory. Also note that the final node can yield batches without waiting for everything to coalesce- as long as the length of all the ragged edges concatenated so far is >= batch size it can yield a batch.

Thats more complex then the 2 pass algorithm you mentioned. I'd be happy with the 2-pass version over nothing for sure!