Merging multiple vortex files together and resumable streams

[asubiotto]

I have n vortex files that I would like to merge into a single file with a soft size limit k. My current approach is to join the scan streams using select_all and buffering arrays up to k bytes using nbytes() as a proxy indicator of the resulting file size.

I would like to modify this approach to be fully streaming so that arrays are written as they become ready, but run into a bit of a snag in the case that k is hit but there are still elements in the stream that should be written to another vortex file. My approach is to wrap the stream in a custom ArrayStreamAdapter that returns Poll::Ready(None) when k is hit (calculated as the sum of nbytes()) reading from the merged stream then reset it for the next poll (pseudocode):

        let source_stream = futures::stream::select_all(streams);
        let limited_stream = ByteLimitedStream::new(source_stream, dtype, byte_threshold);

        while stream_not_exhausted {
            limited_stream.reset_bytes();

            let output_buffer = VortexWriteOptions::default()
                .write(Vec::new(), limited_stream)
                .await?;
        }

However, write takes ownership of limited_stream, meaning I can't reuse it on successive iterations:

vortex/vortex-file/src/writer.rs

Lines 103 to 107 in 4f6ea56

	pub async fn write<W: VortexWrite, S: ArrayStream + Unpin + Send + 'static>(
	self,
	write: W,
	stream: S,
	) -> VortexResult<W> {

I realize this might be me just hitting the limits of my rust knowledge, but I feel like I should be able to do something like this and I'm not sure if I need to somehow modify ByteLimitedStream or whether vortex can facilitate what I'm trying to do.

[gatesn]

I had a think about whether we should add a push-based write_chunk(ArrayRef) API, but I'm not sure it actually helps in this case?

I think you may just need to implement the stream logic you want, which is roughly: take an ArrayStream and turn it into a Stream where you start a new stream whenever some threshold has been reached?

I'm about to push a change that makes the writer return a Stream, which may make your accounting a little easier in that you could terminate the stream and start a new one when the file reaches a certain size.

Am I understanding the problem correctly?

[asubiotto]

I had a think about whether we should add a push-based write_chunk(ArrayRef) API, but I'm not sure it actually helps in this case?

I think this makes sense to me and would solve my problem. This is the most like a traditional Writer trait so I could just flush/close a writer whenever the threshold is reached and start pushing to another writer.

I think you may just need to implement the stream logic you want, which is roughly: take an ArrayStream and turn it into a Stream where you start a new stream whenever some threshold has been reached?

IIUC this is what I've been trying (and failing) to do because write takes ownership of the stream. If I create a new stream on each write wrapping the same ArrayStream (select_all of each file), I push the problem up a layer because that stream itself also needs ownership. Passing a mut ref to the source stream into these threshold-limited streams didn't work because the lifetime seemed to clash with the 'static in the write signature. I do think it must be possible to achieve this some way, but I don't have the Rust knowledge to do so, what am I missing?

I'm about to push a change that makes the writer return a Stream, which may make your accounting a little easier in that you could terminate the stream and start a new one when the file reaches a certain size.

This would also help, because ownership of the stream would be passed back to the caller, but it seems like an awkward API for the general case (i.e. a stream is written to completion)

[asubiotto]

My approach to get this working with the current code has been to wrap source_stream in an Arc<Mutex> so that I can clone it and therefore create a new ByteLimitedStream for each write I need to do. It doesn't feel very pretty because this code is not parallel but it seems like my only alternative to that was using an UnsafeCell and implementing Send