perf: convert scalar OTLP metrics directly to Arrow

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-09-09 17:01:13 +08:00
parent 4d285bf8ee
commit 963be5ceb4
7 changed files with 1376 additions and 190 deletions
+302 -112
View File
@@ -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<bool> {
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<bool> {
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<MetricRowBatcherRef>,
ctx: QueryContextRef,
) -> ServerResult<MetricsIngestOutcome> {
self.plugins
.get::<PermissionCheckerRef>()
.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::<PermissionCheckerRef>()
.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::<Vec<_>>();
// 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<Vec<RowInsertRequests>>,
record_batch_tables: Mutex<Vec<Vec<String>>>,
protocols: Mutex<Vec<MetricRowBatchProtocol>>,
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<MetricRowBatchSubmission> {
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<MetricRecordBatch>,
_ctx: QueryContextRef,
protocol: MetricRowBatchProtocol,
) -> ServerResult<MetricRecordBatchSubmission> {
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<Vec<PermissionTableTargets>>,
}
impl PermissionChecker for CountingPermissionChecker {
@@ -553,9 +672,10 @@ mod tests {
&self,
_user_info: UserInfoRef,
_req: PermissionReq,
_targets: PermissionTableTargets,
targets: PermissionTableTargets,
) -> auth::error::Result<PermissionResp> {
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::<PermissionCheckerRef>(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::<HashSet<_>>();
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::<crate::error::Error>()
.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()
+9
View File
@@ -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!(
+452 -6
View File
@@ -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<MetricRecordBatch>,
/// All emitted logical tables, including metrics without data points.
pub permission_tables: Vec<String>,
/// 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<Option<String>>),
Timestamp(Vec<Option<i64>>),
Field(Vec<Option<f64>>),
}
struct ScalarColumn {
name: String,
values: ScalarColumnValues,
}
struct ScalarTableBuilder {
name: String,
columns: Vec<ScalarColumn>,
column_indexes: HashMap<String, usize>,
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<usize> {
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<T>(&self, name: &str, expected: &str) -> Result<T> {
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<MetricRecordBatch> {
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<Option<ScalarMetricsConversion>> {
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::<String, ScalarTableBuilder>::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::<Result<Vec<_>>>()?;
Ok(Some(ScalarMetricsConversion {
batches,
permission_tables: permission_tables.into_iter().collect(),
rows,
semantic_index,
outcome,
}))
}
fn number_data_point_value(value: &Option<number_data_point::Value>) -> Option<f64> {
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<KeyValue>>,
attribute_type: AttributeType,
metric_ctx: &OtlpMetricCtx,
) -> Result<Vec<(String, String)>> {
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<Value>, 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<String> {
request
.inserts
+32 -3
View File
@@ -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::<Float64Array>()
.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::<arrow::array::StringArray>()
.unwrap();
assert_eq!(
vec![GREPTIME_TEMPORALITY_DELTA; 3],
temporalities.iter().map(Option::unwrap).collect::<Vec<_>>()
);
for temporality in [
AggregationTemporality::Cumulative as i32,
AggregationTemporality::Unspecified as i32,
+446 -49
View File
@@ -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<u64>;
/// 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<MetricRowBatchSubmission> {
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<MetricRecordBatch>,
_ctx: QueryContextRef,
_protocol: MetricRowBatchProtocol,
) -> Result<MetricRecordBatchSubmission> {
Ok(MetricRecordBatchSubmission::UnmatchedPhysicalTable)
}
}
pub type MetricRowBatcherRef = Arc<dyn MetricRowBatcher>;
@@ -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<ColumnSchema>,
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<ColumnSchema>,
batch: RecordBatch,
) -> Result<Self> {
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<ColumnSchema>, 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<String>)>,
}
enum TableResolutionOutcome {
Resolved(TableResolutionPlan),
PlacementMismatch {
table_name: String,
physical_table: String,
},
}
struct PendingBatch {
tables: HashMap<TableId, TableBatch>,
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<u64> {
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<MetricRowBatchSubmission> {
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<u64> {
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<u64> {
if total_rows == 0 {
return Ok(0);
}
@@ -757,6 +932,71 @@ impl PendingRowsBatcher {
}
}
async fn submit_record_batches_with_fallback(
&self,
batches: Vec<MetricRecordBatch>,
ctx: QueryContextRef,
protocol: MetricRowBatchProtocol,
) -> Result<MetricRecordBatchSubmission> {
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::<Vec<_>>();
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::<Vec<_>>(),
&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<Vec<(&str, &[ColumnSchema])>> {
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<TableResolutionOutcome> {
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<TableResolutionPlan> {
) -> Result<TableResolutionOutcome> {
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<u64> {
PendingRowsBatcher::submit(self, requests, ctx, protocol).await
}
async fn submit_with_fallback(
&self,
requests: RowInsertRequests,
ctx: QueryContextRef,
protocol: MetricRowBatchProtocol,
) -> Result<MetricRowBatchSubmission> {
PendingRowsBatcher::submit_with_fallback(self, requests, ctx, protocol).await
}
async fn submit_record_batches_with_fallback(
&self,
batches: Vec<MetricRecordBatch>,
ctx: QueryContextRef,
protocol: MetricRowBatchProtocol,
) -> Result<MetricRecordBatchSubmission> {
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<dyn PhysicalFlushNodeRequester>,
) -> Arc<PendingRowsBatcher> {
@@ -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]
+117 -2
View File
@@ -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<RecordBatchWithTsIdx> {
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 {
+18 -18
View File
@@ -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 |
+------------+-------------------------------+----------------+",
);
}