feat: stream data back for CSV and JSON queries #25927

mgattozzi · 2025-01-28T18:33:03Z

This commit allows us to stream data back for CSV and JSON formatted queries. Prior to this we would buffer up all of the data in memory before sending it back. Now we can make it so that we only buffer in one RecordBatch at a time to reduce memory overhead.

Note that due to the way the APIs for writers work and for how Body in hyper 0.14 works we can't use a streaming body that we can write too. This in turn means we have to use a manually written Future state machine, which works but is far from ideal.

Note this does not include the pretty and parquet files as streamable. I'm attempting to get the pretty one to be streamable, but I don't think that this one and parquet are as likely to be streamable back to the user. In general we might want to discourage these formats from being used.

waynr

I left a couple questions inline, but otherwise this looks good to me so I'm leaving a ✔️.

influxdb3_server/src/http.rs

waynr · 2025-01-28T22:53:56Z

influxdb3_server/src/http.rs

+                                    .with_header(false)
+                                    .build(Vec::new())
+                            };
+                            writer.write(&batch).unwrap();


Prior to this PR and still in the fully-buffered branches of the format match statement it looks like writer.write(&batch)? is being called -- is it not possible to do that here because of the Output type of this future?

Would this work?

Suggested change

writer.write(&batch).unwrap();

if let Err(e) = writer.write(&batch) {

return Poll::Ready(Some(Err(e)));

}

Ah yeah, but I think it would need .into() called: https://docs.rs/datafusion/latest/datafusion/error/enum.DataFusionError.html#impl-From%3CError%3E-for-DataFusionError-1, right? I was initially assuming there was no impl From<std::io::Error> for DatafusionError and that it would need to be convered to the Error defined in this module.

This ended up working out well with the into.

As for the writer.write(&batch) we can't use the ? anymore because you now need to return a Poll type. async is essentially doing all this under the hood, but since we're writing the future ourselves by hand we have this problem where we need to hand wrap these errors to work.

There is no From impl to make ? work for Poll as far as I'm aware.

hiltontj

I left a couple of comments. FWIW - the v1 /query API implementation uses a streamed response via a dedicated type:

influxdb/influxdb3_server/src/http/v1.rs

Lines 81 to 83 in 56ca85e

    
           let stream = QueryResponseStream::new(0, stream, chunk_size, format, epoch, group_by) 
        
               .map_err(QueryError)?; 
        
           let body = Body::wrap_stream(stream);

Not sure if that is at all helpful.

Cargo.toml

hiltontj · 2025-01-30T21:50:20Z

influxdb3_server/src/http.rs

+                                    .with_header(false)
+                                    .build(Vec::new())
+                            };
+                            writer.write(&batch).unwrap();


Would this work?

Suggested change

writer.write(&batch).unwrap();

if let Err(e) = writer.write(&batch) {

return Poll::Ready(Some(Err(e)));

}

influxdb3_server/src/http.rs

hiltontj · 2025-01-30T22:01:37Z

influxdb3_server/src/http.rs

+            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)))


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 for writing parquet in a streaming fashion in the compactor in enterprise, so I wonder if we could adopt a similar approach here?

Unfortunately no, because there's no way to write to a Body with a writer. If there was I wouldn't be writing these wild state machines by hand. I think for tables and parquet we'll just have to leave them be. The main things we expect people to stream at least are done.

Do you know if this is a limitation of hyper 0.14 (that we could solve by going to 1.x) or something that would need to be solved upstream parquet/arrow regardless of hyper version?

I think that streaming for --format=parquet will need to be addressed at some point. I agree for tables though, that is purely for use in the CLI.

This commit allows us to stream data back for CSV and JSON formatted queries. Prior to this we would buffer up all of the data in memory before sending it back. Now we can make it so that we only buffer in one RecordBatch at a time to reduce memory overhead. Note that due to the way the APIs for writers work and for how Body in hyper 0.14 works we can't use a streaming body that we can write too. This in turn means we have to use a manually written Future state machine, which works but is far from ideal. Note this does not include the pretty and parquet files as streamable. I'm attempting to get the pretty one to be streamable, but I don't think that this one and parquet are as likely to be streamable back to the user. In general we might want to discourage these formats from being used.

mgattozzi · 2025-01-31T16:55:46Z

@waynr and @hiltontj I incorporated your suggestions. Let me know if this needs anything else!

hiltontj

I think this looks good, though I do think we need to address streaming parquet format at some point.

mgattozzi · 2025-01-31T18:25:24Z

@hiltontj I think so too, but I think it'll require some upstream work in Arrow to work with hyper's Body which can be a trait implemented in 1.0 so that writers can produce Frames in order to stream the data I think. I'm not sure how to square the circle otherwise from my attempts

hiltontj · 2025-02-03T13:58:05Z

Related: #25955

In #25927 we missed that JSON queries were broken despite having some tests use the format. This fixes JSON queries such that they now properly contain a comma between RecordBatches. This commit also includes tests for the formats that now stream data back (CSV, JSON, and JSON Lines) so that we won't run into this issue again.

mgattozzi requested review from pauldix, waynr, hiltontj, jacksonrnewhouse and praveen-influx January 28, 2025 18:33

waynr approved these changes Jan 28, 2025

View reviewed changes

hiltontj reviewed Jan 30, 2025

View reviewed changes

mgattozzi requested review from waynr and hiltontj January 31, 2025 16:53

mgattozzi force-pushed the mgattozzi/streaming-query branch from 1ac4698 to 72c1847 Compare January 31, 2025 16:53

mgattozzi force-pushed the mgattozzi/streaming-query branch from 72c1847 to 09133b9 Compare January 31, 2025 16:55

hiltontj approved these changes Jan 31, 2025

View reviewed changes

mgattozzi merged commit 20fdc7b into main Jan 31, 2025
13 checks passed

mgattozzi deleted the mgattozzi/streaming-query branch January 31, 2025 18:25

hiltontj mentioned this pull request Feb 3, 2025

Stream query response for Parquet format #25955

Open

mgattozzi mentioned this pull request Feb 6, 2025

fix: broken format for JSON queries and add tests #25980

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: stream data back for CSV and JSON queries #25927

feat: stream data back for CSV and JSON queries #25927

mgattozzi commented Jan 28, 2025

waynr left a comment

waynr Jan 28, 2025

hiltontj Jan 30, 2025

waynr Jan 30, 2025

mgattozzi Jan 31, 2025

hiltontj left a comment

hiltontj Jan 30, 2025

hiltontj Jan 30, 2025

mgattozzi Jan 31, 2025

hiltontj Jan 31, 2025

mgattozzi commented Jan 31, 2025

hiltontj left a comment

mgattozzi commented Jan 31, 2025

hiltontj commented Feb 3, 2025

-                            writer.write(&batch).unwrap();
+                            if let Err(e) = writer.write(&batch) {
+                                return Poll::Ready(Some(Err(e)));
+                            }

	let stream = QueryResponseStream::new(0, stream, chunk_size, format, epoch, group_by)
	.map_err(QueryError)?;
	let body = Body::wrap_stream(stream);

feat: stream data back for CSV and JSON queries #25927

feat: stream data back for CSV and JSON queries #25927

Conversation

mgattozzi commented Jan 28, 2025

waynr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hiltontj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mgattozzi commented Jan 31, 2025

hiltontj left a comment

Choose a reason for hiding this comment

mgattozzi commented Jan 31, 2025

hiltontj commented Feb 3, 2025