SendableRecordBatchStream showing inconsistent behaviour for a custom wrapper implementation

[twitu]

Based on the discussion in #7675, I removed ORDER BY <field> from my query. Now the data is no longer brought into memory to be sorted and directly streamed as chunks from disk. HOWEVER, a weird problem started showing up with the new query.

The underlying SendableRecordBatchStream doesn't pull up new RecordBatch from disk. Awaiting the stream simply hangs the execution. It's clear somehow the custom implementation is interfering with the normal behaviour but it's not clear how that's possible. I've created a minimal reproducible example below.

So a Flattener wraps a stream and puts it in a Vec. The Flattener stream pops out the stream from the Vec polls for the next element and then puts the stream back in the Vec. The actual logic in production uses a binary heap to pop streams in a sorted manner but this a minimal version to reproduce the issue.

SimpleWrapper is a passthrough wrapper just polls the underlying stream and returns the result. I've added it in as a baseline implementation that works.

    pin_project! {
        struct Flattener<S>
        {
            data: Vec<S>,
        }
    }

    pin_project! {
        struct SimpleWrap<S>
        {
            data: S
        }
    }

    impl<S> Stream for Flattener<S>
    where
        S: Stream + Unpin,
    {
        type Item = S::Item;

        fn poll_next(
            self: std::pin::Pin<&mut Self>,
            cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Option<Self::Item>> {
            let this = self.project();
            match this.data.pop() {
                Some(mut stream) => match ready!(stream.poll_next_unpin(cx)) {
                    Some(element) => {
                        this.data.push(stream);
                        Poll::Ready(Some(element))
                    }
                    _ => {
                        dbg!("No next element");
                        Poll::Ready(None)
                    }
                },
                _ => {
                    dbg!("empty vec");
                    Poll::Ready(None)
                }
            }
        }
    }

    impl<S> Stream for SimpleWrap<S>
    where
        S: Stream + Unpin,
    {
        type Item = S::Item;

        fn poll_next(
            self: std::pin::Pin<&mut Self>,
            cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<Option<Self::Item>> {
            let this = self.project();
            if let Some(element) = ready!(this.data.poll_next_unpin(cx)) {
                dbg!("Got next element");
                Poll::Ready(Some(element))
            } else {
                Poll::Ready(None)
            }
        }
    }

Below is the actual test logic, which registers a file on disk as a table and queries for all its records. There's also a mock data stream created from a vec of numbers and put in a flattener.

    let file_path = "../../tests/test_data/quote_tick_data.parquet";
    let table_name = "data";
    let session_ctx = SessionContext::default();
    let parquet_options = ParquetReadOptions::<'_> {
        skip_metadata: Some(false),
        ..Default::default()
    };
    session_ctx
        .register_parquet(table_name, file_path, parquet_options)
        .await
        .unwrap();

    let batch_stream = session_ctx
        .sql(&format!("SELECT * FROM {} ORDER BY ts_init", &table_name))
        // .sql(&format!("SELECT * FROM {}", &table_name))
        .await
        .unwrap()
        .execute_stream()
        .await
        .unwrap();

    let mock_data: Vec<Result<usize, ()>> = (0..100000).map(Ok).collect();
    let mut mock_data = futures::stream::iter(mock_data);



    let mut iter = Flattener {
        data: vec![batch_stream],
    };

    let mut mock_iter = Flattener {
        data: vec![mock_data],
    };

    // let mut iter = SimpleWrap {
    //     data: batch_stream,
    // };

    let step = 10_000;
    while let Some(c) = mock_iter.next().await {
        c.map(|val| if val % step == 0 { dbg!(val) } else { 0 })
            .unwrap_or(0);
    }

    dbg!("Stream batch");
    while let Some(c) = iter.next().await {
        dbg!(c.map(|batch| batch.num_rows()).unwrap_or(0));
    }

Here's the results for the various combinations

Stream	Result
Flattener - record batch no order by	⚔️ terminates immediately finding no elements in the stream
Flattener - record batch order by	⚔️ hangs
SimpleWrap - record batch no order by	✔️
SImpleWrap - record batch order by	✔️
Flattener - mock data	✔️

My guess is that it has the to be the Vec push and pop that's messing with the behavior but it's so weird and not clear why it's messing up the underlying stream.

[alamb]

I didn't quite follow how Flattener is working, but if you want to convert a SendableRecordBatchStream to a Vec<RecordBatch> you could perhaps use `collect: https://docs.rs/datafusion/latest/datafusion/physical_plan/fn.collect.html

[alamb]

BTW

                Some(mut stream) => match ready!(stream.poll_next_unpin(cx)) {

That will return if the underlying stream is not ready, even if there is something in your buffer, right?

[twitu]

You're absolutely right the ready macro was returning even when the stream wasn't ready does skipping values and causing weird behavior. And fixing this thingy led to a much more performance design 😄

[twitu]

About your previous comment. I don't actually want to collect the stream into a RecordBatch because I want to keep the memory usage low by streaming and consuming one RecordBatch at a time.

The larger context is that I'm query multiple files and "sort of" merging their data together, which is why I needed that vec to store all the record batch streams and pop the right one (based on an ordering) and then fetch it's next element. All this while keeping the streaming model of materializing only what's necessary.