`BatchCoalescer` but without automatic batching

[corasaurus-hex]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

I am creating new record batches from a stream of record batches. I need to ensure that identical values in a sorted column are always located within the same record batch. BatchCoalescer isn't a good fit for this but I would like to take advantage of the optimized machinery within BatchCoalescer to accomplish this -- except nothing outside of that crate has access to that machinery.

Describe the solution you'd like

I would like something akin to BatchCoalescer but that gives me the control to decide when to batch buffered records into a batch (and maybe how many of the buffered records should be in the batch?). It must allow me to not know what the largest record batch size will be up front.

Implementation-wise, @timsaucer suggested it could be expressed in terms of a BatchCoalescer<T> that gives you control of batching by you implementing a trait and passing that a value of that trait into the coalescer.

I was thinking that it could be something like a lower-level vec + batch abstraction, where you can add new batches to it and it will re-allocate the internal arrays when the capacity is exhausted or even chain arrays together as more capacity is needed. This can have a function to emit a record batch based on a slice of that data... although maybe that implies the need for something like a ring buffer that can have an offset at the start of the arrays and a capacity that can be grown.

Describe alternatives you've considered

I created a wrapper around BatchCoalescer. I wanted to set an impossibly large target batch size, but the problem is that that size is used to allocate the arrays for columns up front and so using something like usize::MAX isn't possible. I instead have to guess at what the maximum batch size will be with the trade-off that I don't want it to be too large and use too much memory. If I run out of space I create a new BatchCoalescer with a larger capacity and re-push the data within the previous BatchCoalescer into that new BatchCoalescer before pushing my new records into it.

Alternatively I could just not use any sort of optimized machinery and instead slice up record batches myself and then concat them myself as well.

Additional context

@timsaucer also recently had a need for a custom BatchCoalescer except he needs it to spit out batches when it either reaches a certain number of rows or a specific total size in bytes. It seems like there might be more cases out there where re-batching based on different criteria is needed.

[tustvold]

If you don't know the sizes up front it might be more performant to buffer up the record batch and then pass slices to the concat_batches kernel. Otherwise you will end up bump allocating, which is not normally very performant.

[alamb]

If you don't know the sizes up front it might be more performant to buffer up the record batch and then pass slices to the concat_batches kernel. Otherwise you will end up bump allocating, which is not normally very performant.

The "GenericInProgressArray" here already basically already implements buffer + concat: https://github.com/apache/arrow-rs/blob/main/arrow-select/src/coalesce/generic.rs

Maybe we could use that somehow (to avoid having to know the total size upfront) 🤔

Or maybe we could add a setting on BatchCoalescer

[alamb]

I am creating new record batches from a stream of record batches. I need to ensure that identical values in a sorted column are always located within the same record batch.

I am not sure what your usecase is, but without additional constraints this could result in an unbounded use of memory

Another approach I have seen used is to store the same data as a Vec<RecordBatch>, and call slice() to split batches where the sort key changes (so the batches line up nicely with the same keys)

You can use the partition kernel to find where the sorted values change

[corasaurus-hex]

Here's what I landed on and it seems to be very fast for my use case: https://gist.github.com/corasaurus-hex/1754460af46b84a2a13ca034e86e5676

This seems generic enough that it could power a lot of different solutions, even something like BatchCoalescer.

without additional constraints this could result in an unbounded use of memory

Without getting too into the problem, I do have additional constraints, constraints in the physical world, that limit this. The input data can change over time but the pathological cases have an upper bound (I'm not 100% on what that upper bound is but the worst case is ~300k-ish right now and that's likely very close). I'd like to keep using as little memory as possible in the other cases, though, so pre-allocating that amount for every batch isn't ideal.

Another approach I have seen used is to store the same data as a Vec, and call slice() to split batches where the sort key changes (so the batches line up nicely with the same keys)

The code I used above is similar to that but without partition (but I'm using partition elsewhere for a join algorithm, it's lovely). I'm not using partition in this case since I think it'll partition the entire batch and in the vast majority of cases I only need to seek to a position and read a few records past that to find a boundary.

[corasaurus-hex]

the iter_column thing is probably the most suspect in that gist. it works for me in my one specific use case but we'd probably want something with better, faster access patterns for the more general case. maybe something like arrow::compute::unary that can be applied to a range?

[alamb]

the iter_column thing is probably the most suspect in that gist. it works for me in my one specific use case but we'd probably want something with better, faster access patterns for the more general case. maybe something like arrow::compute::unary that can be applied to a range?

You can apply unary to a sliced array to get a range for a single array

If you mean to apply some function to a (virtual) range of rows in a Vec<RecordBatch> that also sounds useful but I am not sure where it would go in arrow

[corasaurus-hex]

The iterator function is there to support this sequence of events:

1. push batch onto the view
2. inspect the view starting at an offset and compare values in order to figure out where to slice
3. slice a batch off the view if an appropriate offset can be found

I'm not sure what a good API might be for that, I haven't actually worked with arrow comparison/unary/binary/etc functions, but I want to do something that compares more values at a time.

I'd love some ideas if you have any!