Repository: influxdata/influxdb
File: influxdb3_server/src/http.rs (Rust)

Code before:
                            writer.write(&batch).unwrap();
                            Poll::Ready(Some(Ok(Bytes::from(writer.into_inner()))))
                        }
                        Poll::Ready(None) => Poll::Ready(None),
                        Poll::Pending => Poll::Pending,
                    }
                }
            }

            let mut future = CsvFuture {
                first_poll: true,
                stream,
            };
            Ok(Body::wrap_stream(futures::stream::poll_fn(move |ctx| {
                future.poll_unpin(ctx)
            })))
        }
        QueryFormat::Json => {
            struct JsonFuture {
                state: State,
                stream: Pin<Box<dyn RecordBatchStream + Send>>,
            }

            enum State {
                FirstPoll,
                Body,
                Done,
            }

            #[derive(Default, Debug)]
            struct JsonMiddle;

            impl arrow_json::writer::JsonFormat for JsonMiddle {
                fn start_row<W: std::io::Write>(
                    &self,
                    writer: &mut W,
                    is_first_row: bool,
                ) -> std::result::Result<(), arrow_schema::ArrowError> {
                    if !is_first_row {
                        writer.write_all(b",")?;
                    }

                    Ok(())
                }
            }

            impl Future for JsonFuture {
                type Output = Option<Result<Bytes, DataFusionError>>;

                fn poll(
                    mut self: Pin<&mut Self>,

Review comment (performance):
So, this still buffers the whole record batch stream in to memory. I am not so sure how to do this for Parquet; I reckon that you wouldn't be able to take a similar approach to how you've done for JSON/CSV, since (I assume) the parquet writer needs to live for the lifetime of the entire stream.

We do however use the [`AsyncArrowWriter`](https://docs.rs/parquet/latest/parquet/arrow/async_writer/struct.AsyncArrowWriter.html) for writing parquet in a streaming fashion in the compactor in enterprise, so I wonder if we could adopt a similar approach here?

Code after:
    match format {
        QueryFormat::Pretty => {
            let batches = stream.try_collect::<Vec<RecordBatch>>().await?;
            Ok(Body::from(Bytes::from(format!(
                "{}",
                pretty::pretty_format_batches(&batches)?
            ))))
        }
        QueryFormat::Parquet => {
            // Grab the first batch so that we can get the schema
            let Some(batch) = stream.next().await.transpose()? else {
                return Ok(Body::empty());
            };
            let schema = batch.schema();

            let mut bytes = Vec::new();
            let mem_pool = Arc::new(UnboundedMemoryPool::default());
            let mut writer = TrackedMemoryArrowWriter::try_new(&mut bytes, schema, mem_pool)?;

            // Write the first batch we got and then continue writing batches
            writer.write(batch)?;
            while let Some(batch) = stream.next().await.transpose()? {
                writer.write(batch)?;
            }
            writer.close()?;
            Ok(Body::from(Bytes::from(bytes)))
        }
        QueryFormat::Csv => {
            struct CsvFuture {
                first_poll: bool,
                stream: Pin<Box<dyn RecordBatchStream + Send>>,
            }

            impl Future for CsvFuture {
                type Output = Option<Result<Bytes, DataFusionError>>;

                fn poll(
                    mut self: Pin<&mut Self>,
                    ctx: &mut std::task::Context<'_>,
                ) -> Poll<Self::Output> {
                    match self.stream.poll_next_unpin(ctx) {
                        Poll::Ready(Some(batch)) => {
                            let batch = match batch {
                                Ok(batch) => batch,
                                Err(e) => return Poll::Ready(Some(Err(e))),
                            };
                            let mut writer = if self.first_poll {
                                self.first_poll = false;
                                arrow_csv::Writer::new(Vec::new())
                            } else {
                                arrow_csv::WriterBuilder::new()