diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs index 1245a86efb..4f5a47cb53 100644 --- a/src/common/grpc/src/flight.rs +++ b/src/common/grpc/src/flight.rs @@ -37,6 +37,35 @@ use vec1::{Vec1, vec1}; use crate::error; use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result}; +/// Encodes a batch into the schema, record-batch header and payload used by bulk inserts. +/// +/// Uses the default Flight compression and rejects dictionary batches, whose +/// additional Flight messages cannot be represented by this three-part format. +pub fn record_batch_to_ipc( + record_batch: DfRecordBatch, +) -> Result<(ProstBytes, ProstBytes, ProstBytes)> { + let mut encoder = FlightEncoder::default(); + let schema = encoder.encode_schema(record_batch.schema().as_ref()); + let mut iter = encoder + .encode(FlightMessage::RecordBatch(record_batch)) + .into_iter(); + let flight_data = iter.next().context(InvalidFlightDataSnafu { + reason: "Failed to encode empty flight data", + })?; + if iter.next().is_some() { + return error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + .fail(); + } + + Ok(( + schema.data_header, + flight_data.data_header, + flight_data.data_body, + )) +} + /// Flight metadata key used to carry flow query extensions as JSON pairs. pub const FLOW_EXTENSIONS_METADATA_KEY: &str = "x-greptime-flow-extensions"; /// Flight metadata key used to carry query snapshot read upper bounds as JSON. @@ -374,8 +403,52 @@ mod test { use datatypes::arrow::buffer::OffsetBuffer; use datatypes::arrow::datatypes::{DataType, Field, Schema}; - use super::*; use crate::Error; + use crate::flight::*; + + #[test] + fn test_record_batch_to_ipc_preserves_wire_bytes() { + let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, true)])); + for values in [vec![], vec![Some(1), None, Some(3)]] { + let batch = + DfRecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(values))]) + .unwrap(); + // Reproduce the previous Prom encoding sequence and compare all bytes. + let mut encoder = FlightEncoder::default(); + let expected_schema = encoder.encode_schema(batch.schema().as_ref()); + let messages = encoder.encode(FlightMessage::RecordBatch(batch.clone())); + assert_eq!(messages.len(), 1); + let expected_batch = messages.first(); + assert_eq!( + record_batch_to_ipc(batch).unwrap(), + ( + expected_schema.data_header, + expected_batch.data_header.clone(), + expected_batch.data_body.clone(), + ) + ); + } + } + + #[test] + fn test_record_batch_to_ipc_rejects_dictionary() { + let schema = Arc::new(Schema::new(vec![Field::new_dictionary( + "tag", + DataType::UInt32, + DataType::Utf8, + true, + )])); + let dictionary = DictionaryArray::new( + UInt32Array::from_value(0, 3), + Arc::new(StringArray::from_iter_values(["x"])), + ); + let batch = DfRecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap(); + assert!(matches!( + record_batch_to_ipc(batch), + Err(Error::NotSupported { feat }) + if feat == "bulk insert RecordBatch with dictionary arrays" + )); + } #[test] fn test_try_decode() -> Result<()> { diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 76d0e4d8cf..4c36296db2 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -28,7 +28,6 @@ use arrow::compute::{concat_batches, filter_record_batch}; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; -use bytes::Bytes; use catalog::CatalogManagerRef; use common_batcher::flush_limiter::FlushLimiter; use common_batcher::flush_policy::FlushTrigger; @@ -37,7 +36,8 @@ use common_batcher::notifier::{Notifier, run_notifier}; use common_batcher::pending_worker::PendingWorker as PendingCore; use common_batcher::request_limiter::RequestLimiter; use common_batcher::worker_registry::WorkerRegistry; -use common_grpc::flight::{FlightEncoder, FlightMessage}; +use common_grpc::error::Error as GrpcError; +use common_grpc::flight::record_batch_to_ipc; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_timestamp, greptime_value}; @@ -1799,7 +1799,7 @@ fn encode_region_write_requests( let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_encode_ipc"]) .start_timer(); - record_batch_to_ipc(resolved.planned.batch)? + record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)? }; let request = RegionRequest { @@ -1841,28 +1841,14 @@ fn notify_waiters(waiters: Vec, result: Result<()>) { } } -fn record_batch_to_ipc(record_batch: RecordBatch) -> Result<(Bytes, Bytes, Bytes)> { - let mut encoder = FlightEncoder::default(); - let schema = encoder.encode_schema(record_batch.schema().as_ref()); - let mut iter = encoder - .encode(FlightMessage::RecordBatch(record_batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - return Err(Error::Internal { - err_msg: "Failed to encode empty flight data".to_string(), - }); - }; - if iter.next().is_some() { - return Err(Error::NotSupported { - feat: "bulk insert RecordBatch with dictionary arrays".to_string(), - }); +fn map_ipc_error(error: GrpcError) -> Error { + match error { + GrpcError::NotSupported { feat } => Error::NotSupported { feat }, + GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason }, + error => Error::Internal { + err_msg: error.to_string(), + }, } - - Ok(( - schema.data_header, - flight_data.data_header, - flight_data.data_body, - )) } #[cfg(test)]