rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-13 14:27:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(common-recordbatch): recognize SeriesScanExec in metrics and preserve typed errors

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-17 11:03:51 +08:00
parent 9ce017653e
commit 3caba1b7c2
1 changed files with 203 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/recordbatch/src/adapter.rs
+203 -9
View File
@@ -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<Box<dyn std::future::Future<Output = Result<SendableRecordBatchStream>> + 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<PlanProperties>,
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<PlanProperties> {
@@ -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::<BoxedError>() 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<Poll<Option<Result<RecordBatch>>>> = 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")]);