feat: Flight supports RecordBatch with dictionary arrays (#6521)

* feat: Flight supports RecordBatch with dictionary arrays

Signed-off-by: luofucong <luofc@foxmail.com>

* resolve PR comments

Signed-off-by: luofucong <luofc@foxmail.com>

---------

Signed-off-by: luofucong <luofc@foxmail.com>

Cargo.lock

@@ -2440,6 +2440,7 @@ dependencies = [
  "tokio-util",
  "tonic 0.12.3",
  "tower 0.5.2",
+ "vec1",
 ]
 
 [[package]]
@@ -13989,6 +13990,12 @@ version = "0.2.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
 
+[[package]]
+name = "vec1"
+version = "1.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eab68b56840f69efb0fefbe3ab6661499217ffdc58e2eef7c3f6f69835386322"
+
 [[package]]
 name = "vergen"
 version = "8.3.2"

src/client/src/database.rs

@@ -42,7 +42,7 @@ use common_telemetry::{error, warn};
 use futures::future;
 use futures_util::{Stream, StreamExt, TryStreamExt};
 use prost::Message;
-use snafu::{ensure, ResultExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue};
 use tonic::transport::Channel;
 
@@ -355,7 +355,10 @@ impl Database {
         let mut flight_message_stream = flight_data_stream.map(move |flight_data| {
             flight_data
                 .map_err(Error::from)
-                .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))
+                .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))?
+                .context(IllegalFlightMessagesSnafu {
+                    reason: "none message",
+                })
         });
 
         let Some(first_flight_message) = flight_message_stream.next().await else {

src/client/src/region.rs

@@ -128,7 +128,10 @@ impl RegionRequester {
         let mut flight_message_stream = flight_data_stream.map(move |flight_data| {
             flight_data
                 .map_err(Error::from)
-                .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))
+                .and_then(|data| decoder.try_decode(&data).context(ConvertFlightDataSnafu))?
+                .context(IllegalFlightMessagesSnafu {
+                    reason: "none message",
+                })
         });
 
         let Some(first_flight_message) = flight_message_stream.next().await else {

src/common/grpc/Cargo.toml

@@ -31,6 +31,7 @@ tokio.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
 tower.workspace = true
+vec1 = "1.12"
 
 [dev-dependencies]
 criterion = "0.4"

src/common/grpc/benches/bench_flight_decoder.rs

@@ -84,9 +84,12 @@ fn prepare_random_record_batch(
 fn prepare_flight_data(num_rows: usize) -> (FlightData, FlightData) {
     let schema = schema();
     let mut encoder = FlightEncoder::default();
-    let schema_data = encoder.encode(FlightMessage::Schema(schema.clone()));
+    let schema_data = encoder.encode_schema(schema.as_ref());
     let rb = prepare_random_record_batch(schema, num_rows);
-    let rb_data = encoder.encode(FlightMessage::RecordBatch(rb));
+    let [rb_data] = encoder
+        .encode(FlightMessage::RecordBatch(rb))
+        .try_into()
+        .unwrap();
     (schema_data, rb_data)
 }
 
@@ -96,7 +99,7 @@ fn decode_flight_data_from_protobuf(schema: &Bytes, payload: &Bytes) -> DfRecord
     let mut decoder = FlightDecoder::default();
     let _schema = decoder.try_decode(&schema).unwrap();
     let message = decoder.try_decode(&payload).unwrap();
-    let FlightMessage::RecordBatch(batch) = message else {
+    let Some(FlightMessage::RecordBatch(batch)) = message else {
         unreachable!("unexpected message");
     };
     batch

src/common/grpc/src/flight.rs

@@ -23,6 +23,7 @@ use arrow_flight::{FlightData, SchemaAsIpc};
 use common_base::bytes::Bytes;
 use common_recordbatch::DfRecordBatch;
 use datatypes::arrow;
+use datatypes::arrow::array::ArrayRef;
 use datatypes::arrow::buffer::Buffer;
 use datatypes::arrow::datatypes::{Schema as ArrowSchema, SchemaRef};
 use datatypes::arrow::error::ArrowError;
@@ -31,6 +32,7 @@ use flatbuffers::FlatBufferBuilder;
 use prost::bytes::Bytes as ProstBytes;
 use prost::Message;
 use snafu::{OptionExt, ResultExt};
+use vec1::{vec1, Vec1};
 
 use crate::error;
 use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result};
@@ -77,9 +79,19 @@ impl FlightEncoder {
         }
     }
 
-    pub fn encode(&mut self, flight_message: FlightMessage) -> FlightData {
+    /// Encode the Arrow schema to [FlightData].
+    pub fn encode_schema(&self, schema: &ArrowSchema) -> FlightData {
+        SchemaAsIpc::new(schema, &self.write_options).into()
+    }
+
+    /// Encode the [FlightMessage] to a list (at least one element) of [FlightData]s.
+    ///
+    /// Normally only when the [FlightMessage] is an Arrow [RecordBatch] with dictionary arrays
+    /// will the encoder produce more than one [FlightData]s. Other types of [FlightMessage] should
+    /// be encoded to exactly one [FlightData].
+    pub fn encode(&mut self, flight_message: FlightMessage) -> Vec1<FlightData> {
         match flight_message {
-            FlightMessage::Schema(schema) => SchemaAsIpc::new(&schema, &self.write_options).into(),
+            FlightMessage::Schema(schema) => vec1![self.encode_schema(schema.as_ref())],
             FlightMessage::RecordBatch(record_batch) => {
                 let (encoded_dictionaries, encoded_batch) = self
                     .data_gen
@@ -90,14 +102,10 @@ impl FlightEncoder {
                     )
                     .expect("DictionaryTracker configured above to not fail on replacement");
 
-                // TODO(LFC): Handle dictionary as FlightData here, when we supported Arrow's Dictionary DataType.
-                // Currently we don't have a datatype corresponding to Arrow's Dictionary DataType,
-                // so there won't be any "dictionaries" here. Assert to be sure about it, and
-                // perform a "testing guard" in case we forgot to handle the possible "dictionaries"
-                // here in the future.
-                debug_assert_eq!(encoded_dictionaries.len(), 0);
-
-                encoded_batch.into()
+                Vec1::from_vec_push(
+                    encoded_dictionaries.into_iter().map(Into::into).collect(),
+                    encoded_batch.into(),
+                )
             }
             FlightMessage::AffectedRows(rows) => {
                 let metadata = FlightMetadata {
@@ -105,12 +113,12 @@ impl FlightEncoder {
                     metrics: None,
                 }
                 .encode_to_vec();
-                FlightData {
+                vec1![FlightData {
                     flight_descriptor: None,
                     data_header: build_none_flight_msg().into(),
                     app_metadata: metadata.into(),
                     data_body: ProstBytes::default(),
-                }
+                }]
             }
             FlightMessage::Metrics(s) => {
                 let metadata = FlightMetadata {
@@ -120,12 +128,12 @@ impl FlightEncoder {
                     }),
                 }
                 .encode_to_vec();
-                FlightData {
+                vec1![FlightData {
                     flight_descriptor: None,
                     data_header: build_none_flight_msg().into(),
                     app_metadata: metadata.into(),
                     data_body: ProstBytes::default(),
-                }
+                }]
             }
         }
     }
@@ -135,6 +143,7 @@ impl FlightEncoder {
 pub struct FlightDecoder {
     schema: Option<SchemaRef>,
     schema_bytes: Option<bytes::Bytes>,
+    dictionaries_by_id: HashMap<i64, ArrayRef>,
 }
 
 impl FlightDecoder {
@@ -145,6 +154,7 @@ impl FlightDecoder {
         Ok(Self {
             schema: Some(Arc::new(arrow_schema)),
             schema_bytes: Some(schema_bytes.clone()),
+            dictionaries_by_id: HashMap::new(),
         })
     }
 
@@ -186,7 +196,13 @@ impl FlightDecoder {
         Ok(result)
     }
 
-    pub fn try_decode(&mut self, flight_data: &FlightData) -> Result<FlightMessage> {
+    /// Try to decode the [FlightData] to a [FlightMessage].
+    ///
+    /// If the [FlightData] is of type `DictionaryBatch` (produced while encoding an Arrow
+    /// [RecordBatch] with dictionary arrays), the decoder will not return any [FlightMessage]s.
+    /// Instead, it will update its internal dictionary cache. Other types of [FlightData] will
+    /// be decoded to exactly one [FlightMessage].
+    pub fn try_decode(&mut self, flight_data: &FlightData) -> Result<Option<FlightMessage>> {
         let message = root_as_message(&flight_data.data_header).map_err(|e| {
             InvalidFlightDataSnafu {
                 reason: e.to_string(),
@@ -198,12 +214,12 @@ impl FlightDecoder {
                 let metadata = FlightMetadata::decode(flight_data.app_metadata.clone())
                     .context(DecodeFlightDataSnafu)?;
                 if let Some(AffectedRows { value }) = metadata.affected_rows {
-                    return Ok(FlightMessage::AffectedRows(value as _));
+                    return Ok(Some(FlightMessage::AffectedRows(value as _)));
                 }
                 if let Some(Metrics { metrics }) = metadata.metrics {
-                    return Ok(FlightMessage::Metrics(
+                    return Ok(Some(FlightMessage::Metrics(
                         String::from_utf8_lossy(&metrics).to_string(),
-                    ));
+                    )));
                 }
                 InvalidFlightDataSnafu {
                     reason: "Expecting FlightMetadata have some meaningful content.",
@@ -219,21 +235,46 @@ impl FlightDecoder {
                 })?);
                 self.schema = Some(arrow_schema.clone());
                 self.schema_bytes = Some(flight_data.data_header.clone());
-                Ok(FlightMessage::Schema(arrow_schema))
+                Ok(Some(FlightMessage::Schema(arrow_schema)))
             }
             MessageHeader::RecordBatch => {
                 let schema = self.schema.clone().context(InvalidFlightDataSnafu {
                     reason: "Should have decoded schema first!",
                 })?;
-                let arrow_batch =
-                    flight_data_to_arrow_batch(flight_data, schema.clone(), &HashMap::new())
-                        .map_err(|e| {
-                            InvalidFlightDataSnafu {
-                                reason: e.to_string(),
-                            }
-                            .build()
+                let arrow_batch = flight_data_to_arrow_batch(
+                    flight_data,
+                    schema.clone(),
+                    &self.dictionaries_by_id,
+                )
+                .map_err(|e| {
+                    InvalidFlightDataSnafu {
+                        reason: e.to_string(),
+                    }
+                    .build()
+                })?;
+                Ok(Some(FlightMessage::RecordBatch(arrow_batch)))
+            }
+            MessageHeader::DictionaryBatch => {
+                let dictionary_batch =
+                    message
+                        .header_as_dictionary_batch()
+                        .context(InvalidFlightDataSnafu {
+                            reason: "could not get dictionary batch from DictionaryBatch message",
                         })?;
-                Ok(FlightMessage::RecordBatch(arrow_batch))
+
+                let schema = self.schema.as_ref().context(InvalidFlightDataSnafu {
+                    reason: "schema message is not present previously",
+                })?;
+
+                reader::read_dictionary(
+                    &flight_data.data_body.clone().into(),
+                    dictionary_batch,
+                    schema,
+                    &mut self.dictionaries_by_id,
+                    &message.version(),
+                )
+                .context(error::ArrowSnafu)?;
+                Ok(None)
             }
             other => {
                 let name = other.variant_name().unwrap_or("UNKNOWN");
@@ -305,14 +346,16 @@ fn build_none_flight_msg() -> Bytes {
 #[cfg(test)]
 mod test {
     use arrow_flight::utils::batches_to_flight_data;
-    use datatypes::arrow::array::Int32Array;
+    use datatypes::arrow::array::{
+        DictionaryArray, Int32Array, StringArray, UInt32Array, UInt8Array,
+    };
     use datatypes::arrow::datatypes::{DataType, Field, Schema};
 
     use super::*;
     use crate::Error;
 
     #[test]
-    fn test_try_decode() {
+    fn test_try_decode() -> Result<()> {
         let schema = Arc::new(ArrowSchema::new(vec![Field::new(
             "n",
             DataType::Int32,
@@ -347,7 +390,7 @@ mod test {
             .to_string()
             .contains("Should have decoded schema first!"));
 
-        let message = decoder.try_decode(d1).unwrap();
+        let message = decoder.try_decode(d1)?.unwrap();
         assert!(matches!(message, FlightMessage::Schema(_)));
         let FlightMessage::Schema(decoded_schema) = message else {
             unreachable!()
@@ -356,19 +399,20 @@ mod test {
 
         let _ = decoder.schema.as_ref().unwrap();
 
-        let message = decoder.try_decode(d2).unwrap();
+        let message = decoder.try_decode(d2)?.unwrap();
         assert!(matches!(message, FlightMessage::RecordBatch(_)));
         let FlightMessage::RecordBatch(actual_batch) = message else {
             unreachable!()
         };
         assert_eq!(actual_batch, batch1);
 
-        let message = decoder.try_decode(d3).unwrap();
+        let message = decoder.try_decode(d3)?.unwrap();
         assert!(matches!(message, FlightMessage::RecordBatch(_)));
         let FlightMessage::RecordBatch(actual_batch) = message else {
             unreachable!()
         };
         assert_eq!(actual_batch, batch2);
+        Ok(())
     }
 
     #[test]
@@ -407,4 +451,86 @@ mod test {
         let actual = flight_messages_to_recordbatches(vec![m1, m2, m3]).unwrap();
         assert_eq!(actual, recordbatches);
     }
+
+    #[test]
+    fn test_flight_encode_decode_with_dictionary_array() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("i", DataType::UInt8, true),
+            Field::new_dictionary("s", DataType::UInt32, DataType::Utf8, true),
+        ]));
+        let batch1 = DfRecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(UInt8Array::from_iter_values(vec![1, 2, 3])) as _,
+                Arc::new(DictionaryArray::new(
+                    UInt32Array::from_value(0, 3),
+                    Arc::new(StringArray::from_iter_values(["x"])),
+                )) as _,
+            ],
+        )
+        .unwrap();
+        let batch2 = DfRecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(UInt8Array::from_iter_values(vec![4, 5, 6, 7, 8])) as _,
+                Arc::new(DictionaryArray::new(
+                    UInt32Array::from_iter_values([0, 1, 2, 2, 3]),
+                    Arc::new(StringArray::from_iter_values(["h", "e", "l", "o"])),
+                )) as _,
+            ],
+        )
+        .unwrap();
+
+        let message_1 = FlightMessage::Schema(schema.clone());
+        let message_2 = FlightMessage::RecordBatch(batch1);
+        let message_3 = FlightMessage::RecordBatch(batch2);
+
+        let mut encoder = FlightEncoder::default();
+        let encoded_1 = encoder.encode(message_1);
+        let encoded_2 = encoder.encode(message_2);
+        let encoded_3 = encoder.encode(message_3);
+        // message 1 is Arrow Schema, should be encoded to one FlightData:
+        assert_eq!(encoded_1.len(), 1);
+        // message 2 and 3 are Arrow RecordBatch with dictionary arrays, should be encoded to
+        // multiple FlightData:
+        assert_eq!(encoded_2.len(), 2);
+        assert_eq!(encoded_3.len(), 2);
+
+        let mut decoder = FlightDecoder::default();
+        let decoded_1 = decoder.try_decode(encoded_1.first())?;
+        let Some(FlightMessage::Schema(actual_schema)) = decoded_1 else {
+            unreachable!()
+        };
+        assert_eq!(actual_schema, schema);
+        let decoded_2 = decoder.try_decode(&encoded_2[0])?;
+        // expected to be a dictionary batch message, decoder should return none:
+        assert!(decoded_2.is_none());
+        let Some(FlightMessage::RecordBatch(decoded_2)) = decoder.try_decode(&encoded_2[1])? else {
+            unreachable!()
+        };
+        let decoded_3 = decoder.try_decode(&encoded_3[0])?;
+        // expected to be a dictionary batch message, decoder should return none:
+        assert!(decoded_3.is_none());
+        let Some(FlightMessage::RecordBatch(decoded_3)) = decoder.try_decode(&encoded_3[1])? else {
+            unreachable!()
+        };
+        let actual = arrow::util::pretty::pretty_format_batches(&[decoded_2, decoded_3])
+            .unwrap()
+            .to_string();
+        let expected = r"
++---+---+
+| i | s |
++---+---+
+| 1 | x |
+| 2 | x |
+| 3 | x |
+| 4 | h |
+| 5 | e |
+| 6 | l |
+| 7 | l |
+| 8 | o |
++---+---+";
+        assert_eq!(actual, expected.trim());
+        Ok(())
+    }
 }

src/common/test-util/src/flight.rs

@@ -19,8 +19,10 @@ use common_recordbatch::DfRecordBatch;
 /// Encodes record batch to a Schema message and a RecordBatch message.
 pub fn encode_to_flight_data(rb: DfRecordBatch) -> (FlightData, FlightData) {
     let mut encoder = FlightEncoder::default();
-    (
-        encoder.encode(FlightMessage::Schema(rb.schema())),
-        encoder.encode(FlightMessage::RecordBatch(rb)),
-    )
+    let schema = encoder.encode_schema(rb.schema_ref().as_ref());
+    let [data] = encoder
+        .encode(FlightMessage::RecordBatch(rb))
+        .try_into()
+        .unwrap();
+    (schema, data)
 }

src/mito2/src/memtable/bulk/part.rs

@@ -93,23 +93,33 @@ impl TryFrom<BulkWalEntry> for BulkPart {
     }
 }
 
-impl From<&BulkPart> for BulkWalEntry {
-    fn from(value: &BulkPart) -> Self {
+impl TryFrom<&BulkPart> for BulkWalEntry {
+    type Error = error::Error;
+
+    fn try_from(value: &BulkPart) -> Result<Self> {
         if let Some(ipc) = &value.raw_data {
-            BulkWalEntry {
+            Ok(BulkWalEntry {
                 sequence: value.sequence,
                 max_ts: value.max_ts,
                 min_ts: value.min_ts,
                 timestamp_index: value.timestamp_index as u32,
                 body: Some(Body::ArrowIpc(ipc.clone())),
-            }
+            })
         } else {
             let mut encoder = FlightEncoder::default();
             let schema_bytes = encoder
-                .encode(FlightMessage::Schema(value.batch.schema()))
+                .encode_schema(value.batch.schema().as_ref())
                 .data_header;
-            let rb_data = encoder.encode(FlightMessage::RecordBatch(value.batch.clone()));
-            BulkWalEntry {
+            let [rb_data] = encoder
+                .encode(FlightMessage::RecordBatch(value.batch.clone()))
+                .try_into()
+                .map_err(|_| {
+                    error::UnsupportedOperationSnafu {
+                        err_msg: "create BulkWalEntry from RecordBatch with dictionary arrays",
+                    }
+                    .build()
+                })?;
+            Ok(BulkWalEntry {
                 sequence: value.sequence,
                 max_ts: value.max_ts,
                 min_ts: value.min_ts,
@@ -119,7 +129,7 @@ impl From<&BulkPart> for BulkWalEntry {
                     data_header: rb_data.data_header,
                     payload: rb_data.data_body,
                 })),
-            }
+            })
         }
     }
 }

src/mito2/src/region/opener.rs

@@ -652,7 +652,13 @@ where
         for bulk_entry in entry.bulk_entries {
             let part = BulkPart::try_from(bulk_entry)?;
             rows_replayed += part.num_rows();
-            region_write_ctx.push_bulk(OptionOutputTx::none(), part);
+            ensure!(
+                region_write_ctx.push_bulk(OptionOutputTx::none(), part),
+                RegionCorruptedSnafu {
+                    region_id,
+                    reason: "unable to replay memtable with bulk entries",
+                }
+            );
         }
 
         // set next_entry_id and write to memtable.

src/mito2/src/region_write_ctx.rs

@@ -252,15 +252,24 @@ impl RegionWriteCtx {
             .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1);
     }
 
-    pub(crate) fn push_bulk(&mut self, sender: OptionOutputTx, mut bulk: BulkPart) {
+    pub(crate) fn push_bulk(&mut self, sender: OptionOutputTx, mut bulk: BulkPart) -> bool {
+        bulk.sequence = self.next_sequence;
+        let entry = match BulkWalEntry::try_from(&bulk) {
+            Ok(entry) => entry,
+            Err(e) => {
+                sender.send(Err(e));
+                return false;
+            }
+        };
+
         self.bulk_notifiers
             .push(WriteNotify::new(sender, bulk.num_rows()));
-        bulk.sequence = self.next_sequence;
 
         // Add bulk wal entry
-        self.wal_entry.bulk_entries.push(BulkWalEntry::from(&bulk));
+        self.wal_entry.bulk_entries.push(entry);
         self.next_sequence += bulk.num_rows() as u64;
         self.bulk_parts.push(bulk);
+        true
     }
 
     pub(crate) async fn write_bulk(&mut self) {

src/mito2/src/worker/handle_write.rs

@@ -393,7 +393,9 @@ impl<S> RegionWorkerLoop<S> {
             }
 
             // Collect requests by region.
-            region_ctx.push_bulk(bulk_req.sender, bulk_req.request);
+            if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request) {
+                return;
+            }
         }
     }
 

src/operator/src/bulk_insert.rs

@@ -27,7 +27,7 @@ use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage};
 use common_grpc::FlightData;
 use common_telemetry::error;
 use common_telemetry::tracing_context::TracingContext;
-use snafu::{OptionExt, ResultExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use store_api::storage::RegionId;
 use table::metadata::TableInfoRef;
 use table::TableRef;
@@ -53,7 +53,10 @@ impl Inserter {
         // Build region server requests
         let message = decoder
             .try_decode(&data)
-            .context(error::DecodeFlightDataSnafu)?;
+            .context(error::DecodeFlightDataSnafu)?
+            .context(error::NotSupportedSnafu {
+                feat: "bulk insert RecordBatch with dictionary arrays",
+            })?;
         let FlightMessage::RecordBatch(record_batch) = message else {
             return Ok(0);
         };
@@ -195,8 +198,20 @@ impl Inserter {
                             let encode_timer = metrics::HANDLE_BULK_INSERT_ELAPSED
                                 .with_label_values(&["encode"])
                                 .start_timer();
-                            let flight_data =
-                                FlightEncoder::default().encode(FlightMessage::RecordBatch(batch));
+                            let mut iter = FlightEncoder::default()
+                                .encode(FlightMessage::RecordBatch(batch))
+                                .into_iter();
+                            let Some(flight_data) = iter.next() else {
+                                // Safety: `iter` on a type of `Vec1`, which is guaranteed to have
+                                // at least one element.
+                                unreachable!()
+                            };
+                            ensure!(
+                                iter.next().is_none(),
+                                error::NotSupportedSnafu {
+                                    feat: "bulk insert RecordBatch with dictionary arrays",
+                                }
+                            );
                             encode_timer.observe_duration();
                             (flight_data.data_header, flight_data.data_body)
                         };

src/servers/src/grpc/flight.rs

@@ -383,9 +383,12 @@ fn to_flight_data_stream(
             Box::pin(stream) as _
         }
         OutputData::AffectedRows(rows) => {
-            let stream = tokio_stream::once(Ok(
-                FlightEncoder::default().encode(FlightMessage::AffectedRows(rows))
-            ));
+            let stream = tokio_stream::iter(
+                FlightEncoder::default()
+                    .encode(FlightMessage::AffectedRows(rows))
+                    .into_iter()
+                    .map(Ok),
+            );
             Box::pin(stream) as _
         }
     }

src/servers/src/grpc/flight/stream.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::VecDeque;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
@@ -40,6 +41,7 @@ pub struct FlightRecordBatchStream {
     join_handle: JoinHandle<()>,
     done: bool,
     encoder: FlightEncoder,
+    buffer: VecDeque<FlightData>,
 }
 
 impl FlightRecordBatchStream {
@@ -64,6 +66,7 @@ impl FlightRecordBatchStream {
             join_handle,
             done: false,
             encoder,
+            buffer: VecDeque::new(),
         }
     }
 
@@ -128,6 +131,9 @@ impl Stream for FlightRecordBatchStream {
         if *this.done {
             Poll::Ready(None)
         } else {
+            if let Some(x) = this.buffer.pop_front() {
+                return Poll::Ready(Some(Ok(x)));
+            }
             match this.rx.poll_next(cx) {
                 Poll::Ready(None) => {
                     *this.done = true;
@@ -135,8 +141,14 @@ impl Stream for FlightRecordBatchStream {
                 }
                 Poll::Ready(Some(result)) => match result {
                     Ok(flight_message) => {
-                        let flight_data = this.encoder.encode(flight_message);
-                        Poll::Ready(Some(Ok(flight_data)))
+                        let mut iter = this.encoder.encode(flight_message).into_iter();
+                        let Some(first) = iter.next() else {
+                            // Safety: `iter` on a type of `Vec1`, which is guaranteed to have
+                            // at least one element.
+                            unreachable!()
+                        };
+                        this.buffer.extend(iter);
+                        Poll::Ready(Some(Ok(first)))
                     }
                     Err(e) => {
                         *this.done = true;
@@ -191,7 +203,7 @@ mod test {
         let decoder = &mut FlightDecoder::default();
         let mut flight_messages = raw_data
             .into_iter()
-            .map(|x| decoder.try_decode(&x).unwrap())
+            .map(|x| decoder.try_decode(&x).unwrap().unwrap())
             .collect::<Vec<FlightMessage>>();
         assert_eq!(flight_messages.len(), 2);
 

tests-integration/src/grpc/flight.rs

@@ -140,7 +140,7 @@ mod test {
         let schema = record_batches[0].schema.arrow_schema().clone();
 
         let stream = futures::stream::once(async move {
-            let mut schema_data = FlightEncoder::default().encode(FlightMessage::Schema(schema));
+            let mut schema_data = FlightEncoder::default().encode_schema(schema.as_ref());
             let metadata = DoPutMetadata::new(0);
             schema_data.app_metadata = serde_json::to_vec(&metadata).unwrap().into();
             // first message in "DoPut" stream should carry table name in flight descriptor
@@ -154,13 +154,15 @@ mod test {
         .chain(
             tokio_stream::iter(record_batches)
                 .enumerate()
-                .map(|(i, x)| {
+                .flat_map(|(i, x)| {
                     let mut encoder = FlightEncoder::default();
                     let message = FlightMessage::RecordBatch(x.into_df_record_batch());
                     let mut data = encoder.encode(message);
                     let metadata = DoPutMetadata::new((i + 1) as i64);
-                    data.app_metadata = serde_json::to_vec(&metadata).unwrap().into();
-                    data
+                    data.iter_mut().for_each(|x| {
+                        x.app_metadata = serde_json::to_vec(&metadata).unwrap().into()
+                    });
+                    tokio_stream::iter(data)
                 })
                 .boxed(),
         )