diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index f8d645bd3e1..19cb403307f 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -22,6 +22,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::task::{Context, Poll}; use common_base::readable_size::ReadableSize; +use common_error::ext::BoxedError; use common_telemetry::tracing::{Span, info_span}; use common_time::util::format_nanoseconds_human_readable; use datafusion::arrow::compute::cast; @@ -55,6 +56,7 @@ use crate::{ }; const REGION_SCAN_EXEC_NAME: &str = "RegionScanExec"; +pub const SERIES_SCAN_EXEC_NAME: &str = "SeriesScanExec"; type FutureStream = Pin> + Send>>; @@ -205,7 +207,9 @@ impl Stream for DfRecordBatchStreamAdapter { Poll::Pending => Poll::Pending, Poll::Ready(Some(recordbatch)) => match recordbatch { Ok(recordbatch) => Poll::Ready(Some(Ok(recordbatch.into_df_record_batch()))), - Err(e) => Poll::Ready(Some(Err(DataFusionError::External(Box::new(e))))), + Err(e) => Poll::Ready(Some(Err(DataFusionError::External(Box::new( + BoxedError::new(e), + ))))), }, Poll::Ready(None) => Poll::Ready(None), } @@ -377,12 +381,16 @@ impl RecordBatchStreamAdapter { } } +fn is_region_scan_name(name: &str) -> bool { + name == REGION_SCAN_EXEC_NAME || name == SERIES_SCAN_EXEC_NAME +} + /// Extracts total `output_bytes` from region scan plan nodes. pub fn region_scan_output_bytes(metrics: &RecordBatchMetrics) -> usize { metrics .plan_metrics .iter() - .filter(|pm| pm.plan_name == REGION_SCAN_EXEC_NAME) + .filter(|pm| is_region_scan_name(&pm.plan_name)) .flat_map(|pm| &pm.metrics) .filter_map(|(name, value)| (name == "output_bytes").then_some(*value)) .sum() @@ -434,7 +442,7 @@ fn collect_lightweight_query_load_metrics_inner( level: usize, record_batch_metrics: &mut RecordBatchMetrics, ) { - let is_region_scan = df_plan.name() == REGION_SCAN_EXEC_NAME; + let is_region_scan = is_region_scan_name(df_plan.name()); let mut region_scan_output_bytes = None; if let Some(metrics) = df_plan.metrics() { @@ -937,9 +945,10 @@ fn convert_map_to_json_binary( #[cfg(test)] mod test { + use std::collections::VecDeque; use std::time::Duration; - use common_error::ext::BoxedError; + use common_error::ext::{ErrorExt, PlainError}; use common_error::mock::MockError; use common_error::status_code::StatusCode; use datafusion::execution::TaskContext; @@ -959,14 +968,15 @@ mod test { use snafu::IntoError; use super::*; - use crate::RecordBatches; use crate::error::Error; + use crate::{RecordBatchStreamWrapper, RecordBatches}; #[derive(Debug)] struct TestMetricsExec { properties: Arc, metrics: ExecutionPlanMetricsSet, format_plan: bool, + plan_name: &'static str, } impl TestMetricsExec { @@ -975,17 +985,23 @@ mod test { } fn new_without_plan_formatting(schema: DfSchemaRef) -> Self { - Self::with_output_bytes_and_formatting(schema, &[24], false) + Self::with_output_bytes_and_formatting(schema, &[24], false, REGION_SCAN_EXEC_NAME) } fn with_output_bytes(schema: DfSchemaRef, output_bytes_by_partition: &[usize]) -> Self { - Self::with_output_bytes_and_formatting(schema, output_bytes_by_partition, true) + Self::with_output_bytes_and_formatting( + schema, + output_bytes_by_partition, + true, + REGION_SCAN_EXEC_NAME, + ) } fn with_output_bytes_and_formatting( schema: DfSchemaRef, output_bytes_by_partition: &[usize], format_plan: bool, + plan_name: &'static str, ) -> Self { let metrics = ExecutionPlanMetricsSet::new(); let elapsed_compute = MetricBuilder::new(&metrics).elapsed_compute(0); @@ -1004,6 +1020,7 @@ mod test { )), metrics, format_plan, + plan_name, } } } @@ -1014,7 +1031,7 @@ mod test { self.format_plan, "non-verbose lightweight partial metrics must not format the plan" ); - write!(f, "RegionScanExec")?; + write!(f, "{}", self.plan_name)?; if matches!(t, DisplayFormatType::Verbose) { write!(f, ": files=[file-1.parquet]")?; } @@ -1024,7 +1041,7 @@ mod test { impl ExecutionPlan for TestMetricsExec { fn name(&self) -> &str { - REGION_SCAN_EXEC_NAME + self.plan_name } fn properties(&self) -> &Arc { @@ -1064,6 +1081,29 @@ mod test { } } + #[test] + fn test_lightweight_query_load_metrics_sums_series_scan_output_bytes() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let plan = TestMetricsExec::with_output_bytes_and_formatting( + schema.arrow_schema().clone(), + &[24, 18], + true, + SERIES_SCAN_EXEC_NAME, + ); + + let metrics = collect_lightweight_query_load_metrics(&plan, Some(42)); + + assert_eq!(region_scan_output_bytes(&metrics), 42); + assert_eq!(metrics.plan_metrics[0].plan_name, SERIES_SCAN_EXEC_NAME); + + let full_metrics = collect_full_metrics(&plan, false, Some(42)); + assert_eq!(region_scan_output_bytes(&full_metrics), 42); + } + #[test] fn test_lightweight_query_load_metrics_sums_region_scan_output_bytes() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -1313,6 +1353,160 @@ mod test { ); } + /// Invariant: both stream adapters must forward the original error as a + /// typed `BoxedError` inside `DataFusionError::External`, so downstream + /// status code recovery (the `Error::PollStream` arm of `Error::status_code`) + /// keeps the original status code. + #[test] + fn test_df_record_batch_stream_adapter_preserves_typed_errors() { + fn chain_has_marker(err: &(dyn std::error::Error + 'static), marker: &str) -> bool { + let mut current = Some(err); + while let Some(error) = current { + if error.to_string().contains(marker) { + return true; + } + current = error.source(); + } + false + } + + fn assert_typed_external( + df_error: &DataFusionError, + expected: StatusCode, + context: &str, + marker: &str, + ) { + let DataFusionError::External(payload) = df_error else { + panic!("{context}: expected DataFusionError::External, got {df_error:?}"); + }; + let Some(boxed) = payload.downcast_ref::() else { + panic!("{context}: external payload lost its typed BoxedError: {payload:?}"); + }; + assert_eq!( + boxed.status_code(), + expected, + "{context}: original status code lost" + ); + assert!( + chain_has_marker(boxed, marker), + "{context}: original error lost from the source chain: {payload:?}" + ); + } + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as _], + ) + .unwrap(); + let expected_df_batch = batch.clone().into_df_record_batch(); + + // A stream is consumed once, so every adapter under test gets its own + // fresh equivalent stream. Sequence: Pending -> batch -> typed error -> EOF. + let make_stream = |marker: &str, status: StatusCode| -> SendableRecordBatchStream { + let typed_error = error::ExternalSnafu + .into_error(BoxedError::new(PlainError::new(marker.to_string(), status))); + let mut items: VecDeque>>> = VecDeque::from([ + Poll::Pending, + Poll::Ready(Some(Ok(batch.clone()))), + Poll::Ready(Some(Err(typed_error))), + Poll::Ready(None), + ]); + Box::pin(RecordBatchStreamWrapper::new( + schema.clone(), + futures::stream::poll_fn(move |_| items.pop_front().unwrap_or(Poll::Ready(None))), + )) + }; + + let waker = futures::task::noop_waker_ref(); + let mut cx = Context::from_waker(waker); + + for status in [StatusCode::Internal, StatusCode::RequestOutdated] { + // Unique per case, so a match proves the *original* error survived. + let marker = format!("df-record-batch-stream-adapter-typed-error-marker-{status}"); + + // Greptime -> DataFusion forward adapter. + let mut df_adapter = DfRecordBatchStreamAdapter::new(make_stream(&marker, status)); + assert_eq!(df_adapter.schema(), schema.arrow_schema().clone()); + assert!( + matches!(Pin::new(&mut df_adapter).poll_next(&mut cx), Poll::Pending), + "Pending must be preserved" + ); + let Poll::Ready(Some(Ok(df_batch))) = Pin::new(&mut df_adapter).poll_next(&mut cx) + else { + panic!("expected the forwarded batch"); + }; + assert_eq!(df_batch.schema(), schema.arrow_schema().clone()); + assert_eq!(df_batch, expected_df_batch); + let Poll::Ready(Some(Err(df_error))) = Pin::new(&mut df_adapter).poll_next(&mut cx) + else { + panic!("expected the forwarded error"); + }; + assert_typed_external( + &df_error, + status, + &format!("forward adapter, status {status}"), + &marker, + ); + assert!( + matches!( + Pin::new(&mut df_adapter).poll_next(&mut cx), + Poll::Ready(None) + ), + "EOF must be preserved" + ); + + // DataFusion -> Greptime adapter over another fresh equivalent + // stream: Pending, batch, schema and EOF stay unchanged, and the + // error is reported exactly once. + let df_stream: DfSendableRecordBatchStream = Box::pin(DfRecordBatchStreamAdapter::new( + make_stream(&marker, status), + )); + let mut roundtrip = RecordBatchStreamAdapter::try_new(df_stream).unwrap(); + assert_eq!(roundtrip.schema().arrow_schema(), schema.arrow_schema()); + assert!( + matches!(Pin::new(&mut roundtrip).poll_next(&mut cx), Poll::Pending), + "Pending must be preserved" + ); + let Poll::Ready(Some(Ok(roundtrip_batch))) = + Pin::new(&mut roundtrip).poll_next(&mut cx) + else { + panic!("expected the roundtripped batch"); + }; + assert_eq!(roundtrip_batch.schema, roundtrip.schema()); + assert_eq!(roundtrip_batch.df_record_batch(), &expected_df_batch); + let Poll::Ready(Some(Err(roundtrip_error))) = + Pin::new(&mut roundtrip).poll_next(&mut cx) + else { + panic!("expected the roundtripped error"); + }; + let Error::PollStream { + error: df_error, .. + } = &roundtrip_error + else { + panic!("expected a poll stream error, got {roundtrip_error:?}"); + }; + assert_eq!(roundtrip_error.status_code(), status); + assert_typed_external( + df_error, + status, + &format!("roundtrip adapter, status {status}"), + &marker, + ); + assert!( + matches!( + Pin::new(&mut roundtrip).poll_next(&mut cx), + Poll::Ready(None) + ), + "the roundtrip stream must end after the error" + ); + } + } + #[test] fn test_convert_map_to_json_binary() { let keys = StringArray::from(vec![Some("a"), Some("b"), Some("c"), Some("x")]);