Skip to main content

reqwest_streams/
arrow_body.rs

1//! Streaming an Apache Arrow IPC request body.
2
3use crate::stream_body::{ReqwestStreamBody, StreamBodyRequest};
4use arrow::array::RecordBatch;
5use arrow::datatypes::SchemaRef;
6use futures::Stream;
7use http_streams_core::ArrowRecordBatchIpcStreamFormat;
8
9/// Extension trait for [`reqwest::RequestBuilder`] that streams an Arrow IPC request body.
10///
11/// Unlike decoding, encoding needs the schema up front: it is written once, ahead of the first
12/// batch. See [`ReqwestStreamBody`] for the HTTP caveats.
13pub trait ArrowIpcStreamRequest {
14    /// Streams `stream` as an Arrow IPC stream, setting
15    /// `Content-Type: application/vnd.apache.arrow.stream`.
16    fn arrow_ipc_stream_body<S>(self, schema: SchemaRef, stream: S) -> reqwest::RequestBuilder
17    where
18        S: Stream<Item = RecordBatch> + Send + 'static;
19
20    /// Streams a fallible `stream` as an Arrow IPC stream.
21    fn try_arrow_ipc_stream_body<S, E>(
22        self,
23        schema: SchemaRef,
24        stream: S,
25    ) -> reqwest::RequestBuilder
26    where
27        S: Stream<Item = Result<RecordBatch, E>> + Send + 'static,
28        E: Into<Box<dyn std::error::Error + Send + Sync>> + Send + 'static;
29}
30
31impl ArrowIpcStreamRequest for reqwest::RequestBuilder {
32    fn arrow_ipc_stream_body<S>(self, schema: SchemaRef, stream: S) -> reqwest::RequestBuilder
33    where
34        S: Stream<Item = RecordBatch> + Send + 'static,
35    {
36        self.stream_body(ReqwestStreamBody::new(
37            ArrowRecordBatchIpcStreamFormat::new(schema),
38            stream,
39        ))
40    }
41
42    fn try_arrow_ipc_stream_body<S, E>(
43        self,
44        schema: SchemaRef,
45        stream: S,
46    ) -> reqwest::RequestBuilder
47    where
48        S: Stream<Item = Result<RecordBatch, E>> + Send + 'static,
49        E: Into<Box<dyn std::error::Error + Send + Sync>> + Send + 'static,
50    {
51        self.stream_body(ReqwestStreamBody::try_new(
52            ArrowRecordBatchIpcStreamFormat::new(schema),
53            stream,
54        ))
55    }
56}