viirya commented on code in PR #3391: URL: https://github.com/apache/arrow-rs/pull/3391#discussion_r1058660172
########## 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 Review Comment: ```suggestion /// Specify application specific metadata included in the ``` ########## arrow-flight/src/decode.rs: ########## @@ -0,0 +1,396 @@ +// 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 crate::{utils::flight_data_to_arrow_batch, FlightData}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::Schema; +use bytes::Bytes; +use futures::{ready, stream::BoxStream, Stream, StreamExt}; +use std::{ + collections::HashMap, convert::TryFrom, fmt::Debug, pin::Pin, sync::Arc, task::Poll, +}; + +use crate::error::{FlightError, Result}; + +/// Decodes a [Stream] of [`FlightData`] back into +/// [`RecordBatch`]es. This can be used to decode the response from an +/// Arrow Flight server +/// +/// # Note +/// To access the lower level Flight messages (e.g. to access +/// [`FlightData::app_metadata`]), you can call [`Self::into_inner`] +/// and use the [`FlightDataDecoder`] directly. +/// +/// # Example: +/// ```no_run +/// # async fn f() -> Result<(), arrow_flight::error::FlightError>{ +/// # use bytes::Bytes; +/// // make a do_get request +/// use arrow_flight::{ +/// error::Result, +/// decode::FlightRecordBatchStream, +/// Ticket, +/// flight_service_client::FlightServiceClient +/// }; +/// use tonic::transport::Channel; +/// use futures::stream::{StreamExt, TryStreamExt}; +/// +/// let client: FlightServiceClient<Channel> = // make client.. +/// # unimplemented!(); +/// +/// let request = tonic::Request::new( +/// Ticket { ticket: Bytes::new() } +/// ); +/// +/// // Get a stream of FlightData; +/// let flight_data_stream = client +/// .do_get(request) +/// .await? +/// .into_inner(); +/// +/// // Decode stream of FlightData to RecordBatches +/// let record_batch_stream = FlightRecordBatchStream::new_from_flight_data( +/// // convert tonic::Status to FlightError +/// flight_data_stream.map_err(|e| e.into()) +/// ); +/// +/// // Read back RecordBatches +/// while let Some(batch) = record_batch_stream.next().await { +/// match batch { +/// Ok(batch) => { /* process batch */ }, +/// Err(e) => { /* handle error */ }, +/// }; +/// } +/// +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug)] +pub struct FlightRecordBatchStream { + inner: FlightDataDecoder, + got_schema: bool, +} + +impl FlightRecordBatchStream { + /// Create a new [`FlightRecordBatchStream`] from a decoded stream + pub fn new(inner: FlightDataDecoder) -> Self { + Self { + inner, + got_schema: false, + } + } + + /// Create a new [`FlightRecordBatchStream`] from a stream of [`FlightData`] + pub fn new_from_flight_data<S>(inner: S) -> Self + where + S: Stream<Item = Result<FlightData>> + Send + 'static, + { + Self { + inner: FlightDataDecoder::new(inner), + got_schema: false, + } + } + + /// Has a message defining the schema been received yet? + pub fn got_schema(&self) -> bool { + self.got_schema + } + + /// Consume self and return the wrapped [`FlightDataDecoder`] + pub fn into_inner(self) -> FlightDataDecoder { + self.inner + } +} +impl futures::Stream for FlightRecordBatchStream { + type Item = Result<RecordBatch>; + + /// Returns the next [`RecordBatch`] available in this stream, or `None` if + /// there are no further results available. + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + let res = ready!(self.inner.poll_next_unpin(cx)); + match res { + // Inner exhausted + None => { + return Poll::Ready(None); + } + Some(Err(e)) => { + return Poll::Ready(Some(Err(e))); + } + // translate data + Some(Ok(data)) => match data.payload { + DecodedPayload::Schema(_) if self.got_schema => { + return Poll::Ready(Some(Err(FlightError::protocol( + "Unexpectedly saw multiple Schema messages in FlightData stream", + )))); + } + DecodedPayload::Schema(_) => { + self.got_schema = true; + // Need next message, poll inner again + } + DecodedPayload::RecordBatch(batch) => { + return Poll::Ready(Some(Ok(batch))); + } + DecodedPayload::None => { + // Need next message + } + }, + } + } + } +} + +/// Wrapper around a stream of [`FlightData`] that handles the details +/// of decoding low level Flight messages into [`Schema`] and +/// [`RecordBatch`]es, including details such as dictionaries. +/// +/// # Protocol Details +/// +/// The client handles flight messages as followes: +/// +/// - **None:** This message has no effect. This is useful to +/// transmit metadata without any actual payload. +/// +/// - **Schema:** The schema is (re-)set. Dictionaries are cleared and +/// the decoded schema is returned. +/// +/// - **Dictionary Batch:** A new dictionary for a given column is registered. An existing +/// dictionary for the same column will be overwritten. This +/// message is NOT visible. +/// +/// - **Record Batch:** Record batch is created based on the current +/// schema and dictionaries. This fails if no schema was transmitted +/// yet. +/// +/// All other message types (at the time of writing: e.g. tensor and +/// sparse tensor) lead to an error. +/// +/// Example usecases +/// +/// 1. Using this low level stream it is possible to receive a steam +/// of RecordBatches in FlightData that have different schemas by +/// handling multiple schema messages separately. +pub struct FlightDataDecoder { + /// Underlying data stream + response: BoxStream<'static, Result<FlightData>>, + /// Decoding state + state: Option<FlightStreamState>, + /// seen the end of the inner stream? Review Comment: nit: ```suggestion /// Seen the end of the inner stream? ``` ########## 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 Review Comment: ```suggestion /// Max size of batches to encode ``` ########## 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: Once reaching here, I guess the queue is guaranteed to be empty? If so, seems we can just `return Poll::Ready(None);` here? -- 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]
