alamb commented on code in PR #3391: URL: https://github.com/apache/arrow-rs/pull/3391#discussion_r1059008071
########## arrow-flight/src/encode.rs: ########## @@ -0,0 +1,503 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{collections::VecDeque, fmt::Debug, pin::Pin, sync::Arc, task::Poll}; + +use crate::{error::FlightError, error::Result, FlightData, SchemaAsIpc}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use futures::{ready, stream::BoxStream, Stream, StreamExt}; + +/// Creates a [`Stream`](futures::Stream) of [`FlightData`]s from a +/// `Stream` of [`Result`]<[`RecordBatch`], [`FlightError`]>. +/// +/// This can be used to implement [`FlightService::do_get`] in an +/// Arrow Flight implementation; +/// +/// # Caveats +/// 1. [`DictionaryArray`](arrow_array::array::DictionaryArray)s +/// are converted to their underlying types prior to transport, due to +/// <https://github.com/apache/arrow-rs/issues/3389>. +/// +/// # Example +/// ```no_run +/// # use std::sync::Arc; +/// # use arrow_array::{ArrayRef, RecordBatch, UInt32Array}; +/// # async fn f() { +/// # let c1 = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]); +/// # let record_batch = RecordBatch::try_from_iter(vec![ +/// # ("a", Arc::new(c1) as ArrayRef) +/// # ]) +/// # .expect("cannot create record batch"); +/// use arrow_flight::encode::FlightDataEncoderBuilder; +/// +/// // Get an input stream of Result<RecordBatch, FlightError> +/// let input_stream = futures::stream::iter(vec![Ok(record_batch)]); +/// +/// // Build a stream of `Result<FlightData>` (e.g. to return for do_get) +/// let flight_data_stream = FlightDataEncoderBuilder::new() +/// .build(input_stream); +/// +/// // Create a tonic `Response` that can be returned from a Flight server +/// let response = tonic::Response::new(flight_data_stream); +/// # } +/// ``` +/// +/// [`FlightService::do_get`]: crate::flight_service_server::FlightService::do_get +#[derive(Debug)] +pub struct FlightDataEncoderBuilder { + /// The maximum message size (see details on [`Self::with_max_message_size`]). + max_batch_size: usize, + /// Ipc writer options + options: IpcWriteOptions, + /// Metadata to add to the schema message + app_metadata: Bytes, +} + +/// Default target size for record batches to send. +/// +/// Note this value would normally be 4MB, but the size calculation is +/// somehwhat inexact, so we set it to 2MB. +pub const GRPC_TARGET_MAX_BATCH_SIZE: usize = 2097152; + +impl Default for FlightDataEncoderBuilder { + fn default() -> Self { + Self { + max_batch_size: GRPC_TARGET_MAX_BATCH_SIZE, + options: IpcWriteOptions::default(), + app_metadata: Bytes::new(), + } + } +} + +impl FlightDataEncoderBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Set the (approximate) maximum encoded [`RecordBatch`] size to + /// limit the gRPC message size. Defaults fo 2MB. + /// + /// The encoder splits up [`RecordBatch`]s (preserving order) to + /// limit individual messages to approximately this size. The size + /// is approximate because there additional encoding overhead on + /// top of the underlying data itself. + /// + pub fn with_max_message_size(mut self, max_batch_size: usize) -> Self { + self.max_batch_size = max_batch_size; + self + } + + /// Specfy application specific metadata included in the + /// [`FlightData::app_metadata`] field of the the first Schema + /// message + pub fn with_metadata(mut self, app_metadata: Bytes) -> Self { + self.app_metadata = app_metadata; + self + } + + /// Set the [`IpcWriteOptions`] used to encode the [`RecordBatch`]es for transport. + pub fn with_options(mut self, options: IpcWriteOptions) -> Self { + self.options = options; + self + } + + /// Return a [`Stream`](futures::Stream) of [`FlightData`], + /// consuming self. More details on [`FlightDataEncoder`] + pub fn build<S>(self, input: S) -> FlightDataEncoder + where + S: Stream<Item = Result<RecordBatch>> + Send + 'static, + { + let Self { + max_batch_size, + options, + app_metadata, + } = self; + + FlightDataEncoder::new(input.boxed(), max_batch_size, options, app_metadata) + } +} + +/// Stream that encodes a stream of record batches to flight data. +/// +/// See [`FlightDataEncoderBuilder`] for details and example. +pub struct FlightDataEncoder { + /// Input stream + inner: BoxStream<'static, Result<RecordBatch>>, + /// schema, set after the first batch + schema: Option<SchemaRef>, + /// Max sixe of batches to encode + max_batch_size: usize, + /// do the encoding / tracking of dictionaries + encoder: FlightIpcEncoder, + /// optional metadata to add to schema FlightData + app_metadata: Option<Bytes>, + /// data queued up to send but not yet sent + queue: VecDeque<FlightData>, + /// Is this strema done (inner is empty or errored) + done: bool, +} + +impl FlightDataEncoder { + fn new( + inner: BoxStream<'static, Result<RecordBatch>>, + max_batch_size: usize, + options: IpcWriteOptions, + app_metadata: Bytes, + ) -> Self { + Self { + inner, + schema: None, + max_batch_size, + encoder: FlightIpcEncoder::new(options), + app_metadata: Some(app_metadata), + queue: VecDeque::new(), + done: false, + } + } + + /// Place the `FlightData` in the queue to send + fn queue_message(&mut self, data: FlightData) { + self.queue.push_back(data); + } + + /// Place the `FlightData` in the queue to send + fn queue_messages(&mut self, datas: Vec<FlightData>) { + for data in datas { + self.queue_message(data) + } + } + + /// Encodes batch into one or more `FlightData` messages in self.queue + fn encode_batch(&mut self, batch: RecordBatch) -> Result<()> { + let schema = match self.schema.take() { + Some(schema) => schema, + None => { + let batch_schema = batch.schema(); + // The first message is the schema message, and all + // batches have the same schema + let schema = Arc::new(prepare_schema_for_flight(&batch_schema)); + let mut schema_flight_data = self.encoder.encode_schema(&schema); + + // attach any metadata requested + if let Some(app_metadata) = self.app_metadata.take() { + schema_flight_data.app_metadata = app_metadata; + } + self.queue_message(schema_flight_data); + schema + } + }; + + // remember schema + self.schema = Some(schema.clone()); + + // encode the batch + let batch = prepare_batch_for_flight(&batch, schema)?; + + for batch in split_batch_for_grpc_response(batch, self.max_batch_size) { + let (flight_dictionaries, flight_batch) = + self.encoder.encode_batch(&batch)?; + + self.queue_messages(flight_dictionaries); + self.queue_message(flight_batch); + } + + Ok(()) + } +} + +impl Stream for FlightDataEncoder { + type Item = Result<FlightData>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + loop { + if self.done && self.queue.is_empty() { + return Poll::Ready(None); + } + + // Any messages queued to send? + if let Some(data) = self.queue.pop_front() { + return Poll::Ready(Some(Ok(data))); + } + + // Get next batch + let batch = ready!(self.inner.poll_next_unpin(cx)); + + match batch { + None => { + // inner is done + self.done = true; + } Review Comment: Yes that is correct. This case is also handled on the next loop iteration, but I think making it explicit is good too. I did so in 4f30ab71c -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
