refactor: streamline pipeline ingestion row building (#8236)

* refactor: minor pipeline refactor

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* refactor: use suffix as group key

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

---------

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

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<VrlValue>,
+) -> Result<usize> {
+    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::<Vec<VrlValue>>();
+    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!({

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<VrlValue>,
     pipeline_ctx: &PipelineContext<'_>,
+    max_nested_levels: usize,
 ) -> Result<(SchemaInfo, HashMap<ContextOpt, Vec<Row>>)> {
     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<Arc<table::Table>>,
     pipeline_ctx: &PipelineContext<'_>,
 ) -> Result<HashMap<ContextOpt, Rows>> {
-    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::<HashMap<ContextOpt, Rows>>())
-    })
+    let column_schemas = schema.column_schemas()?;
+    Ok(opt_map
+        .into_iter()
+        .map(|(opt, rows)| {
+            (
+                opt,
+                Rows {
+                    schema: column_schemas.clone(),
+                    rows,
+                },
+            )
+        })
+        .collect::<HashMap<ContextOpt, Rows>>())
 }
 
 /// 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, &params, 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() {

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<ContextOpt, HashMap<String, Vec<Row>>> = HashMap::new();
     let mut dispatched: BTreeMap<DispatchedTo, Vec<VrlValue>> = 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<Row> 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>) -> 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)
     }
 }