diff --git a/src/pipeline/benches/processor.rs b/src/pipeline/benches/processor.rs index e088c89c6b..dd0c3f1160 100644 --- a/src/pipeline/benches/processor.rs +++ b/src/pipeline/benches/processor.rs @@ -17,7 +17,10 @@ use std::sync::Arc; use criterion::{Criterion, criterion_group, criterion_main}; use pipeline::error::Result; -use pipeline::{Content, Pipeline, PipelineContext, SchemaInfo, parse, setup_pipeline}; +use pipeline::{ + Content, GreptimePipelineParams, Pipeline, PipelineContext, PipelineDefinition, SchemaInfo, + identity_pipeline, parse, setup_pipeline, +}; use serde_json::Deserializer; use vrl::value::Value as VrlValue; @@ -40,6 +43,14 @@ fn processor_mut( Ok(result) } +fn identity_pipeline_rows( + pipeline_ctx: &PipelineContext<'_>, + input_values: Vec, +) -> Result { + let rows = identity_pipeline(input_values, None, pipeline_ctx)?; + Ok(rows.values().map(|rows| rows.rows.len()).sum()) +} + fn prepare_pipeline() -> Pipeline { let pipeline_yaml = r#" --- @@ -293,6 +304,42 @@ fn criterion_benchmark(c: &mut Criterion) { }); group.finish(); + let identity_input_value = (0..128) + .map(|i| { + serde_json::json!({ + "service": "frontend", + "host": format!("host-{i}"), + "status": 200, + "message": { + "path": "/api/v1/ingest", + "latency_ms": i, + }, + "labels": ["http", "pipeline"], + }) + .into() + }) + .collect::>(); + let identity_pipeline_def = PipelineDefinition::GreptimeIdentityPipeline(None); + let identity_pipeline_param = GreptimePipelineParams::default(); + let identity_pipeline_ctx = PipelineContext::new( + &identity_pipeline_def, + &identity_pipeline_param, + session::context::Channel::Unknown, + ); + + let mut group = c.benchmark_group("identity pipeline"); + group.sample_size(50); + group.bench_function("processor mut", |b| { + b.iter(|| { + identity_pipeline_rows( + black_box(&identity_pipeline_ctx), + black_box(identity_input_value.clone()), + ) + .unwrap(); + }) + }); + group.finish(); + let vrl_input_value = (0..128) .map(|i| { serde_json::json!({ diff --git a/src/pipeline/src/etl/transform/transformer/greptime.rs b/src/pipeline/src/etl/transform/transformer/greptime.rs index a1941aca05..e48026840f 100644 --- a/src/pipeline/src/etl/transform/transformer/greptime.rs +++ b/src/pipeline/src/etl/transform/transformer/greptime.rs @@ -772,6 +772,7 @@ fn vrl_value_to_jsonb_value<'a>(value: &'a VrlValue) -> jsonb::Value<'a> { fn identity_pipeline_inner( pipeline_maps: Vec, pipeline_ctx: &PipelineContext<'_>, + max_nested_levels: usize, ) -> Result<(SchemaInfo, HashMap>)> { let skip_error = pipeline_ctx.pipeline_param.skip_error(); let mut schema_info = SchemaInfo::default(); @@ -801,7 +802,9 @@ fn identity_pipeline_inner( let mut opt_map = HashMap::new(); let len = pipeline_maps.len(); - for mut pipeline_map in pipeline_maps { + for pipeline_map in pipeline_maps { + let mut pipeline_map = + unwrap_or_continue_if_err!(flatten_object(pipeline_map, max_nested_levels), skip_error); let opt = unwrap_or_continue_if_err!( ContextOpt::from_pipeline_map_to_opt(&mut pipeline_map), skip_error @@ -826,10 +829,8 @@ fn identity_pipeline_inner( column_count, row.values.len() ); - let diff = column_count - row.values.len(); - for _ in 0..diff { - row.values.push(GreptimeValue { value_data: None }); - } + row.values + .resize(column_count, GreptimeValue { value_data: None }); } } @@ -849,40 +850,31 @@ pub fn identity_pipeline( table: Option>, pipeline_ctx: &PipelineContext<'_>, ) -> Result> { - let skip_error = pipeline_ctx.pipeline_param.skip_error(); let max_nested_levels = pipeline_ctx.pipeline_param.max_nested_levels(); - // Always flatten JSON objects and stringify arrays - let mut input = Vec::with_capacity(array.len()); - for item in array.into_iter() { - let result = - unwrap_or_continue_if_err!(flatten_object(item, max_nested_levels), skip_error); - input.push(result); - } - identity_pipeline_inner(input, pipeline_ctx).and_then(|(mut schema, opt_map)| { - if let Some(table) = table { - let table_info = table.table_info(); - for tag_name in table_info.meta.row_key_column_names() { - if let Some(index) = schema.index.get(tag_name) { - schema.schema[*index].semantic_type = SemanticType::Tag; - } + let (mut schema, opt_map) = identity_pipeline_inner(array, pipeline_ctx, max_nested_levels)?; + if let Some(table) = table { + let table_info = table.table_info(); + for tag_name in table_info.meta.row_key_column_names() { + if let Some(index) = schema.index.get(tag_name) { + schema.schema[*index].semantic_type = SemanticType::Tag; } } + } - let column_schemas = schema.column_schemas()?; - Ok(opt_map - .into_iter() - .map(|(opt, rows)| { - ( - opt, - Rows { - schema: column_schemas.clone(), - rows, - }, - ) - }) - .collect::>()) - }) + let column_schemas = schema.column_schemas()?; + Ok(opt_map + .into_iter() + .map(|(opt, rows)| { + ( + opt, + Rows { + schema: column_schemas.clone(), + rows, + }, + ) + }) + .collect::>()) } /// Consumes the JSON object and consumes it into a single-level object. @@ -1100,23 +1092,26 @@ mod tests { ]; let tag_column_names = ["name".to_string(), "address".to_string()]; - let rows = - identity_pipeline_inner(array.iter().map(|v| v.into()).collect(), &pipeline_ctx) - .map(|(mut schema, mut rows)| { - for name in tag_column_names { - if let Some(index) = schema.index.get(&name) { - schema.schema[*index].semantic_type = SemanticType::Tag; - } - } + let rows = identity_pipeline_inner( + array.iter().map(|v| v.into()).collect(), + &pipeline_ctx, + pipeline_ctx.pipeline_param.max_nested_levels(), + ) + .map(|(mut schema, mut rows)| { + for name in tag_column_names { + if let Some(index) = schema.index.get(&name) { + schema.schema[*index].semantic_type = SemanticType::Tag; + } + } - assert!(rows.len() == 1); - let rows = rows.remove(&ContextOpt::default()).unwrap(); + assert!(rows.len() == 1); + let rows = rows.remove(&ContextOpt::default()).unwrap(); - Rows { - schema: schema.column_schemas().unwrap(), - rows, - } - }); + Rows { + schema: schema.column_schemas().unwrap(), + rows, + } + }); assert!(rows.is_ok()); let rows = rows.unwrap(); @@ -1226,6 +1221,46 @@ mod tests { } } + #[test] + fn test_identity_pipeline_skip_error_flattens_valid_rows() { + let params = GreptimePipelineParams::from_map(ahash::HashMap::from_iter([( + "skip_error".to_string(), + "true".to_string(), + )])); + let pipeline_def = PipelineDefinition::GreptimeIdentityPipeline(None); + let pipeline_ctx = PipelineContext::new(&pipeline_def, ¶ms, Channel::Unknown); + let array = vec![ + serde_json::json!({ + "service": "frontend", + "nested": { + "status": 200, + "path": "/v1/ingest" + }, + "labels": ["pipeline", "identity"] + }) + .into(), + VrlValue::Bytes("invalid_string".into()), + serde_json::json!({ + "service": "frontend", + "nested": { + "status": 201, + "path": "/v1/ingest" + }, + "labels": ["pipeline", "identity"] + }) + .into(), + ]; + + let mut rows_by_opt = identity_pipeline(array, None, &pipeline_ctx).unwrap(); + let rows = rows_by_opt.remove(&ContextOpt::default()).unwrap(); + + assert_eq!(rows.rows.len(), 2); + assert_eq!(rows.schema.len(), rows.rows[0].values.len()); + assert!(rows.schema.iter().any(|s| s.column_name == "nested.status")); + assert!(rows.schema.iter().any(|s| s.column_name == "nested.path")); + assert!(rows.schema.iter().any(|s| s.column_name == "labels")); + } + use ahash::HashMap as AHashMap; #[test] fn test_values_to_rows_skip_error_handling() { diff --git a/src/servers/src/pipeline.rs b/src/servers/src/pipeline.rs index 9de2d63b97..7f6fa20f1f 100644 --- a/src/servers/src/pipeline.rs +++ b/src/servers/src/pipeline.rs @@ -17,12 +17,12 @@ use std::sync::Arc; use ahash::{HashMap, HashMapExt}; use api::v1::helper::time_index_column_schema; -use api::v1::{ColumnDataType, RowInsertRequest, Rows, Value}; +use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value}; use common_time::timestamp::TimeUnit; use pipeline::{ - ContextReq, DispatchedTo, GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, Pipeline, PipelineContext, - PipelineDefinition, PipelineExecOutput, SchemaInfo, TransformedOutput, TransformerMode, - identity_pipeline, unwrap_or_continue_if_err, + ContextOpt, ContextReq, DispatchedTo, GREPTIME_INTERNAL_IDENTITY_PIPELINE_NAME, Pipeline, + PipelineContext, PipelineDefinition, PipelineExecOutput, SchemaInfo, TransformedOutput, + TransformerMode, identity_pipeline, unwrap_or_continue_if_err, }; use session::context::{Channel, QueryContextRef}; use snafu::ResultExt; @@ -114,7 +114,7 @@ async fn run_custom_pipeline( values: pipeline_maps, } = pipeline_req; let arr_len = pipeline_maps.len(); - let mut transformed_map = HashMap::new(); + let mut transformed_map: HashMap>> = HashMap::new(); let mut dispatched: BTreeMap> = BTreeMap::new(); let mut schema_info = match pipeline.transformer() { @@ -156,11 +156,11 @@ async fn run_custom_pipeline( PipelineExecOutput::Transformed(TransformedOutput { rows_by_context }) => { // Process each ContextOpt group separately for (opt, rows_with_suffix) in rows_by_context { - // Group rows by table name within each context + let rows_by_suffix = transformed_map.entry(opt).or_default(); + // Group rows by table suffix within each context. for (row, table_suffix) in rows_with_suffix { - let act_table_name = table_suffix_to_table_name(&table_name, table_suffix); - transformed_map - .entry((opt.clone(), act_table_name)) + rows_by_suffix + .entry(table_suffix.unwrap_or_default()) .or_insert_with(|| Vec::with_capacity(arr_len)) .push(row); } @@ -177,28 +177,28 @@ async fn run_custom_pipeline( let mut results = ContextReq::default(); - // Process transformed outputs. Each entry in transformed_map contains - // Vec grouped by (opt, table_name). + // Process transformed outputs. Rows are grouped by context first, then table suffix. let column_count = schema_info.schema.len(); - for ((opt, table_name), mut rows) in transformed_map { - // Pad rows to match final schema size (schema may have evolved during processing) - for row in &mut rows { - let diff = column_count.saturating_sub(row.values.len()); - for _ in 0..diff { - row.values.push(Value { value_data: None }); - } - } + let column_schemas = schema_info.column_schemas()?; + for (opt, rows_by_suffix) in transformed_map { + let row_requests = rows_by_suffix.into_iter().map(|(table_suffix, mut rows)| { + let table_name = table_suffix_to_table_name(&table_name, &table_suffix); + + // Pad rows to match final schema size (schema may have evolved during processing) + for row in &mut rows { + row.values.resize(column_count, Value { value_data: None }); + } - results.add_row( - &opt, RowInsertRequest { rows: Some(Rows { rows, - schema: schema_info.column_schemas()?, + schema: column_schemas.clone(), }), - table_name: table_name.clone(), - }, - ); + table_name, + } + }); + + results.add_rows(opt, row_requests); } // if current pipeline contains dispatcher and has several rules, we may @@ -245,9 +245,10 @@ async fn run_custom_pipeline( } #[inline] -fn table_suffix_to_table_name(table_name: &String, table_suffix: Option) -> String { - match table_suffix { - Some(suffix) => format!("{}{}", table_name, suffix), - None => table_name.clone(), +fn table_suffix_to_table_name(table_name: &str, table_suffix: &str) -> String { + if table_suffix.is_empty() { + table_name.to_string() + } else { + format!("{}{}", table_name, table_suffix) } }