diff --git a/Cargo.lock b/Cargo.lock index 1e0d5a09e23..6fbfa8e9fd9 100644 --- a/Cargo.lock +++ b/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" diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 48e9a208def..a3898b2a206 100644 --- a/src/client/src/database.rs +++ b/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 { diff --git a/src/client/src/region.rs b/src/client/src/region.rs index fa634f6b60d..d4a7e5b36c0 100644 --- a/src/client/src/region.rs +++ b/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 { diff --git a/src/common/grpc/Cargo.toml b/src/common/grpc/Cargo.toml index b56167ee155..32d096c87ab 100644 --- a/src/common/grpc/Cargo.toml +++ b/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" diff --git a/src/common/grpc/benches/bench_flight_decoder.rs b/src/common/grpc/benches/bench_flight_decoder.rs index 9e4de209828..5e72e55fce0 100644 --- a/src/common/grpc/benches/bench_flight_decoder.rs +++ b/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 diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs index 36a9a93decd..a5c57162df9 100644 --- a/src/common/grpc/src/flight.rs +++ b/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 { 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, schema_bytes: Option, + dictionaries_by_id: HashMap, } 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 { + /// 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> { 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(()) + } } diff --git a/src/common/test-util/src/flight.rs b/src/common/test-util/src/flight.rs index 2735017eea0..2f0e63f7665 100644 --- a/src/common/test-util/src/flight.rs +++ b/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) } diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 295f5cb7f07..9d5f7ef8862 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -93,23 +93,33 @@ impl TryFrom 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 { 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, })), - } + }) } } } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 885963f21ee..1b355606182 100644 --- a/src/mito2/src/region/opener.rs +++ b/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. diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index c05c66b9efb..395a33ec20c 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/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) { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index d3389ccefa7..aa98ef1ee59 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -393,7 +393,9 @@ impl RegionWorkerLoop { } // 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; + } } } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 35a031dd43a..dfd02398a9b 100644 --- a/src/operator/src/bulk_insert.rs +++ b/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) }; diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 37c2cf05f25..8500181cc8d 100644 --- a/src/servers/src/grpc/flight.rs +++ b/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 _ } } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 99498bf3162..1937fb63462 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/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, } 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::>(); assert_eq!(flight_messages.len(), 2); diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 85953dd1ee6..27e5fbc88d3 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/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(), )