diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 0bd3a610602..8f9155e2ae4 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -39,8 +39,8 @@ use servers::metrics::METRIC_ROW_BATCH_ELIGIBILITY_FALLBACKS; use servers::otlp; use servers::otlp::trace::span::TraceSpanGroup; use servers::pending_rows_batcher::{ - MetricRowBatchProtocol, MetricRowBatcherRef, is_scalar_metric_batchable, - metric_table_matches_physical_table, + MetricRecordBatchSubmission, MetricRowBatchProtocol, MetricRowBatchSubmission, + MetricRowBatcherRef, is_scalar_metric_batchable, }; use servers::query_handler::{ MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, @@ -55,52 +55,10 @@ use table::requests::{ use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; - -impl Instance { - async fn metric_table_matches_physical_table( - &self, - table_name: &str, - physical_table: &str, - ctx: &QueryContextRef, - ) -> ServerResult { - let table = self - .catalog_manager() - .table( - ctx.current_catalog(), - &ctx.current_schema(), - table_name, - Some(ctx.as_ref()), - ) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu)?; - Ok(table.is_none_or(|table| { - metric_table_matches_physical_table(&table.table_info(), physical_table) - })) - } - - async fn existing_metric_tables_match_physical_table( - &self, - requests: &api::v1::RowInsertRequests, - ctx: &QueryContextRef, - ) -> ServerResult { - let physical_table = ctx - .extension(PHYSICAL_TABLE_PARAM) - .unwrap_or(GREPTIME_PHYSICAL_TABLE); - - for request in &requests.inserts { - if !self - .metric_table_matches_physical_table(&request.table_name, physical_table, ctx) - .await? - { - return Ok(false); - } - } - - Ok(true) - } -} +use crate::metrics::{ + OTLP_LOGS_ROWS, OTLP_METRICS_INGEST_STAGE_ELAPSED, OTLP_METRICS_ROWS, + OTLP_RESOURCE_INFO_WRITE_ERRORS, +}; fn trace_permission_targets( table_name: &str, @@ -133,6 +91,23 @@ fn trace_permission_targets( PermissionTableTargets::resolved(targets) } +fn metric_write_context( + ctx: &QueryContextRef, + semantic_index: &otlp::metrics::SemanticIndex, + is_legacy: bool, +) -> QueryContextRef { + let mut ctx = ctx.as_ref().clone(); + ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); + ctx.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); + if let Some(index) = semantic_index.encode(&ctx.current_schema()) { + ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); + } + if !is_legacy { + ctx.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string()); + } + Arc::new(ctx) +} + #[async_trait] impl OpenTelemetryProtocolHandler for Instance { #[tracing::instrument(skip_all)] @@ -142,11 +117,16 @@ impl OpenTelemetryProtocolHandler for Instance { metric_row_batcher: Option, ctx: QueryContextRef, ) -> ServerResult { - self.plugins - .get::() - .as_ref() - .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE)) - .context(AuthSnafu)?; + { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["request_permission"]) + .start_timer(); + self.plugins + .get::() + .as_ref() + .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE)) + .context(AuthSnafu)?; + } let interceptor_ref = self .plugins @@ -162,7 +142,12 @@ impl OpenTelemetryProtocolHandler for Instance { .collect::>(); // See [`OtlpMetricCtx`] for details - let is_legacy = self.check_otlp_legacy(&input_names, &ctx).await?; + let is_legacy = { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["legacy_check"]) + .start_timer(); + self.check_otlp_legacy(&input_names, &ctx).await? + }; let mut metric_ctx = ctx .protocol_ctx() @@ -172,13 +157,84 @@ impl OpenTelemetryProtocolHandler for Instance { metric_ctx.is_legacy = is_legacy; metric_ctx.resource_info = self.otlp_resource_info; + let mut batching_candidate = + !metric_ctx.is_legacy && metric_ctx.with_metric_engine && metric_row_batcher.is_some(); + if batching_candidate { + let direct_conversion = { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["direct_conversion"]) + .start_timer(); + otlp::metrics::try_to_scalar_record_batches(&request, &metric_ctx)? + }; + if let Some(otlp::metrics::ScalarMetricsConversion { + batches, + permission_tables, + rows, + semantic_index, + mut outcome, + }) = direct_conversion + { + if outcome.accepted_data_points == 0 { + return Ok(outcome); + } + let targets = PermissionTableTargets::resolved( + permission_tables + .into_iter() + .map(|table_name| { + PermissionTableTarget::new( + ctx.current_catalog(), + ctx.current_schema(), + table_name, + ) + }) + .collect(), + ); + self.check_table_permission(&ctx, PermissionReq::Action(OTLP_WRITE), targets) + .context(AuthSnafu)?; + let write_ctx = metric_write_context(&ctx, &semantic_index, is_legacy); + let Some(metric_row_batcher) = metric_row_batcher.as_ref() else { + return Err(error::InternalSnafu { + err_msg: "OTLP batching candidate has no batcher".to_string(), + } + .build()); + }; + let submission = { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["batch_submit"]) + .start_timer(); + metric_row_batcher + .submit_record_batches_with_fallback( + batches, + write_ctx, + MetricRowBatchProtocol::Otlp, + ) + .await? + }; + if submission == MetricRecordBatchSubmission::Submitted { + self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; + OTLP_METRICS_ROWS.inc_by(rows as u64); + outcome.write_cost = 0; + return Ok(outcome); + } + METRIC_ROW_BATCH_ELIGIBILITY_FALLBACKS + .with_label_values(&[MetricRowBatchProtocol::Otlp.as_str()]) + .inc(); + batching_candidate = false; + } + } + let otlp::metrics::MetricsConversion { requests, rows, semantic_index, resource_info, mut outcome, - } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + } = { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["conversion"]) + .start_timer(); + otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)? + }; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP metrics data points: {}", @@ -190,64 +246,73 @@ impl OpenTelemetryProtocolHandler for Instance { return Ok(outcome); } - self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) - .context(AuthSnafu)?; + { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["row_permission"]) + .start_timer(); + self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) + .context(AuthSnafu)?; + } self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; OTLP_METRICS_ROWS.inc_by(rows as u64); - let ctx = { - let mut c = (*ctx).clone(); - c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); - c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); - // Per-table metric specifics + resource/scope lineage ride this - // internal channel; the auto-create path folds them per schema and - // table name. - if let Some(index) = semantic_index.encode(&c.current_schema()) { - c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); - } - if !is_legacy { - c.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string()); - } - Arc::new(c) - }; + let ctx = metric_write_context(&ctx, &semantic_index, is_legacy); - let batching_candidate = - !metric_ctx.is_legacy && metric_ctx.with_metric_engine && metric_row_batcher.is_some(); - let batchable = batching_candidate - && is_scalar_metric_batchable(&requests)? - && self - .existing_metric_tables_match_physical_table(&requests, &ctx) - .await?; - let write_cost = - if let Some(metric_row_batcher) = metric_row_batcher.as_ref().filter(|_| batchable) { + let scalar_batchable = if batching_candidate { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["batchability"]) + .start_timer(); + is_scalar_metric_batchable(&requests)? + } else { + false + }; + let direct_requests = if let Some(metric_row_batcher) = + metric_row_batcher.as_ref().filter(|_| scalar_batchable) + { + let submission = { + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["batch_submit"]) + .start_timer(); metric_row_batcher - .submit(requests, ctx.clone(), MetricRowBatchProtocol::Otlp) - .await?; - 0 - } else { - if batching_candidate { - METRIC_ROW_BATCH_ELIGIBILITY_FALLBACKS - .with_label_values(&[MetricRowBatchProtocol::Otlp.as_str()]) - .inc(); - } - // OTLP tables have one sample field in both the legacy and physical paths. - let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx.clone(), false, true) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - } else { - let physical_table = ctx - .extension(PHYSICAL_TABLE_PARAM) - .unwrap_or(GREPTIME_PHYSICAL_TABLE) - .to_string(); - self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) - .await - .map_err(BoxedError::new) - .context(error::ExecuteGrpcQuerySnafu) - }?; - output.meta.cost + .submit_with_fallback(requests, ctx.clone(), MetricRowBatchProtocol::Otlp) + .await? }; + match submission { + MetricRowBatchSubmission::Submitted => None, + MetricRowBatchSubmission::UnmatchedPhysicalTable(requests) => Some(requests), + } + } else { + Some(requests) + }; + let write_cost = if let Some(requests) = direct_requests { + if batching_candidate { + METRIC_ROW_BATCH_ELIGIBILITY_FALLBACKS + .with_label_values(&[MetricRowBatchProtocol::Otlp.as_str()]) + .inc(); + } + let _timer = OTLP_METRICS_INGEST_STAGE_ELAPSED + .with_label_values(&["direct_write"]) + .start_timer(); + // OTLP tables have one sample field in both the legacy and physical paths. + let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + self.handle_row_inserts(requests, ctx.clone(), false, true) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + } else { + let physical_table = ctx + .extension(PHYSICAL_TABLE_PARAM) + .unwrap_or(GREPTIME_PHYSICAL_TABLE) + .to_string(); + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) + .await + .map_err(BoxedError::new) + .context(error::ExecuteGrpcQuerySnafu) + }?; + output.meta.cost + } else { + 0 + }; outcome.write_cost = write_cost; // Derived enrichment, written after the metric data is committed: @@ -393,6 +458,7 @@ impl OpenTelemetryProtocolHandler for Instance { #[cfg(test)] mod tests { + use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -422,7 +488,9 @@ mod tests { }; use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource as MetricResource; use servers::metrics::METRIC_ROW_BATCH_ELIGIBILITY_FALLBACKS; - use servers::pending_rows_batcher::{MetricRowBatchProtocol, MetricRowBatcher}; + use servers::pending_rows_batcher::{ + MetricRecordBatch, MetricRecordBatchSubmission, MetricRowBatchProtocol, MetricRowBatcher, + }; use session::context::QueryContext; use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; @@ -436,9 +504,11 @@ mod tests { #[derive(Default)] struct FakeMetricRowBatcher { requests: Mutex>, + record_batch_tables: Mutex>>, protocols: Mutex>, flush_error: Option<&'static str>, accepted_rows: u64, + reject_placement: bool, } impl FakeMetricRowBatcher { @@ -455,6 +525,13 @@ mod tests { ..Default::default() } } + + fn rejecting_placement() -> Self { + Self { + reject_placement: true, + ..Default::default() + } + } } #[async_trait] @@ -475,6 +552,47 @@ mod tests { } Ok(self.accepted_rows) } + + async fn submit_with_fallback( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> ServerResult { + if self.reject_placement { + self.requests.lock().unwrap().push(requests.clone()); + self.protocols.lock().unwrap().push(protocol); + return Ok(MetricRowBatchSubmission::UnmatchedPhysicalTable(requests)); + } + self.submit(requests, ctx, protocol) + .await + .map(|_| MetricRowBatchSubmission::Submitted) + } + + async fn submit_record_batches_with_fallback( + &self, + batches: Vec, + _ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> ServerResult { + self.record_batch_tables.lock().unwrap().push( + batches + .iter() + .map(|batch| batch.table_name().to_string()) + .collect(), + ); + self.protocols.lock().unwrap().push(protocol); + if self.reject_placement { + return Ok(MetricRecordBatchSubmission::UnmatchedPhysicalTable); + } + if let Some(message) = self.flush_error { + return Err(servers::error::InternalSnafu { + err_msg: message.to_string(), + } + .build()); + } + Ok(MetricRecordBatchSubmission::Submitted) + } } async fn test_instance() -> Instance { @@ -537,6 +655,7 @@ mod tests { #[derive(Default)] struct CountingPermissionChecker { checks: AtomicUsize, + targets: Mutex>, } impl PermissionChecker for CountingPermissionChecker { @@ -553,9 +672,10 @@ mod tests { &self, _user_info: UserInfoRef, _req: PermissionReq, - _targets: PermissionTableTargets, + targets: PermissionTableTargets, ) -> auth::error::Result { self.checks.fetch_add(1, Ordering::Relaxed); + self.targets.lock().unwrap().push(targets); Ok(PermissionResp::Allow) } } @@ -665,9 +785,79 @@ mod tests { assert_eq!(1, outcome.accepted_data_points); assert_eq!(0, outcome.rejected_data_points); assert_eq!(0, outcome.write_cost); - let requests = batcher.requests.lock().unwrap(); - assert_eq!(1, requests.len()); - assert_eq!("temperature", requests[0].inserts[0].table_name); + assert!(batcher.requests.lock().unwrap().is_empty()); + assert_eq!( + vec![vec!["temperature".to_string()]], + *batcher.record_batch_tables.lock().unwrap() + ); + assert_eq!( + vec![MetricRowBatchProtocol::Otlp], + *batcher.protocols.lock().unwrap() + ); + } + + #[tokio::test] + async fn test_direct_metrics_permission_includes_empty_metric_table() { + let checker = Arc::new(CountingPermissionChecker::default()); + let plugins = Plugins::new(); + plugins.insert::(checker.clone()); + let instance = test_instance_with(FrontendOptions::default(), plugins).await; + let batcher = Arc::new(FakeMetricRowBatcher::successful(1)); + let mut empty = gauge_metric("empty_temperature"); + let Some(metric::Data::Gauge(gauge)) = empty.data.as_mut() else { + panic!("expected gauge") + }; + gauge.data_points.clear(); + + instance + .metrics( + metrics_request(vec![gauge_metric("temperature"), empty]), + Some(batcher), + metric_query_ctx(true), + ) + .await + .unwrap(); + + let targets = checker.targets.lock().unwrap(); + let PermissionTableTargets::Resolved(targets) = &targets[0] else { + panic!("expected resolved table targets") + }; + let table_names = targets + .iter() + .map(|target| target.table.as_str()) + .collect::>(); + assert_eq!( + HashSet::from(["temperature", "empty_temperature"]), + table_names + ); + } + + #[tokio::test] + async fn test_metrics_falls_back_when_batcher_rejects_placement() { + let instance = test_instance().await; + let batcher = Arc::new(FakeMetricRowBatcher::rejecting_placement()); + + let error = instance + .metrics( + gauge_request("temperature"), + Some(batcher.clone()), + metric_query_ctx(true), + ) + .await + .unwrap_err(); + + let servers::error::Error::ExecuteGrpcQuery { source, .. } = error else { + panic!("expected ExecuteGrpcQuery from Inserter, got {error:?}"); + }; + assert!( + source + .as_any() + .downcast_ref::() + .is_some_and(|source| matches!(source, crate::error::Error::TableOperation { .. })), + "expected frontend TableOperation from Inserter, got {source:?}" + ); + assert!(batcher.requests.lock().unwrap().is_empty()); + assert_eq!(1, batcher.record_batch_tables.lock().unwrap().len()); assert_eq!( vec![MetricRowBatchProtocol::Otlp], *batcher.protocols.lock().unwrap() diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index f649db0b46f..f41243f5b34 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -45,6 +45,15 @@ lazy_static! { ) .unwrap(); + /// Elapsed time of OTLP metrics ingest stages on the frontend node. + pub static ref OTLP_METRICS_INGEST_STAGE_ELAPSED: HistogramVec = register_histogram_vec!( + "greptime_frontend_otlp_metrics_ingest_stage_elapsed", + "Elapsed time of OTLP metrics ingest stages on the frontend node", + &["stage"], + vec![0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 60.0] + ) + .unwrap(); + /// Failed writes of the synthesized OTLP resource descriptor table; these /// surface as OTLP partial-success warnings, not request failures. pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!( diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 0f2dea314c8..37cc2ed3a9d 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -12,13 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use ahash::HashSet; +use std::sync::Arc; + +use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{ Count as PromCount, ResetHint, ZeroCount as PromZeroCount, }; use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as PromHistogram}; use api::v1::value::ValueData; -use api::v1::{RowInsertRequests, SemanticType, Value}; +use api::v1::{ColumnDataType, ColumnSchema, RowInsertRequests, SemanticType, Value}; +use arrow::array::{ArrayRef, Float64Array, StringBuilder, TimestampNanosecondArray}; +use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit}; +use arrow::record_batch::RecordBatch; use common_grpc::precision::Precision; use common_query::native_histogram::{ encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, @@ -34,6 +39,7 @@ use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetrics use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; use otel_arrow_rust::proto::opentelemetry::metrics::v1::{metric, number_data_point, *}; use session::protocol_ctx::{MetricType, OtlpMetricCtx}; +use snafu::ResultExt; use table::requests::{ METADATA_QUALITY_DECLARED, METRIC_TEMPORALITY_CUMULATIVE, METRIC_TEMPORALITY_DELTA, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_ORIGINAL_NAME, SEMANTIC_METRIC_TEMPORALITY, @@ -42,6 +48,7 @@ use table::requests::{ use crate::error::{self, Result}; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; +use crate::pending_rows_batcher::MetricRecordBatch; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; @@ -119,6 +126,212 @@ pub struct MetricsConversion { pub outcome: MetricsIngestOutcome, } +/// A Gauge/Sum-only OTLP conversion that materializes Arrow arrays directly. +pub struct ScalarMetricsConversion { + /// Per-table scalar metric batches. + pub batches: Vec, + /// All emitted logical tables, including metrics without data points. + pub permission_tables: Vec, + /// Total metric row count. + pub rows: usize, + /// Semantic metadata used by logical-table auto-creation. + pub semantic_index: SemanticIndex, + /// OTLP accepted/rejected data point counts. + pub outcome: MetricsIngestOutcome, +} + +enum ScalarColumnValues { + Tag(Vec>), + Timestamp(Vec>), + Field(Vec>), +} + +struct ScalarColumn { + name: String, + values: ScalarColumnValues, +} + +struct ScalarTableBuilder { + name: String, + columns: Vec, + column_indexes: HashMap, + rows: usize, +} + +impl ScalarTableBuilder { + fn new(name: String) -> Self { + Self { + name, + columns: Vec::with_capacity(APPROXIMATE_COLUMN_COUNT), + column_indexes: HashMap::with_capacity(APPROXIMATE_COLUMN_COUNT), + rows: 0, + } + } + + fn begin_row(&mut self) { + for column in &mut self.columns { + match &mut column.values { + ScalarColumnValues::Tag(values) => values.push(None), + ScalarColumnValues::Timestamp(values) => values.push(None), + ScalarColumnValues::Field(values) => values.push(None), + } + } + self.rows += 1; + } + + fn write_tag(&mut self, name: &str, value: &str) -> Result<()> { + let index = self.ensure_column(name, ScalarColumnKind::Tag)?; + let ScalarColumnValues::Tag(values) = &mut self.columns[index].values else { + return self.column_kind_error(name, "tag"); + }; + values[self.rows - 1] = Some(value.to_string()); + Ok(()) + } + + fn write_timestamp(&mut self, timestamp_nanos: u64) -> Result<()> { + let timestamp = i64::try_from(timestamp_nanos).map_err(|_| { + error::TimestampOverflowSnafu { + error: format!( + "timestamp {timestamp_nanos} overflow with precision {}", + Precision::Nanosecond + ), + } + .build() + })?; + let name = greptime_timestamp(); + let index = self.ensure_column(name, ScalarColumnKind::Timestamp)?; + let ScalarColumnValues::Timestamp(values) = &mut self.columns[index].values else { + return self.column_kind_error(name, "timestamp"); + }; + values[self.rows - 1] = Some(timestamp); + Ok(()) + } + + fn write_field(&mut self, value: f64) -> Result<()> { + let name = greptime_value(); + let index = self.ensure_column(name, ScalarColumnKind::Field)?; + let ScalarColumnValues::Field(values) = &mut self.columns[index].values else { + return self.column_kind_error(name, "field"); + }; + values[self.rows - 1] = Some(value); + Ok(()) + } + + fn ensure_column(&mut self, name: &str, kind: ScalarColumnKind) -> Result { + if let Some(index) = self.column_indexes.get(name).copied() { + let matches = matches!( + (&self.columns[index].values, kind), + (ScalarColumnValues::Tag(_), ScalarColumnKind::Tag) + | ( + ScalarColumnValues::Timestamp(_), + ScalarColumnKind::Timestamp + ) + | (ScalarColumnValues::Field(_), ScalarColumnKind::Field) + ); + if !matches { + return self.column_kind_error(name, kind.as_str()); + } + return Ok(index); + } + + let values = match kind { + ScalarColumnKind::Tag => ScalarColumnValues::Tag(vec![None; self.rows]), + ScalarColumnKind::Timestamp => ScalarColumnValues::Timestamp(vec![None; self.rows]), + ScalarColumnKind::Field => ScalarColumnValues::Field(vec![None; self.rows]), + }; + let index = self.columns.len(); + self.columns.push(ScalarColumn { + name: name.to_string(), + values, + }); + self.column_indexes.insert(name.to_string(), index); + Ok(index) + } + + fn column_kind_error(&self, name: &str, expected: &str) -> Result { + error::InvalidOtlpMetricInputSnafu { + reason: format!( + "OTLP metric column '{}' in table '{}' conflicts with the {} column", + name, self.name, expected + ), + } + .fail() + } + + fn has_field(&self) -> bool { + self.columns + .iter() + .any(|column| matches!(column.values, ScalarColumnValues::Field(_))) + } + + fn finish(self) -> Result { + let mut request_schema = Vec::with_capacity(self.columns.len()); + let mut fields = Vec::with_capacity(self.columns.len()); + let mut arrays = Vec::with_capacity(self.columns.len()); + for column in self.columns { + let (datatype, semantic_type, arrow_type, array): (_, _, _, ArrayRef) = + match column.values { + ScalarColumnValues::Tag(values) => { + let mut builder = StringBuilder::new(); + for value in values { + builder.append_option(value.as_deref()); + } + ( + ColumnDataType::String, + SemanticType::Tag, + ArrowDataType::Utf8, + Arc::new(builder.finish()), + ) + } + ScalarColumnValues::Timestamp(values) => ( + ColumnDataType::TimestampNanosecond, + SemanticType::Timestamp, + ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + Arc::new(TimestampNanosecondArray::from(values)), + ), + ScalarColumnValues::Field(values) => ( + ColumnDataType::Float64, + SemanticType::Field, + ArrowDataType::Float64, + Arc::new(Float64Array::from(values)), + ), + }; + request_schema.push(ColumnSchema { + column_name: column.name.clone(), + datatype: datatype as i32, + semantic_type: semantic_type as i32, + ..Default::default() + }); + fields.push(Field::new( + column.name, + arrow_type, + semantic_type != SemanticType::Timestamp, + )); + arrays.push(array); + } + let batch = RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), arrays) + .context(error::ArrowSnafu)?; + MetricRecordBatch::try_new(self.name, request_schema, batch) + } +} + +#[derive(Clone, Copy)] +enum ScalarColumnKind { + Tag, + Timestamp, + Field, +} + +impl ScalarColumnKind { + fn as_str(self) -> &'static str { + match self { + Self::Tag => "tag", + Self::Timestamp => "timestamp", + Self::Field => "field", + } + } +} + /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// /// See @@ -214,6 +427,139 @@ pub fn to_grpc_insert_requests( }) } +/// Converts requests containing only Gauge and Sum metrics directly into Arrow +/// record batches. Returns `None` when the whole request must use the general +/// OTLP conversion path. +pub fn try_to_scalar_record_batches( + request: &ExportMetricsServiceRequest, + metric_ctx: &OtlpMetricCtx, +) -> Result> { + if metric_ctx.is_legacy + || metric_ctx.resource_info + || request + .resource_metrics + .iter() + .flat_map(|resource| &resource.scope_metrics) + .flat_map(|scope| &scope.metrics) + .filter_map(|metric| metric.data.as_ref()) + .any(|data| !matches!(data, metric::Data::Gauge(_) | metric::Data::Sum(_))) + { + return Ok(None); + } + + let mut tables = HashMap::::new(); + let mut semantic_index = SemanticIndex::default(); + let mut outcome = MetricsIngestOutcome::default(); + let mut direct_ctx = metric_ctx.clone(); + let mut permission_tables = HashSet::new(); + + for resource in &request.resource_metrics { + let resource_attrs = resource.resource.as_ref().map(|resource| { + let mut attrs = resource.attributes.clone(); + process_resource_attrs(&mut attrs, &direct_ctx); + attrs + }); + let resource_tags = normalize_attributes( + resource_attrs.as_ref(), + AttributeType::Resource, + &direct_ctx, + )?; + + for scope in &resource.scope_metrics { + let scope_attrs = process_scope_attrs(scope, &direct_ctx); + let scope_tags = + normalize_attributes(scope_attrs.as_ref(), AttributeType::Scope, &direct_ctx)?; + + for metric in &scope.metrics { + let Some(data) = metric.data.as_ref() else { + continue; + }; + direct_ctx.set_metric_type(from_metric_type(data)); + let name = translate_metric_name( + metric, + &direct_ctx.metric_type, + direct_ctx.metric_translation_strategy, + ); + let (data_points, is_delta) = match data { + metric::Data::Gauge(gauge) => (gauge.data_points.as_slice(), false), + metric::Data::Sum(sum) => ( + sum.data_points.as_slice(), + matches!( + AggregationTemporality::try_from(sum.aggregation_temporality), + Ok(AggregationTemporality::Delta) + ), + ), + _ => return Ok(None), + }; + permission_tables.insert(name.clone()); + if data_points.is_empty() { + continue; + } + + add_accepted_data_points(&mut outcome, data_points.len())?; + let table = tables + .entry(name.clone()) + .or_insert_with(|| ScalarTableBuilder::new(name.clone())); + for data_point in data_points { + table.begin_row(); + for (key, value) in &resource_tags { + table.write_tag(key, value)?; + } + for (key, value) in &scope_tags { + table.write_tag(key, value)?; + } + for (key, value) in normalize_attributes( + Some(&data_point.attributes), + AttributeType::DataPoint, + &direct_ctx, + )? { + table.write_tag(&key, &value)?; + } + table.write_timestamp(data_point.time_unix_nano)?; + if is_delta { + table.write_tag( + OTLP_AGGREGATION_TEMPORALITY_LABEL, + GREPTIME_TEMPORALITY_DELTA, + )?; + } + if matches!(data, metric::Data::Sum(_)) + && has_no_recorded_value(data_point.flags) + { + table.write_field(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))?; + } else if let Some(value) = number_data_point_value(&data_point.value) { + table.write_field(value)?; + } + } + record_metric_semantics(&mut semantic_index, metric, &name, &direct_ctx); + } + } + } + + if tables.values().any(|table| !table.has_field()) { + return Ok(None); + } + let rows = tables.values().map(|table| table.rows).sum(); + let batches = tables + .into_values() + .map(ScalarTableBuilder::finish) + .collect::>>()?; + Ok(Some(ScalarMetricsConversion { + batches, + permission_tables: permission_tables.into_iter().collect(), + rows, + semantic_index, + outcome, + })) +} + +fn number_data_point_value(value: &Option) -> Option { + match value { + Some(number_data_point::Value::AsInt(value)) => Some(*value as f64), + Some(number_data_point::Value::AsDouble(value)) => Some(*value), + None => None, + } +} + fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { for request in &requests.inserts { let Some(rows) = &request.rows else { @@ -929,8 +1275,19 @@ fn write_attributes( attribute_type: AttributeType, metric_ctx: &OtlpMetricCtx, ) -> Result<()> { + let tags = normalize_attributes(attrs, attribute_type, metric_ctx)?; + row_writer::write_tags(writer, tags.into_iter(), row)?; + + Ok(()) +} + +fn normalize_attributes( + attrs: Option<&Vec>, + attribute_type: AttributeType, + metric_ctx: &OtlpMetricCtx, +) -> Result> { let Some(attrs) = attrs else { - return Ok(()); + return Ok(Vec::new()); }; let mut tags = Vec::with_capacity(attrs.len()); @@ -962,9 +1319,7 @@ fn write_attributes( } tags.push((key, value)); } - row_writer::write_tags(writer, tags.into_iter(), row)?; - - Ok(()) + Ok(tags) } fn write_timestamp(table: &mut TableData, row: &mut Vec, time_nano: u64) -> Result<()> { @@ -1573,6 +1928,97 @@ mod tests { } } + #[test] + fn test_scalar_record_batch_matches_row_conversion() { + set_default_prefix(None).unwrap(); + let mut request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "temperature", + ); + let metric::Data::Gauge(gauge) = request.resource_metrics[0].scope_metrics[0].metrics[0] + .data + .as_mut() + .unwrap() + else { + panic!("expected gauge") + }; + gauge.data_points[0].time_unix_nano = 1_000_123; + gauge.data_points[0] + .attributes + .push(keyvalue("sensor.id", "s-1")); + gauge.data_points.push(NumberDataPoint { + attributes: vec![keyvalue("zone", "west")], + time_unix_nano: 2_000_456, + value: Some(Value::AsDouble(2.5)), + ..Default::default() + }); + + let direct = try_to_scalar_record_batches(&request, &OtlpMetricCtx::default()) + .unwrap() + .unwrap(); + let rows = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) + .unwrap() + .requests; + let old = rows + .inserts + .iter() + .find(|insert| insert.table_name == "temperature") + .unwrap() + .rows + .as_ref() + .unwrap(); + let (_, request_schema, direct_batch) = direct + .batches + .into_iter() + .find(|batch| batch.table_name() == "temperature") + .unwrap() + .into_parts(); + + assert_eq!(old.schema, request_schema); + let old_batch = crate::prom_row_builder::rows_to_aligned_record_batch( + old, + direct_batch.schema().as_ref(), + ) + .unwrap() + .into_parts() + .0; + assert_eq!(old_batch, direct_batch); + assert_eq!(2, direct.rows); + assert_eq!(2, direct.outcome.accepted_data_points); + } + + #[test] + fn test_scalar_record_batch_falls_back_without_a_sample_field() { + let mut request = gauge_request(Vec::new(), "temperature"); + let metric::Data::Gauge(gauge) = request.resource_metrics[0].scope_metrics[0].metrics[0] + .data + .as_mut() + .unwrap() + else { + panic!("expected gauge"); + }; + gauge.data_points[0].value = None; + + assert!( + try_to_scalar_record_batches(&request, &OtlpMetricCtx::default()) + .unwrap() + .is_none() + ); + } + + #[test] + fn test_scalar_record_batch_falls_back_for_non_scalar_metric() { + let mut request = gauge_request(Vec::new(), "temperature"); + request.resource_metrics[0].scope_metrics[0].metrics[0].data = + Some(metric::Data::Summary(Summary::default())); + + assert!( + try_to_scalar_record_batches(&request, &OtlpMetricCtx::default()) + .unwrap() + .is_none() + ); + } + fn column_names(request: &RowInsertRequests, table: &str) -> Vec { request .inserts diff --git a/src/servers/src/otlp/metrics/tests/delta.rs b/src/servers/src/otlp/metrics/tests/delta.rs index cc3e5b04cf7..c77ffd14622 100644 --- a/src/servers/src/otlp/metrics/tests/delta.rs +++ b/src/servers/src/otlp/metrics/tests/delta.rs @@ -72,9 +72,11 @@ fn test_raw_delta_sum_identity_and_stale_marker() { })), ..Default::default() }; - let conversion = - to_grpc_insert_requests(metrics_request(vec![metric]), &mut OtlpMetricCtx::default()) - .unwrap(); + let request = metrics_request(vec![metric]); + let direct = try_to_scalar_record_batches(&request, &OtlpMetricCtx::default()) + .unwrap() + .unwrap(); + let conversion = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(3, conversion.outcome.accepted_data_points); assert_eq!(0, conversion.outcome.rejected_data_points); @@ -106,6 +108,33 @@ fn test_raw_delta_sum_identity_and_stale_marker() { ); } + let (_, direct_schema, direct_batch) = direct + .batches + .into_iter() + .find(|batch| batch.table_name() == "requests_total") + .unwrap() + .into_parts(); + assert_eq!(rows.schema, direct_schema); + let values = direct_batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(10.0, values.value(0)); + assert_eq!(20.5, values.value(1)); + assert_eq!(PROMETHEUS_STALE_NAN_BITS, values.value(2).to_bits()); + let temporalities = direct_batch + .column_by_name(OTLP_AGGREGATION_TEMPORALITY_LABEL) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + vec![GREPTIME_TEMPORALITY_DELTA; 3], + temporalities.iter().map(Option::unwrap).collect::>() + ); + for temporality in [ AggregationTemporality::Cumulative as i32, AggregationTemporality::Unspecified as i32, diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 40937127a73..7445c1f9ac7 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -61,7 +61,7 @@ use crate::metrics::{ }; use crate::prom_row_builder::{ build_metric_create_table_schema_from_proto, identify_missing_columns_from_proto, - rows_to_aligned_record_batch, + metric_record_batch_to_aligned, rows_to_aligned_record_batch, }; const PHYSICAL_TABLE_KEY: &str = "physical_table"; @@ -130,6 +130,7 @@ impl PendingRowsBatcherOptions { const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3; const MAX_CONCURRENT_FLOW_NOTIFICATIONS: usize = 8; +const MAX_CONCURRENT_TABLE_RESOLUTIONS: usize = 32; /// Batches protocol-neutral metric rows for ingestion. #[async_trait] @@ -141,6 +142,30 @@ pub trait MetricRowBatcher: Send + Sync { ctx: QueryContextRef, protocol: MetricRowBatchProtocol, ) -> Result; + + /// Submits rows, or returns them unchanged when the batcher cannot safely + /// use the requested physical table. + async fn submit_with_fallback( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { + self.submit(requests, ctx, protocol) + .await + .map(|_| MetricRowBatchSubmission::Submitted) + } + + /// Submits pre-built scalar metric Arrow batches. Implementations that do + /// not support this fast path request a fallback to ordinary row inserts. + async fn submit_record_batches_with_fallback( + &self, + _batches: Vec, + _ctx: QueryContextRef, + _protocol: MetricRowBatchProtocol, + ) -> Result { + Ok(MetricRecordBatchSubmission::UnmatchedPhysicalTable) + } } pub type MetricRowBatcherRef = Arc; @@ -160,6 +185,108 @@ impl MetricRowBatchProtocol { } } +/// Result of a metric row batch submission that supports Inserter fallback. +pub enum MetricRowBatchSubmission { + /// The batcher accepted and persisted the rows. + Submitted, + /// The logical tables do not belong to the requested physical table. + UnmatchedPhysicalTable(RowInsertRequests), +} + +/// Result of submitting pre-built scalar metric Arrow batches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MetricRecordBatchSubmission { + /// The batcher accepted and persisted the batches. + Submitted, + /// At least one logical table uses a different physical table. + UnmatchedPhysicalTable, +} + +/// A scalar metric table represented directly as Arrow arrays. +#[derive(Debug)] +pub struct MetricRecordBatch { + table_name: String, + request_schema: Vec, + batch: RecordBatch, +} + +impl MetricRecordBatch { + /// Creates a batch after validating its scalar metric schema and Arrow fields. + pub fn try_new( + table_name: String, + request_schema: Vec, + batch: RecordBatch, + ) -> Result { + ensure!( + !table_name.is_empty() && is_scalar_metric_schema(&request_schema), + error::InvalidParameterSnafu { + reason: format!("Invalid scalar metric record batch for table '{table_name}'") + } + ); + ensure!( + request_schema.len() == batch.num_columns(), + error::InvalidParameterSnafu { + reason: format!( + "Schema width mismatch for table '{}', expected {}, got {}", + table_name, + request_schema.len(), + batch.num_columns() + ) + } + ); + let batch_schema = batch.schema(); + for (column, field) in request_schema.iter().zip(batch_schema.fields()) { + let expected = match ColumnDataType::try_from(column.datatype) { + Ok(ColumnDataType::String) => ArrowDataType::Utf8, + Ok(ColumnDataType::Float64) => ArrowDataType::Float64, + Ok(ColumnDataType::TimestampMillisecond) => { + ArrowDataType::Timestamp(TimeUnit::Millisecond, None) + } + Ok(ColumnDataType::TimestampNanosecond) => { + ArrowDataType::Timestamp(TimeUnit::Nanosecond, None) + } + _ => { + return error::InvalidParameterSnafu { + reason: format!( + "Unsupported datatype {} in scalar metric table '{}'", + column.datatype, table_name + ), + } + .fail(); + } + }; + ensure!( + column.column_name == field.name().as_str() && expected == *field.data_type(), + error::InvalidParameterSnafu { + reason: format!( + "Arrow field mismatch for table '{}', expected '{}: {}', got '{}: {}'", + table_name, + column.column_name, + expected, + field.name(), + field.data_type() + ) + } + ); + } + Ok(Self { + table_name, + request_schema, + batch, + }) + } + + /// Returns the target logical table name. + pub fn table_name(&self) -> &str { + &self.table_name + } + + #[cfg(test)] + pub(crate) fn into_parts(self) -> (String, Vec, RecordBatch) { + (self.table_name, self.request_schema, self.batch) + } +} + struct WaitFlushResultTimer { started_at: Instant, protocol: MetricRowBatchProtocol, @@ -399,6 +526,14 @@ struct TableResolutionPlan { tables_to_alter: Vec<(String, Vec)>, } +enum TableResolutionOutcome { + Resolved(TableResolutionPlan), + PlacementMismatch { + table_name: String, + physical_table: String, + }, +} + struct PendingBatch { tables: HashMap, total_row_count: usize, @@ -584,20 +719,14 @@ pub fn metric_table_matches_physical_table(table_info: &TableInfo, physical_tabl .is_some_and(|table| table == physical_table) } -fn ensure_logical_table_matches_physical_table( - table_info: &TableInfo, - physical_table: &str, -) -> Result<()> { - ensure!( - metric_table_matches_physical_table(table_info, physical_table), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "Metric table '{}' does not belong to physical table '{}'", - table_info.name, physical_table - ), - } - ); - Ok(()) +fn logical_table_placement_error(table_name: &str, physical_table: &str) -> Error { + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Metric table '{}' does not belong to physical table '{}'", + table_name, physical_table + ), + } + .build() } /// Metric pending rows batcher. @@ -670,6 +799,40 @@ impl PendingRowsBatcher { requests: RowInsertRequests, ctx: QueryContextRef, protocol: MetricRowBatchProtocol, + ) -> Result { + let plan = match self.plan_request_tables(&requests, &ctx).await? { + TableResolutionOutcome::Resolved(plan) => plan, + TableResolutionOutcome::PlacementMismatch { + table_name, + physical_table, + } => return Err(logical_table_placement_error(&table_name, &physical_table)), + }; + self.submit_resolved(requests, ctx, protocol, plan).await + } + + async fn submit_with_fallback( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { + let plan = match self.plan_request_tables(&requests, &ctx).await? { + TableResolutionOutcome::Resolved(plan) => plan, + TableResolutionOutcome::PlacementMismatch { .. } => { + return Ok(MetricRowBatchSubmission::UnmatchedPhysicalTable(requests)); + } + }; + self.submit_resolved(requests, ctx, protocol, plan) + .await + .map(|_| MetricRowBatchSubmission::Submitted) + } + + async fn submit_resolved( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + plan: TableResolutionPlan, ) -> Result { METRIC_ROW_BATCH_SUBMISSIONS .with_label_values(&[protocol.as_str()]) @@ -678,8 +841,20 @@ impl PendingRowsBatcher { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_build_and_align"]) .start_timer(); - self.build_and_align_table_batches(requests, &ctx).await? + self.build_and_align_table_batches(requests, &ctx, plan) + .await? }; + self.submit_aligned(table_batches, total_rows, ctx, protocol) + .await + } + + async fn submit_aligned( + &self, + table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>, + total_rows: usize, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { if total_rows == 0 { return Ok(0); } @@ -757,6 +932,71 @@ impl PendingRowsBatcher { } } + async fn submit_record_batches_with_fallback( + &self, + batches: Vec, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { + let unique_tables = batches + .iter() + .filter(|batch| batch.batch.num_rows() > 0) + .map(|batch| (batch.table_name.as_str(), batch.request_schema.as_slice())) + .collect::>(); + let mut plan = match self + .plan_table_resolution( + ctx.current_catalog(), + &ctx.current_schema(), + &ctx, + &unique_tables, + ) + .await? + { + TableResolutionOutcome::Resolved(plan) => plan, + TableResolutionOutcome::PlacementMismatch { .. } => { + return Ok(MetricRecordBatchSubmission::UnmatchedPhysicalTable); + } + }; + + METRIC_ROW_BATCH_SUBMISSIONS + .with_label_values(&[protocol.as_str()]) + .inc(); + let catalog = ctx.current_catalog(); + let schema = ctx.current_schema(); + self.create_missing_tables_and_refresh_schemas( + catalog, + &schema, + &ctx, + &unique_tables, + &mut plan, + ) + .await?; + self.alter_tables_and_refresh_schemas(catalog, &schema, &ctx, &mut plan) + .await?; + + let total_rows = batches.iter().map(|batch| batch.batch.num_rows()).sum(); + let mut aligned = Vec::with_capacity(batches.len()); + for batch in batches { + if batch.batch.num_rows() == 0 { + continue; + } + let (region_schema, table_id) = plan + .region_schemas + .get(&batch.table_name) + .cloned() + .with_context(|| error::UnexpectedResultSnafu { + reason: format!("Region schema not resolved for table: {}", batch.table_name), + })?; + let record_batch = + metric_record_batch_to_aligned(&batch.batch, region_schema.as_ref())?; + aligned.push((batch.table_name, table_id, record_batch)); + } + + self.submit_aligned(aligned, total_rows, ctx, protocol) + .await + .map(|_| MetricRecordBatchSubmission::Submitted) + } + /// Converts proto `RowInsertRequests` directly into aligned `RecordBatch`es /// in a single pass, handling table creation, schema alteration, column /// renaming, reordering, and null-filling without building intermediate @@ -765,6 +1005,7 @@ impl PendingRowsBatcher { &self, requests: RowInsertRequests, ctx: &QueryContextRef, + mut plan: TableResolutionPlan, ) -> Result<(Vec<(String, u32, RecordBatchWithTsIdx)>, usize)> { let catalog = ctx.current_catalog().to_string(); let schema = ctx.current_schema(); @@ -774,16 +1015,14 @@ impl PendingRowsBatcher { return Ok((Vec::new(), 0)); } - let unique_tables = Self::collect_unique_table_schemas(&table_rows)?; - let mut plan = self - .plan_table_resolution(&catalog, &schema, ctx, &unique_tables) - .await?; - self.create_missing_tables_and_refresh_schemas( &catalog, &schema, ctx, - &table_rows, + &table_rows + .iter() + .map(|(name, rows)| (name.as_str(), rows.schema.as_slice())) + .collect::>(), &mut plan, ) .await?; @@ -820,20 +1059,27 @@ impl PendingRowsBatcher { /// Returns unique `(table_name, proto_schema)` pairs while keeping the /// first-seen schema for duplicate table names. fn collect_unique_table_schemas( - table_rows: &[(String, Rows)], + requests: &RowInsertRequests, ) -> Result> { - let mut unique_tables: Vec<(&str, &[ColumnSchema])> = Vec::with_capacity(table_rows.len()); + let mut unique_tables: Vec<(&str, &[ColumnSchema])> = + Vec::with_capacity(requests.inserts.len()); let mut seen = HashSet::new(); - for (table_name, rows) in table_rows { - if seen.insert(table_name.as_str()) { - unique_tables.push((table_name.as_str(), &rows.schema)); + for request in &requests.inserts { + let Some(rows) = &request.rows else { + continue; + }; + if rows.rows.is_empty() { + continue; + } + + if seen.insert(request.table_name.as_str()) { + unique_tables.push((request.table_name.as_str(), &rows.schema)); } else { - // table_rows should group rows by table name. return error::InvalidPromRemoteRequestSnafu { msg: format!( "Found duplicated table name in RowInsertRequest: {}", - table_name + request.table_name ), } .fail(); @@ -843,6 +1089,21 @@ impl PendingRowsBatcher { Ok(unique_tables) } + async fn plan_request_tables( + &self, + requests: &RowInsertRequests, + ctx: &QueryContextRef, + ) -> Result { + let unique_tables = Self::collect_unique_table_schemas(requests)?; + self.plan_table_resolution( + ctx.current_catalog(), + &ctx.current_schema(), + ctx, + &unique_tables, + ) + .await + } + /// Resolves table metadata and classifies each table into existing, /// to-create, and to-alter groups used by subsequent DDL steps. async fn plan_table_resolution( @@ -851,7 +1112,7 @@ impl PendingRowsBatcher { schema: &str, ctx: &QueryContextRef, unique_tables: &[(&str, &[ColumnSchema])], - ) -> Result { + ) -> Result { let physical_table = ctx .extension(PHYSICAL_TABLE_KEY) .unwrap_or(GREPTIME_PHYSICAL_TABLE); @@ -865,11 +1126,17 @@ impl PendingRowsBatcher { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["align_resolve_table"]) .start_timer(); - futures::future::join_all(unique_tables.iter().map(|(table_name, _)| { - self.catalog_manager - .table(catalog, schema, table_name, Some(ctx.as_ref())) - })) - .await + let mut resolved_tables = Vec::with_capacity(unique_tables.len()); + for tables in unique_tables.chunks(MAX_CONCURRENT_TABLE_RESOLUTIONS) { + resolved_tables.extend( + futures::future::join_all(tables.iter().map(|(table_name, _)| { + self.catalog_manager + .table(catalog, schema, table_name, Some(ctx.as_ref())) + })) + .await, + ); + } + resolved_tables }; for ((table_name, rows_schema), table_result) in unique_tables.iter().zip(resolved_tables) { @@ -877,8 +1144,13 @@ impl PendingRowsBatcher { if let Some(table) = table { let table_info = table.table_info(); - if self.with_metric_engine { - ensure_logical_table_matches_physical_table(&table_info, physical_table)?; + if self.with_metric_engine + && !metric_table_matches_physical_table(&table_info, physical_table) + { + return Ok(TableResolutionOutcome::PlacementMismatch { + table_name: (*table_name).to_string(), + physical_table: physical_table.to_string(), + }); } let table_id = table_info.ident.table_id; let region_schema = table_info.meta.schema.arrow_schema().clone(); @@ -907,7 +1179,7 @@ impl PendingRowsBatcher { } } - Ok(plan) + Ok(TableResolutionOutcome::Resolved(plan)) } /// Batch-creates missing tables, refreshes their schema metadata, and @@ -917,7 +1189,7 @@ impl PendingRowsBatcher { catalog: &str, schema: &str, ctx: &QueryContextRef, - table_rows: &[(String, Rows)], + table_schemas: &[(&str, &[ColumnSchema])], plan: &mut TableResolutionPlan, ) -> Result<()> { if plan.tables_to_create.is_empty() { @@ -970,7 +1242,12 @@ impl PendingRowsBatcher { let physical_table = ctx .extension(PHYSICAL_TABLE_KEY) .unwrap_or(GREPTIME_PHYSICAL_TABLE); - ensure_logical_table_matches_physical_table(&table_info, physical_table)?; + if !metric_table_matches_physical_table(&table_info, physical_table) { + return Err(logical_table_placement_error( + &table_info.name, + physical_table, + )); + } } let table_id = table_info.ident.table_id; let region_schema = table_info.meta.schema.arrow_schema().clone(); @@ -978,7 +1255,7 @@ impl PendingRowsBatcher { .insert(table_name.clone(), (region_schema, table_id)); } - Self::enqueue_alter_for_new_tables(table_rows, plan)?; + Self::enqueue_alter_for_new_tables(table_schemas, plan)?; Ok(()) } @@ -986,7 +1263,7 @@ impl PendingRowsBatcher { /// For newly created tables, re-checks all row schemas and appends alter /// operations when additional tag columns are still missing. fn enqueue_alter_for_new_tables( - table_rows: &[(String, Rows)], + table_schemas: &[(&str, &[ColumnSchema])], plan: &mut TableResolutionPlan, ) -> Result<()> { let created_tables: HashSet<&str> = plan @@ -995,16 +1272,17 @@ impl PendingRowsBatcher { .map(|(table_name, _)| table_name.as_str()) .collect(); - for (table_name, rows) in table_rows { - if !created_tables.contains(table_name.as_str()) { + for (table_name, request_schema) in table_schemas { + if !created_tables.contains(table_name) { continue; } - let Some((region_schema, _)) = plan.region_schemas.get(table_name) else { + let Some((region_schema, _)) = plan.region_schemas.get(*table_name) else { continue; }; - let missing_columns = identify_missing_columns_from_proto(&rows.schema, region_schema)?; + let missing_columns = + identify_missing_columns_from_proto(request_schema, region_schema)?; if missing_columns.is_empty() || plan .tables_to_alter @@ -1015,7 +1293,7 @@ impl PendingRowsBatcher { } plan.tables_to_alter - .push((table_name.clone(), missing_columns)); + .push(((*table_name).to_string(), missing_columns)); } Ok(()) @@ -1173,6 +1451,24 @@ impl MetricRowBatcher for PendingRowsBatcher { ) -> Result { PendingRowsBatcher::submit(self, requests, ctx, protocol).await } + + async fn submit_with_fallback( + &self, + requests: RowInsertRequests, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { + PendingRowsBatcher::submit_with_fallback(self, requests, ctx, protocol).await + } + + async fn submit_record_batches_with_fallback( + &self, + batches: Vec, + ctx: QueryContextRef, + protocol: MetricRowBatchProtocol, + ) -> Result { + PendingRowsBatcher::submit_record_batches_with_fallback(self, batches, ctx, protocol).await + } } impl Drop for PendingRowsBatcher { @@ -2151,7 +2447,8 @@ mod tests { StructValue, Value, }; use arrow::array::{ - BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray, TimestampNanosecondArray, + BinaryArray, BooleanArray, Float64Array, StringArray, TimestampMillisecondArray, + TimestampNanosecondArray, }; use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use arrow::record_batch::RecordBatch; @@ -2197,7 +2494,8 @@ mod tests { use super::{ BatchKey, CatalogManagerPhysicalFlushAdapter, Error, FlushBatch, FlushRegionWrite, - FlushWaiter, MetricRowBatchProtocol, NodeManagerPhysicalFlushAdapter, + FlushWaiter, MetricRecordBatch, MetricRecordBatchSubmission, MetricRowBatchProtocol, + MetricRowBatchSubmission, NodeManagerPhysicalFlushAdapter, PartitionManagerPhysicalFlushAdapter, PendingBatch, PendingRowsBatchMode, PendingRowsBatcher, PendingRowsSchemaAlterer, PendingWorker, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, @@ -3031,6 +3329,34 @@ mod tests { requests } + fn synchronous_metric_record_batch() -> MetricRecordBatch { + let request_schema = synchronous_submit_request().inserts[0] + .rows + .as_ref() + .unwrap() + .schema + .clone(); + let schema = ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + false, + ), + Field::new("value", ArrowDataType::Float64, true), + Field::new("tag1", ArrowDataType::Utf8, true), + ]); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000])), + Arc::new(Float64Array::from(vec![1.0])), + Arc::new(StringArray::from(vec!["host-1"])), + ], + ) + .unwrap(); + MetricRecordBatch::try_new("cpu".to_string(), request_schema, batch).unwrap() + } + fn test_synchronous_submit_batcher( node_requester: Arc, ) -> Arc { @@ -3144,6 +3470,77 @@ mod tests { assert_eq!(0, writes.load(Ordering::SeqCst)); } + #[tokio::test] + async fn test_submit_with_fallback_returns_mismatched_request_without_writing() { + let writes = Arc::new(AtomicUsize::new(0)); + let batcher = test_synchronous_submit_batcher_with_physical_table( + Arc::new(MockFlushNodeRequester { + writes: writes.clone(), + fail: false, + }), + "other_physical_table", + ); + + let submission = batcher + .submit_with_fallback( + synchronous_submit_request(), + session::context::QueryContext::arc(), + MetricRowBatchProtocol::Otlp, + ) + .await + .unwrap(); + + let MetricRowBatchSubmission::UnmatchedPhysicalTable(requests) = submission else { + panic!("expected placement mismatch to fall back") + }; + assert_eq!("cpu", requests.inserts[0].table_name); + assert_eq!(0, writes.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_submit_record_batch_writes_without_proto_rows() { + let (batcher, writes) = test_successful_synchronous_submit_batcher(); + + let submission = batcher + .submit_record_batches_with_fallback( + vec![synchronous_metric_record_batch()], + session::context::QueryContext::arc(), + MetricRowBatchProtocol::Otlp, + ) + .await + .unwrap(); + + assert_eq!(MetricRecordBatchSubmission::Submitted, submission); + assert_eq!(1, writes.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn test_submit_record_batch_returns_placement_mismatch() { + let writes = Arc::new(AtomicUsize::new(0)); + let batcher = test_synchronous_submit_batcher_with_physical_table( + Arc::new(MockFlushNodeRequester { + writes: writes.clone(), + fail: false, + }), + "other_physical_table", + ); + + let submission = batcher + .submit_record_batches_with_fallback( + vec![synchronous_metric_record_batch()], + session::context::QueryContext::arc(), + MetricRowBatchProtocol::Otlp, + ) + .await + .unwrap(); + + assert_eq!( + MetricRecordBatchSubmission::UnmatchedPhysicalTable, + submission + ); + assert_eq!(0, writes.load(Ordering::SeqCst)); + } + static METRIC_ASSERTION_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); #[tokio::test] diff --git a/src/servers/src/prom_row_builder.rs b/src/servers/src/prom_row_builder.rs index c7596f94b21..4e364135ed2 100644 --- a/src/servers/src/prom_row_builder.rs +++ b/src/servers/src/prom_row_builder.rs @@ -199,6 +199,75 @@ pub(crate) fn rows_to_aligned_record_batch( RecordBatchWithTsIdx::try_new(batch, timestamp_index) } +/// Aligns an Arrow scalar metric batch to an existing logical table schema. +/// +/// Source arrays are reused when possible; only missing target columns allocate +/// null arrays. Timestamp and value columns are matched by semantic shape so +/// existing tables may use non-default names. +pub(crate) fn metric_record_batch_to_aligned( + source: &RecordBatch, + target_schema: &ArrowSchema, +) -> Result { + let source_schema = source.schema(); + let (target_ts_name, target_field_name, _) = unzip_logical_region_schema(target_schema)?; + let timestamp_index = target_schema + .column_with_name(&target_ts_name) + .map(|(index, _)| index) + .with_context(|| error::UnexpectedResultSnafu { + reason: format!( + "Failed to resolve timestamp column '{}' in target schema", + target_ts_name + ), + })?; + let target_timestamp_type = target_schema.field(timestamp_index).data_type(); + let mut source_map = HashMap::with_capacity(source.num_columns()); + + for (index, field) in source_schema.fields().iter().enumerate() { + let effective_name = match field.data_type() { + ArrowDataType::Float64 => target_field_name.as_str(), + ArrowDataType::Timestamp(_, _) => { + ensure!( + field.data_type() == target_timestamp_type, + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Metric row timestamp type mismatch, expected {}, actual {}", + target_timestamp_type, + field.data_type() + ) + } + ); + target_ts_name.as_str() + } + ArrowDataType::Utf8 => field.name(), + other => { + return error::InvalidPromRemoteRequestSnafu { + msg: format!( + "Unexpected metric record batch type {}, field name: {}", + other, + field.name() + ), + } + .fail(); + } + }; + source_map.insert(effective_name, index); + } + + let columns = target_schema + .fields() + .iter() + .map(|field| { + source_map + .get(field.name().as_str()) + .map(|index| source.column(*index).clone()) + .unwrap_or_else(|| new_null_array(field.data_type(), source.num_rows())) + }) + .collect(); + let batch = RecordBatch::try_new(Arc::new(target_schema.clone()), columns) + .context(error::ArrowSnafu)?; + RecordBatchWithTsIdx::try_new(batch, timestamp_index) +} + /// Identify tag columns in the proto `rows_schema` that are absent from the /// target region schema, without building an intermediate `RecordBatch`. pub(crate) fn identify_missing_columns_from_proto( @@ -330,18 +399,64 @@ fn build_arrow_array( #[cfg(test)] mod tests { + use std::sync::Arc; + use api::v1::value::ValueData; use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value}; use arrow::array::{ - Array, Float64Array, StringArray, TimestampMillisecondArray, TimestampNanosecondArray, + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, + TimestampNanosecondArray, }; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use arrow::record_batch::RecordBatch; use super::{ build_metric_create_table_schema_from_proto, identify_missing_columns_from_proto, - rows_to_aligned_record_batch, + metric_record_batch_to_aligned, rows_to_aligned_record_batch, }; + #[test] + fn test_metric_record_batch_aligns_without_proto_materialization() { + let source = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new( + "greptime_timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("host", DataType::Utf8, true), + Field::new("greptime_value", DataType::Float64, true), + ])), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![2_000_123])) as ArrayRef, + Arc::new(StringArray::from(vec!["h1"])) as ArrayRef, + Arc::new(Float64Array::from(vec![42.0])) as ArrayRef, + ], + ) + .unwrap(); + let target = ArrowSchema::new(vec![ + Field::new("region", DataType::Utf8, true), + Field::new("my_value", DataType::Float64, true), + Field::new( + "my_ts", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("host", DataType::Utf8, true), + ]); + + let (aligned, timestamp_index) = metric_record_batch_to_aligned(&source, &target) + .unwrap() + .into_parts(); + + assert_eq!(2, timestamp_index); + assert_eq!(aligned.schema().as_ref(), &target); + assert_eq!(1, aligned.column(0).null_count()); + assert!(Arc::ptr_eq(source.column(2), aligned.column(1))); + assert!(Arc::ptr_eq(source.column(0), aligned.column(2))); + assert!(Arc::ptr_eq(source.column(1), aligned.column(3))); + } + #[test] fn test_rows_to_aligned_record_batch_renames_and_reorders() { let rows = Rows { diff --git a/tests-integration/src/otlp.rs b/tests-integration/src/otlp.rs index 4f458a6abc2..f25faa159b1 100644 --- a/tests-integration/src/otlp.rs +++ b/tests-integration/src/otlp.rs @@ -65,7 +65,7 @@ mod test { let mut output = instance .do_query( "CREATE TABLE fixed_delta_total (\ - \"stream\" STRING, greptime_timestamp TIMESTAMP(3) NOT NULL, \ + \"stream\" STRING, greptime_timestamp TIMESTAMP(9) NOT NULL, \ greptime_value DOUBLE, TIME INDEX (greptime_timestamp), \ PRIMARY KEY (\"stream\")) ENGINE=mito", ctx.clone(), @@ -123,7 +123,7 @@ mod test { let mut output = instance .do_query( "CREATE TABLE raw_delta_mito_total (\ - \"stream\" STRING, greptime_timestamp TIMESTAMP(3) NOT NULL, \ + \"stream\" STRING, greptime_timestamp TIMESTAMP(9) NOT NULL, \ greptime_value DOUBLE, TIME INDEX (greptime_timestamp), \ PRIMARY KEY (\"stream\")) ENGINE=mito", ctx.clone(), @@ -404,12 +404,12 @@ mod test { assert_eq!( recordbatches.pretty_print().unwrap(), "\ -+----------------+---------------------+----------------+ -| container_name | greptime_timestamp | greptime_value | -+----------------+---------------------+----------------+ -| testserver | 1970-01-01T00:00:00 | 105.0 | -| testsevrer | 1970-01-01T00:00:00 | 100.0 | -+----------------+---------------------+----------------+", ++----------------+-------------------------------+----------------+ +| container_name | greptime_timestamp | greptime_value | ++----------------+-------------------------------+----------------+ +| testsevrer | 1970-01-01T00:00:00.000000100 | 100.0 | +| testserver | 1970-01-01T00:00:00.000000105 | 105.0 | ++----------------+-------------------------------+----------------+", ); let mut output = instance @@ -449,11 +449,11 @@ mod test { assert_eq!( recordbatches.pretty_print().unwrap(), "\ -+------------+---------------------+----------------+ -| host | greptime_timestamp | greptime_value | -+------------+---------------------+----------------+ -| testserver | 1970-01-01T00:00:00 | 51.0 | -+------------+---------------------+----------------+", ++------------+-------------------------------+----------------+ +| host | greptime_timestamp | greptime_value | ++------------+-------------------------------+----------------+ +| testserver | 1970-01-01T00:00:00.000000100 | 51.0 | ++------------+-------------------------------+----------------+", ); let mut output = instance @@ -470,11 +470,11 @@ mod test { assert_eq!( recordbatches.pretty_print().unwrap(), "\ -+------------+---------------------+----------------+ -| host | greptime_timestamp | greptime_value | -+------------+---------------------+----------------+ -| testserver | 1970-01-01T00:00:00 | 4.0 | -+------------+---------------------+----------------+", ++------------+-------------------------------+----------------+ +| host | greptime_timestamp | greptime_value | ++------------+-------------------------------+----------------+ +| testserver | 1970-01-01T00:00:00.000000100 | 4.0 | ++------------+-------------------------------+----------------+", ); }