Skip to main content

pipeline/etl/transform/transformer/
greptime.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod coerce;
16
17use std::borrow::Cow;
18use std::collections::{BTreeMap, HashSet};
19use std::sync::Arc;
20
21use ahash::{HashMap, HashMapExt};
22use api::helper::{ColumnDataTypeWrapper, encode_json_value};
23use api::v1::column_def::{collect_column_options, options_from_column_schema};
24use api::v1::value::ValueData;
25use api::v1::{ColumnDataType, SemanticType};
26use arrow_schema::extension::ExtensionType;
27use coerce::{coerce_columns, coerce_value};
28use common_query::prelude::{greptime_timestamp, greptime_value};
29use common_telemetry::warn;
30use datatypes::data_type::ConcreteDataType;
31use datatypes::extension::json::{Json2ExtensionType, parse_legacy_json2_settings};
32use datatypes::json::JsonSettings;
33use datatypes::value::Value;
34use greptime_proto::v1::{ColumnSchema, Row, Rows, Value as GreptimeValue};
35use itertools::Itertools;
36use jsonb::Number;
37use once_cell::sync::OnceCell;
38use serde_json as serde_json_crate;
39use session::context::Channel;
40use snafu::OptionExt;
41use table::Table;
42use vrl::prelude::{Bytes, VrlValueConvert};
43use vrl::value::value::StdError;
44use vrl::value::{KeyString, Value as VrlValue};
45
46use crate::error::{
47    ArrayElementMustBeObjectSnafu, CoerceIncompatibleTypesSnafu,
48    IdentifyPipelineColumnTypeMismatchSnafu, InvalidTimestampSnafu, Result,
49    TimeIndexMustBeNonNullSnafu, TransformColumnNameMustBeUniqueSnafu,
50    TransformMultipleTimestampIndexSnafu, TransformTimestampIndexCountSnafu, ValueMustBeMapSnafu,
51};
52use crate::etl::PipelineDocVersion;
53use crate::etl::ctx_req::ContextOpt;
54use crate::etl::field::{Field, Fields};
55use crate::etl::transform::index::Index;
56use crate::etl::transform::{Transform, Transforms};
57use crate::{PipelineContext, truthy, unwrap_or_continue_if_err};
58
59const DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING: usize = 10;
60
61/// Row with potentially designated table suffix.
62pub type RowWithTableSuffix = (Row, Option<String>);
63
64/// fields not in the columns will be discarded
65/// to prevent automatic column creation in GreptimeDB
66#[derive(Debug, Clone)]
67pub struct GreptimeTransformer {
68    transforms: Transforms,
69    schema: Vec<ColumnSchema>,
70}
71
72/// Parameters that can be used to configure the greptime pipelines.
73#[derive(Debug, Default)]
74pub struct GreptimePipelineParams {
75    /// The original options for configuring the greptime pipelines.
76    /// This should not be used directly, instead, use the parsed shortcut option values.
77    options: HashMap<String, String>,
78
79    /// Whether to skip error when processing the pipeline.
80    pub skip_error: OnceCell<bool>,
81    /// Max nested levels when flattening JSON object. Defaults to
82    /// `DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING` when not provided.
83    pub max_nested_levels: OnceCell<usize>,
84}
85
86impl GreptimePipelineParams {
87    /// Create a `GreptimePipelineParams` from params string which is from the http header with key `x-greptime-pipeline-params`
88    /// The params is in the format of `key1=value1&key2=value2`,for example:
89    /// x-greptime-pipeline-params: max_nested_levels=5
90    pub fn from_params(params: Option<&str>) -> Self {
91        let options = Self::parse_header_str_to_map(params);
92
93        Self {
94            options,
95            skip_error: OnceCell::new(),
96            max_nested_levels: OnceCell::new(),
97        }
98    }
99
100    pub fn from_map(options: HashMap<String, String>) -> Self {
101        Self {
102            options,
103            skip_error: OnceCell::new(),
104            max_nested_levels: OnceCell::new(),
105        }
106    }
107
108    pub fn parse_header_str_to_map(params: Option<&str>) -> HashMap<String, String> {
109        if let Some(params) = params {
110            if params.is_empty() {
111                HashMap::new()
112            } else {
113                params
114                    .split('&')
115                    .filter_map(|s| s.split_once('='))
116                    .map(|(k, v)| (k.to_string(), v.to_string()))
117                    .collect::<HashMap<String, String>>()
118            }
119        } else {
120            HashMap::new()
121        }
122    }
123
124    /// Whether to skip error when processing the pipeline.
125    pub fn skip_error(&self) -> bool {
126        *self
127            .skip_error
128            .get_or_init(|| self.options.get("skip_error").map(truthy).unwrap_or(false))
129    }
130
131    /// Max nested levels for JSON flattening. If not provided or invalid,
132    /// falls back to `DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING`.
133    pub fn max_nested_levels(&self) -> usize {
134        *self.max_nested_levels.get_or_init(|| {
135            self.options
136                .get("max_nested_levels")
137                .and_then(|s| s.parse::<usize>().ok())
138                .filter(|v| *v > 0)
139                .unwrap_or(DEFAULT_MAX_NESTED_LEVELS_FOR_JSON_FLATTENING)
140        })
141    }
142}
143
144impl GreptimeTransformer {
145    /// Add a default timestamp column to the transforms
146    fn add_greptime_timestamp_column(transforms: &mut Transforms) {
147        let type_ = ColumnDataType::TimestampNanosecond;
148        let default = None;
149
150        let transform = Transform {
151            fields: Fields::one(Field::new(greptime_timestamp().to_string(), None)),
152            type_,
153            json_settings: None,
154            default,
155            index: Some(Index::Time),
156            index_options: None,
157            on_failure: Some(crate::etl::transform::OnFailure::Default),
158            tag: false,
159        };
160        transforms.push(transform);
161    }
162
163    /// Generate the schema for the GreptimeTransformer
164    fn init_schemas(transforms: &Transforms) -> Result<Vec<ColumnSchema>> {
165        let mut schema = vec![];
166        for transform in transforms.iter() {
167            schema.extend(coerce_columns(transform)?);
168        }
169        Ok(schema)
170    }
171}
172
173impl GreptimeTransformer {
174    pub fn new(mut transforms: Transforms, doc_version: &PipelineDocVersion) -> Result<Self> {
175        // empty check is done in the caller
176        let mut column_names_set = HashSet::new();
177        let mut timestamp_columns = vec![];
178
179        for transform in transforms.iter() {
180            let target_fields_set = transform
181                .fields
182                .iter()
183                .map(|f| f.target_or_input_field())
184                .collect::<HashSet<_>>();
185
186            let intersections: Vec<_> = column_names_set.intersection(&target_fields_set).collect();
187            if !intersections.is_empty() {
188                let duplicates = intersections.iter().join(",");
189                return TransformColumnNameMustBeUniqueSnafu { duplicates }.fail();
190            }
191
192            column_names_set.extend(target_fields_set);
193
194            if let Some(idx) = transform.index
195                && idx == Index::Time
196            {
197                match transform.fields.len() {
198                    //Safety unwrap is fine here because we have checked the length of real_fields
199                    1 => timestamp_columns.push(transform.fields.first().unwrap().input_field()),
200                    _ => {
201                        return TransformMultipleTimestampIndexSnafu {
202                            columns: transform.fields.iter().map(|x| x.input_field()).join(", "),
203                        }
204                        .fail();
205                    }
206                }
207            }
208        }
209
210        let schema = match timestamp_columns.len() {
211            0 if doc_version == &PipelineDocVersion::V1 => {
212                // compatible with v1, add a default timestamp column
213                GreptimeTransformer::add_greptime_timestamp_column(&mut transforms);
214                GreptimeTransformer::init_schemas(&transforms)?
215            }
216            1 => GreptimeTransformer::init_schemas(&transforms)?,
217            count => {
218                let columns = timestamp_columns.iter().join(", ");
219                return TransformTimestampIndexCountSnafu { count, columns }.fail();
220            }
221        };
222        Ok(GreptimeTransformer { transforms, schema })
223    }
224
225    pub fn transform_mut(
226        &self,
227        pipeline_map: &mut VrlValue,
228        is_v1: bool,
229    ) -> Result<Vec<GreptimeValue>> {
230        self.transform_mut_with_schema(pipeline_map, is_v1, &SchemaInfo::default(), None)
231    }
232
233    pub(crate) fn transform_mut_with_schema(
234        &self,
235        pipeline_map: &mut VrlValue,
236        is_v1: bool,
237        schema_info: &SchemaInfo,
238        table_suffix: Option<&str>,
239    ) -> Result<Vec<GreptimeValue>> {
240        let mut values = vec![GreptimeValue { value_data: None }; self.schema.len()];
241        let mut output_index = 0;
242        for transform in self.transforms.iter() {
243            for field in transform.fields.iter() {
244                let column_name = field.input_field();
245
246                let pipeline_map = pipeline_map.as_object_mut().context(ValueMustBeMapSnafu)?;
247                // let keep us `get` here to be compatible with v1
248                match pipeline_map.get(column_name) {
249                    Some(v) => {
250                        let json_settings = if transform.type_ == ColumnDataType::Json
251                            && matches!(v, VrlValue::Array(_) | VrlValue::Object(_))
252                        {
253                            schema_info.json_settings_for_column(
254                                field.target_or_input_field(),
255                                table_suffix,
256                            )?
257                        } else {
258                            None
259                        };
260                        let value_data = coerce_value(v, transform, json_settings.as_ref())?;
261                        // every transform fields has only one output field
262                        values[output_index] = GreptimeValue { value_data };
263                    }
264                    None => {
265                        let value_data = match transform.on_failure {
266                            Some(crate::etl::transform::OnFailure::Default) => {
267                                match transform.get_default() {
268                                    Some(default) => Some(default.clone()),
269                                    None => transform.get_default_value_when_data_is_none(),
270                                }
271                            }
272                            Some(crate::etl::transform::OnFailure::Ignore) => None,
273                            None => None,
274                        };
275                        if transform.is_timeindex() && value_data.is_none() {
276                            return TimeIndexMustBeNonNullSnafu.fail();
277                        }
278                        values[output_index] = GreptimeValue { value_data };
279                    }
280                }
281                output_index += 1;
282                if !is_v1 {
283                    // remove the column from the pipeline_map
284                    // so that the auto-transform can use the rest fields
285                    pipeline_map.remove(column_name);
286                }
287            }
288        }
289        Ok(values)
290    }
291
292    pub fn transforms(&self) -> &Transforms {
293        &self.transforms
294    }
295
296    pub fn schemas(&self) -> &Vec<greptime_proto::v1::ColumnSchema> {
297        &self.schema
298    }
299
300    pub fn has_json_transform(&self) -> bool {
301        self.transforms
302            .iter()
303            .any(|transform| transform.type_ == ColumnDataType::Json)
304    }
305
306    pub fn transforms_mut(&mut self) -> &mut Transforms {
307        &mut self.transforms
308    }
309}
310
311#[derive(Clone)]
312pub struct ColumnMetadata {
313    column_schema: datatypes::schema::ColumnSchema,
314    semantic_type: SemanticType,
315    json_settings: OnceCell<JsonSettings>,
316}
317
318impl From<ColumnSchema> for ColumnMetadata {
319    fn from(value: ColumnSchema) -> Self {
320        let datatype = value.datatype();
321        let semantic_type = value.semantic_type();
322        let ColumnSchema {
323            column_name,
324            datatype: _,
325            semantic_type: _,
326            datatype_extension,
327            options,
328        } = value;
329
330        let column_schema = datatypes::schema::ColumnSchema::new(
331            column_name,
332            ColumnDataTypeWrapper::new(datatype, datatype_extension).into(),
333            semantic_type != SemanticType::Timestamp,
334        );
335
336        let metadata = collect_column_options(options.as_ref());
337        let column_schema = column_schema.with_metadata(metadata);
338
339        Self {
340            column_schema,
341            semantic_type,
342            json_settings: OnceCell::new(),
343        }
344    }
345}
346
347impl ColumnMetadata {
348    fn json_settings(&self) -> Result<&JsonSettings> {
349        self.json_settings.get_or_try_init(|| {
350            if let Some(extension) = self.column_schema.extension_type::<Json2ExtensionType>()? {
351                Ok(extension.metadata().json_settings().clone())
352            } else {
353                Ok(parse_legacy_json2_settings(self.column_schema.metadata())?.unwrap_or_default())
354            }
355        })
356    }
357}
358
359impl TryFrom<ColumnMetadata> for ColumnSchema {
360    type Error = api::error::Error;
361
362    fn try_from(value: ColumnMetadata) -> std::result::Result<Self, Self::Error> {
363        let ColumnMetadata {
364            column_schema,
365            semantic_type,
366            ..
367        } = value;
368
369        let options = options_from_column_schema(&column_schema);
370
371        let (datatype, datatype_extension) =
372            ColumnDataTypeWrapper::try_from(column_schema.data_type).map(|x| x.into_parts())?;
373
374        Ok(ColumnSchema {
375            column_name: column_schema.name,
376            datatype: datatype as _,
377            semantic_type: semantic_type as _,
378            datatype_extension,
379            options,
380        })
381    }
382}
383
384/// This is used to record the current state schema information and a sequential cache of field names.
385/// As you traverse the user input JSON, this will change.
386/// It will record a superset of all user input schemas.
387#[derive(Default)]
388pub struct SchemaInfo {
389    /// schema info
390    pub schema: Vec<ColumnMetadata>,
391    /// index of the column name
392    pub index: HashMap<String, usize>,
393    /// Tables already looked up, keyed by their resolved suffix. Missing tables are cached as None.
394    tables: HashMap<String, Option<Arc<Table>>>,
395}
396
397impl SchemaInfo {
398    pub fn with_capacity(capacity: usize) -> Self {
399        Self {
400            schema: Vec::with_capacity(capacity),
401            index: HashMap::with_capacity(capacity),
402            tables: HashMap::new(),
403        }
404    }
405
406    pub fn from_schema_list(schema_list: Vec<ColumnSchema>) -> Self {
407        let mut index = HashMap::new();
408        for (i, schema) in schema_list.iter().enumerate() {
409            index.insert(schema.column_name.clone(), i);
410        }
411        Self {
412            schema: schema_list.into_iter().map(Into::into).collect(),
413            index,
414            tables: HashMap::new(),
415        }
416    }
417
418    pub fn set_table(&mut self, table: Option<Arc<Table>>) {
419        self.set_table_for_suffix(String::new(), table);
420    }
421
422    pub fn has_table_for_suffix(&self, table_suffix: &str) -> bool {
423        self.tables.contains_key(table_suffix)
424    }
425
426    pub fn set_table_for_suffix(&mut self, table_suffix: String, table: Option<Arc<Table>>) {
427        self.tables.insert(table_suffix, table);
428    }
429
430    fn table_for_suffix(&self, table_suffix: Option<&str>) -> Option<&Arc<Table>> {
431        let table_suffix = table_suffix.unwrap_or_default();
432        match self.tables.get(table_suffix) {
433            Some(table) => table.as_ref(),
434            None if !table_suffix.is_empty() => self.tables.get("").and_then(Option::as_ref),
435            None => None,
436        }
437    }
438
439    fn find_column_schema_in_table(&self, column_name: &str) -> Option<ColumnMetadata> {
440        if let Some(table) = self.table_for_suffix(None)
441            && let Some(i) = table.schema_ref().column_index_by_name(column_name)
442        {
443            let column_schema = table.schema_ref().column_schemas()[i].clone();
444
445            let semantic_type = if column_schema.is_time_index() {
446                SemanticType::Timestamp
447            } else if table.table_info().meta.primary_key_indices.contains(&i) {
448                SemanticType::Tag
449            } else {
450                SemanticType::Field
451            };
452
453            Some(ColumnMetadata {
454                column_schema,
455                semantic_type,
456                json_settings: OnceCell::new(),
457            })
458        } else {
459            None
460        }
461    }
462
463    fn json_settings_for_column(
464        &self,
465        column_name: &str,
466        table_suffix: Option<&str>,
467    ) -> Result<Option<JsonSettings>> {
468        let Some(column_schema) = self
469            .table_for_suffix(table_suffix)
470            .and_then(|table| table.schema_ref().column_schema_by_name(column_name))
471        else {
472            return Ok(None);
473        };
474        if !column_schema.data_type.is_json2() {
475            return Ok(None);
476        }
477
478        if let Some(extension) = column_schema.extension_type::<Json2ExtensionType>()? {
479            Ok(Some(extension.metadata().json_settings().clone()))
480        } else {
481            Ok(Some(
482                parse_legacy_json2_settings(column_schema.metadata())?.unwrap_or_default(),
483            ))
484        }
485    }
486
487    pub fn column_schemas(&self) -> api::error::Result<Vec<ColumnSchema>> {
488        self.schema
489            .iter()
490            .map(|x| x.clone().try_into())
491            .collect::<api::error::Result<Vec<_>>>()
492    }
493}
494
495fn resolve_schema(
496    index: Option<usize>,
497    pipeline_context: &PipelineContext,
498    column: &str,
499    value_type: &ConcreteDataType,
500    schema_info: &mut SchemaInfo,
501) -> Result<()> {
502    if let Some(index) = index {
503        let column_type = &mut schema_info.schema[index].column_schema.data_type;
504        match (column_type, value_type) {
505            (column_type, value_type) if column_type == value_type => Ok(()),
506            // Auto-transformed arrays and objects use JSONB-encoded binary values.
507            (ConcreteDataType::Json(json_type), ConcreteDataType::Binary(_))
508                if !json_type.is_json2() =>
509            {
510                Ok(())
511            }
512            (ConcreteDataType::Json(column_type), ConcreteDataType::Json(value_type))
513                if column_type.is_json2() && value_type.is_json2() =>
514            {
515                Ok(())
516            }
517            (column_type, value_type) => IdentifyPipelineColumnTypeMismatchSnafu {
518                column,
519                expected: column_type.to_string(),
520                actual: value_type.to_string(),
521            }
522            .fail(),
523        }
524    } else {
525        let column_schema = schema_info
526            .find_column_schema_in_table(column)
527            .unwrap_or_else(|| {
528                let semantic_type = decide_semantic(pipeline_context, column);
529                let column_schema = datatypes::schema::ColumnSchema::new(
530                    column,
531                    value_type.clone(),
532                    semantic_type != SemanticType::Timestamp,
533                );
534                ColumnMetadata {
535                    column_schema,
536                    semantic_type,
537                    json_settings: OnceCell::new(),
538                }
539            });
540        let key = column.to_string();
541        schema_info.schema.push(column_schema);
542        schema_info.index.insert(key, schema_info.schema.len() - 1);
543        Ok(())
544    }
545}
546
547fn calc_ts(p_ctx: &PipelineContext, values: &VrlValue) -> Result<Option<ValueData>> {
548    match p_ctx.channel {
549        Channel::Prometheus => {
550            let ts = values
551                .as_object()
552                .and_then(|m| m.get(greptime_timestamp()))
553                .and_then(|ts| ts.try_into_i64().ok())
554                .unwrap_or_default();
555            Ok(Some(ValueData::TimestampMillisecondValue(ts)))
556        }
557        _ => {
558            let custom_ts = p_ctx.pipeline_definition.get_custom_ts();
559            match custom_ts {
560                Some(ts) => {
561                    let ts_field = values.as_object().and_then(|m| m.get(ts.get_column_name()));
562                    Some(ts.get_timestamp_value(ts_field)).transpose()
563                }
564                None => Ok(Some(ValueData::TimestampNanosecondValue(
565                    chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default(),
566                ))),
567            }
568        }
569    }
570}
571
572/// Converts VRL values to Greptime rows grouped by their ContextOpt.
573/// # Returns
574/// A HashMap where keys are `ContextOpt` and values are vectors of (row, table_suffix) pairs.
575/// Single object input produces one ContextOpt group with one row.
576/// Array input groups rows by their per-element ContextOpt values.
577///
578/// # Errors
579/// - `ArrayElementMustBeObject` if an array element is not an object
580pub(crate) fn values_to_rows(
581    schema_info: &mut SchemaInfo,
582    mut values: VrlValue,
583    pipeline_ctx: &PipelineContext<'_>,
584    row: Option<Vec<GreptimeValue>>,
585    need_calc_ts: bool,
586    tablesuffix_template: Option<&crate::tablesuffix::TableSuffixTemplate>,
587) -> Result<std::collections::HashMap<ContextOpt, Vec<RowWithTableSuffix>>> {
588    let skip_error = pipeline_ctx.pipeline_param.skip_error();
589    let VrlValue::Array(arr) = values else {
590        // Single object: extract ContextOpt and table_suffix
591        let mut result = std::collections::HashMap::new();
592
593        let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &values);
594        let opt = match ContextOpt::from_pipeline_map_to_opt(&mut values) {
595            Ok(r) => r,
596            Err(e) => return if skip_error { Ok(result) } else { Err(e) },
597        };
598
599        let row = match values_to_row(schema_info, values, pipeline_ctx, row, need_calc_ts) {
600            Ok(r) => r,
601            Err(e) => return if skip_error { Ok(result) } else { Err(e) },
602        };
603        result.insert(opt, vec![(row, table_suffix)]);
604        return Ok(result);
605    };
606
607    let mut rows_by_context: std::collections::HashMap<ContextOpt, Vec<RowWithTableSuffix>> =
608        std::collections::HashMap::new();
609    for (index, mut value) in arr.into_iter().enumerate() {
610        if !value.is_object() {
611            unwrap_or_continue_if_err!(
612                ArrayElementMustBeObjectSnafu {
613                    index,
614                    actual_type: value.kind_str().to_string(),
615                }
616                .fail(),
617                skip_error
618            );
619        }
620
621        // Extract ContextOpt and table_suffix for this element
622        let table_suffix = ContextOpt::resolve_table_suffix(tablesuffix_template, &value);
623        let opt = unwrap_or_continue_if_err!(
624            ContextOpt::from_pipeline_map_to_opt(&mut value),
625            skip_error
626        );
627        let transformed_row = unwrap_or_continue_if_err!(
628            values_to_row(schema_info, value, pipeline_ctx, row.clone(), need_calc_ts),
629            skip_error
630        );
631        rows_by_context
632            .entry(opt)
633            .or_default()
634            .push((transformed_row, table_suffix));
635    }
636    Ok(rows_by_context)
637}
638
639/// `need_calc_ts` happens in two cases:
640/// 1. full greptime_identity
641/// 2. auto-transform without transformer
642///
643/// if transform is present in custom pipeline in v2 mode
644/// we dont need to calc ts again, nor do we need to check ts column name
645pub(crate) fn values_to_row(
646    schema_info: &mut SchemaInfo,
647    values: VrlValue,
648    pipeline_ctx: &PipelineContext<'_>,
649    row: Option<Vec<GreptimeValue>>,
650    need_calc_ts: bool,
651) -> Result<Row> {
652    let mut row: Vec<GreptimeValue> =
653        row.unwrap_or_else(|| Vec::with_capacity(schema_info.schema.len()));
654    let custom_ts = pipeline_ctx.pipeline_definition.get_custom_ts();
655
656    if need_calc_ts {
657        // calculate timestamp value based on the channel
658        let ts = calc_ts(pipeline_ctx, &values)?;
659        row.push(GreptimeValue { value_data: ts });
660    }
661
662    row.resize(schema_info.schema.len(), GreptimeValue { value_data: None });
663
664    // skip ts column
665    let ts_column_name = custom_ts
666        .as_ref()
667        .map_or(greptime_timestamp(), |ts| ts.get_column_name());
668
669    let values = values.into_object().context(ValueMustBeMapSnafu)?;
670
671    for (column_name, value) in values {
672        if need_calc_ts && column_name.as_str() == ts_column_name {
673            continue;
674        }
675
676        resolve_value(
677            value,
678            column_name.into(),
679            &mut row,
680            schema_info,
681            pipeline_ctx,
682        )?;
683    }
684    Ok(Row { values: row })
685}
686
687fn decide_semantic(p_ctx: &PipelineContext, column_name: &str) -> SemanticType {
688    if p_ctx.channel == Channel::Prometheus && column_name != greptime_value() {
689        SemanticType::Tag
690    } else {
691        SemanticType::Field
692    }
693}
694
695fn resolve_value(
696    value: VrlValue,
697    column_name: String,
698    row: &mut Vec<GreptimeValue>,
699    schema_info: &mut SchemaInfo,
700    p_ctx: &PipelineContext,
701) -> Result<()> {
702    let index = schema_info.index.get(&column_name).copied();
703
704    let value_data = match value {
705        VrlValue::Null => return Ok(()),
706
707        VrlValue::Integer(v) => {
708            // safe unwrap after type matched
709            resolve_schema(
710                index,
711                p_ctx,
712                &column_name,
713                &ConcreteDataType::int64_datatype(),
714                schema_info,
715            )?;
716            Some(ValueData::I64Value(v))
717        }
718
719        VrlValue::Float(v) => {
720            // safe unwrap after type matched
721            resolve_schema(
722                index,
723                p_ctx,
724                &column_name,
725                &ConcreteDataType::float64_datatype(),
726                schema_info,
727            )?;
728            Some(ValueData::F64Value(v.into()))
729        }
730
731        VrlValue::Boolean(v) => {
732            resolve_schema(
733                index,
734                p_ctx,
735                &column_name,
736                &ConcreteDataType::boolean_datatype(),
737                schema_info,
738            )?;
739            Some(ValueData::BoolValue(v))
740        }
741
742        VrlValue::Bytes(v) => {
743            resolve_schema(
744                index,
745                p_ctx,
746                &column_name,
747                &ConcreteDataType::string_datatype(),
748                schema_info,
749            )?;
750            Some(ValueData::StringValue(
751                String::from_utf8_lossy(&v).into_owned(),
752            ))
753        }
754
755        VrlValue::Regex(v) => {
756            warn!(
757                "Persisting regex value in the table, this should not happen, column_name: {}",
758                column_name
759            );
760            resolve_schema(
761                index,
762                p_ctx,
763                &column_name,
764                &ConcreteDataType::string_datatype(),
765                schema_info,
766            )?;
767            Some(ValueData::StringValue(v.to_string()))
768        }
769
770        VrlValue::Timestamp(ts) => {
771            let ns = ts.timestamp_nanos_opt().context(InvalidTimestampSnafu {
772                input: ts.to_rfc3339(),
773            })?;
774            resolve_schema(
775                index,
776                p_ctx,
777                &column_name,
778                &ConcreteDataType::timestamp_nanosecond_datatype(),
779                schema_info,
780            )?;
781            Some(ValueData::TimestampNanosecondValue(ns))
782        }
783
784        VrlValue::Array(_) | VrlValue::Object(_) => {
785            let index = index.or_else(|| {
786                let column = schema_info.find_column_schema_in_table(&column_name)?;
787                let index = schema_info.schema.len();
788                schema_info.schema.push(column);
789                schema_info.index.insert(column_name.clone(), index);
790                Some(index)
791            });
792            // TODO(LFC): Default to JSON2 for auto-created tables.
793            let json2_index = index.filter(|&index| {
794                matches!(
795                    &schema_info.schema[index].column_schema.data_type,
796                    ConcreteDataType::Json(column_type) if column_type.is_json2()
797                )
798            });
799
800            let value = if let Some(index) = json2_index {
801                let value: serde_json::Value = value.try_into().map_err(|e: StdError| {
802                    CoerceIncompatibleTypesSnafu { msg: e.to_string() }.build()
803                })?;
804                let value = schema_info.schema[index].json_settings()?.encode(value)?;
805
806                resolve_schema(
807                    Some(index),
808                    p_ctx,
809                    &column_name,
810                    &ConcreteDataType::json2(Default::default()),
811                    schema_info,
812                )?;
813
814                let Value::Json(value) = value else {
815                    unreachable!()
816                };
817                ValueData::JsonValue(encode_json_value(*value))
818            } else {
819                resolve_schema(
820                    index,
821                    p_ctx,
822                    &column_name,
823                    &ConcreteDataType::binary_datatype(),
824                    schema_info,
825                )?;
826
827                let value = vrl_value_to_jsonb_value(&value);
828                ValueData::BinaryValue(value.to_vec())
829            };
830            Some(value)
831        }
832    };
833
834    let value = GreptimeValue { value_data };
835    if let Some(index) = index {
836        row[index] = value;
837    } else {
838        row.push(value);
839    }
840    Ok(())
841}
842
843fn vrl_value_to_jsonb_value<'a>(value: &'a VrlValue) -> jsonb::Value<'a> {
844    match value {
845        VrlValue::Bytes(bytes) => jsonb::Value::String(String::from_utf8_lossy(bytes)),
846        VrlValue::Regex(value_regex) => jsonb::Value::String(Cow::Borrowed(value_regex.as_str())),
847        VrlValue::Integer(i) => jsonb::Value::Number(Number::Int64(*i)),
848        VrlValue::Float(not_nan) => jsonb::Value::Number(Number::Float64(not_nan.into_inner())),
849        VrlValue::Boolean(b) => jsonb::Value::Bool(*b),
850        VrlValue::Timestamp(date_time) => jsonb::Value::String(Cow::Owned(date_time.to_rfc3339())),
851        VrlValue::Object(btree_map) => jsonb::Value::Object(
852            btree_map
853                .iter()
854                .map(|(key, value)| (key.to_string(), vrl_value_to_jsonb_value(value)))
855                .collect(),
856        ),
857        VrlValue::Array(values) => jsonb::Value::Array(
858            values
859                .iter()
860                .map(|value| vrl_value_to_jsonb_value(value))
861                .collect(),
862        ),
863        VrlValue::Null => jsonb::Value::Null,
864    }
865}
866
867fn identity_pipeline_inner(
868    pipeline_maps: Vec<VrlValue>,
869    pipeline_ctx: &PipelineContext<'_>,
870    max_nested_levels: usize,
871) -> Result<(SchemaInfo, HashMap<ContextOpt, Vec<Row>>)> {
872    let skip_error = pipeline_ctx.pipeline_param.skip_error();
873    let mut schema_info = SchemaInfo::default();
874    let custom_ts = pipeline_ctx.pipeline_definition.get_custom_ts();
875
876    // set time index column schema first
877    let column_schema = datatypes::schema::ColumnSchema::new(
878        custom_ts
879            .map(|ts| ts.get_column_name().to_string())
880            .unwrap_or_else(|| greptime_timestamp().to_string()),
881        custom_ts
882            .map(|c| ConcreteDataType::from(ColumnDataTypeWrapper::new(c.get_datatype(), None)))
883            .unwrap_or_else(|| {
884                if pipeline_ctx.channel == Channel::Prometheus {
885                    ConcreteDataType::timestamp_millisecond_datatype()
886                } else {
887                    ConcreteDataType::timestamp_nanosecond_datatype()
888                }
889            }),
890        false,
891    );
892    schema_info.schema.push(ColumnMetadata {
893        column_schema,
894        semantic_type: SemanticType::Timestamp,
895        json_settings: OnceCell::new(),
896    });
897
898    let mut opt_map = HashMap::new();
899    let len = pipeline_maps.len();
900
901    for pipeline_map in pipeline_maps {
902        let mut pipeline_map =
903            unwrap_or_continue_if_err!(flatten_object(pipeline_map, max_nested_levels), skip_error);
904        let opt = unwrap_or_continue_if_err!(
905            ContextOpt::from_pipeline_map_to_opt(&mut pipeline_map),
906            skip_error
907        );
908        let row = unwrap_or_continue_if_err!(
909            values_to_row(&mut schema_info, pipeline_map, pipeline_ctx, None, true),
910            skip_error
911        );
912
913        opt_map
914            .entry(opt)
915            .or_insert_with(|| Vec::with_capacity(len))
916            .push(row);
917    }
918
919    let column_count = schema_info.schema.len();
920    for row in opt_map.values_mut() {
921        for row in row.iter_mut() {
922            assert!(
923                column_count >= row.values.len(),
924                "column_count: {}, row.values.len(): {}",
925                column_count,
926                row.values.len()
927            );
928            row.values
929                .resize(column_count, GreptimeValue { value_data: None });
930        }
931    }
932
933    Ok((schema_info, opt_map))
934}
935
936/// Identity pipeline for Greptime
937/// This pipeline will convert the input JSON array to Greptime Rows
938/// params table is used to set the semantic type of the row key column to Tag
939/// 1. The pipeline will add a default timestamp column to the schema
940/// 2. The pipeline not resolve NULL value
941/// 3. The pipeline assumes that the json format is fixed
942/// 4. The pipeline will return an error if the same column datatype is mismatched
943/// 5. The pipeline will analyze the schema of each json record and merge them to get the final schema.
944pub fn identity_pipeline(
945    array: Vec<VrlValue>,
946    table: Option<Arc<table::Table>>,
947    pipeline_ctx: &PipelineContext<'_>,
948) -> Result<HashMap<ContextOpt, Rows>> {
949    let max_nested_levels = pipeline_ctx.pipeline_param.max_nested_levels();
950
951    let (mut schema, opt_map) = identity_pipeline_inner(array, pipeline_ctx, max_nested_levels)?;
952    if let Some(table) = table {
953        let table_info = table.table_info();
954        for tag_name in table_info.meta.row_key_column_names() {
955            if let Some(index) = schema.index.get(tag_name) {
956                schema.schema[*index].semantic_type = SemanticType::Tag;
957            }
958        }
959    }
960
961    let column_schemas = schema.column_schemas()?;
962    Ok(opt_map
963        .into_iter()
964        .map(|(opt, rows)| {
965            (
966                opt,
967                Rows {
968                    schema: column_schemas.clone(),
969                    rows,
970                },
971            )
972        })
973        .collect::<HashMap<ContextOpt, Rows>>())
974}
975
976/// Consumes the JSON object and consumes it into a single-level object.
977///
978/// The `max_nested_levels` parameter is used to limit how deep to flatten nested JSON objects.
979/// When the maximum level is reached, the remaining nested structure is serialized to a JSON
980/// string and stored at the current flattened key.
981pub fn flatten_object(object: VrlValue, max_nested_levels: usize) -> Result<VrlValue> {
982    let mut flattened = BTreeMap::new();
983    let object = object.into_object().context(ValueMustBeMapSnafu)?;
984
985    if !object.is_empty() {
986        // it will use recursion to flatten the object.
987        do_flatten_object(&mut flattened, None, object, 1, max_nested_levels);
988    }
989
990    Ok(VrlValue::Object(flattened))
991}
992
993pub(crate) fn vrl_value_to_serde_json(value: &VrlValue) -> serde_json_crate::Value {
994    match value {
995        VrlValue::Null => serde_json_crate::Value::Null,
996        VrlValue::Boolean(b) => serde_json_crate::Value::Bool(*b),
997        VrlValue::Integer(i) => serde_json_crate::Value::Number((*i).into()),
998        VrlValue::Float(not_nan) => serde_json_crate::Number::from_f64(not_nan.into_inner())
999            .map(serde_json_crate::Value::Number)
1000            .unwrap_or(serde_json_crate::Value::Null),
1001        VrlValue::Bytes(bytes) => {
1002            serde_json_crate::Value::String(String::from_utf8_lossy(bytes).into_owned())
1003        }
1004        VrlValue::Regex(re) => serde_json_crate::Value::String(re.as_str().to_string()),
1005        VrlValue::Timestamp(ts) => serde_json_crate::Value::String(ts.to_rfc3339()),
1006        VrlValue::Array(arr) => {
1007            serde_json_crate::Value::Array(arr.iter().map(vrl_value_to_serde_json).collect())
1008        }
1009        VrlValue::Object(map) => serde_json_crate::Value::Object(
1010            map.iter()
1011                .map(|(k, v)| (k.to_string(), vrl_value_to_serde_json(v)))
1012                .collect(),
1013        ),
1014    }
1015}
1016
1017fn do_flatten_object(
1018    dest: &mut BTreeMap<KeyString, VrlValue>,
1019    base: Option<&str>,
1020    object: BTreeMap<KeyString, VrlValue>,
1021    current_level: usize,
1022    max_nested_levels: usize,
1023) {
1024    for (key, value) in object {
1025        let new_key = base.map_or_else(
1026            || key.clone(),
1027            |base_key| format!("{base_key}.{key}").into(),
1028        );
1029
1030        match value {
1031            VrlValue::Object(object) => {
1032                if current_level >= max_nested_levels {
1033                    // Reached the maximum level; stringify the remaining object.
1034                    let json_string = serde_json_crate::to_string(&vrl_value_to_serde_json(
1035                        &VrlValue::Object(object),
1036                    ))
1037                    .unwrap_or_else(|_| String::from("{}"));
1038                    dest.insert(new_key, VrlValue::Bytes(Bytes::from(json_string)));
1039                } else {
1040                    do_flatten_object(
1041                        dest,
1042                        Some(&new_key),
1043                        object,
1044                        current_level + 1,
1045                        max_nested_levels,
1046                    );
1047                }
1048            }
1049            // Arrays are stringified to ensure no JSON column types in the result.
1050            VrlValue::Array(_) => {
1051                let json_string = serde_json_crate::to_string(&vrl_value_to_serde_json(&value))
1052                    .unwrap_or_else(|_| String::from("[]"));
1053                dest.insert(new_key, VrlValue::Bytes(Bytes::from(json_string)));
1054            }
1055            // Other leaf types are inserted as-is.
1056            _ => {
1057                dest.insert(new_key, value);
1058            }
1059        }
1060    }
1061}
1062
1063#[cfg(test)]
1064mod tests {
1065    use api::v1::SemanticType;
1066    use common_recordbatch::RecordBatch;
1067    use datatypes::extension::json::JsonMetadata;
1068    use datatypes::json::JsonTypeHint;
1069    use datatypes::schema::{ColumnSchema as DatatypeColumnSchema, Schema};
1070    use table::test_util::MemTable;
1071
1072    use super::*;
1073    use crate::{Content, PipelineDefinition, identity_pipeline, parse, setup_pipeline};
1074
1075    #[test]
1076    fn test_auto_transform_existing_json_column()
1077    -> std::result::Result<(), Box<dyn std::error::Error>> {
1078        let pipeline = parse(&Content::Yaml(
1079            r#"
1080version: 2
1081transform:
1082  - field: ts
1083    type: epoch, ms
1084    index: timestamp
1085"#,
1086        ))?;
1087        let (pipeline, mut schema_info, pipeline_definition, pipeline_params) =
1088            setup_pipeline!(pipeline);
1089        let pipeline_context =
1090            PipelineContext::new(&pipeline_definition, &pipeline_params, Channel::Unknown);
1091        schema_info.set_table(Some(MemTable::table(
1092            "events",
1093            RecordBatch::new_empty(Arc::new(Schema::new(vec![DatatypeColumnSchema::new(
1094                "payload",
1095                ConcreteDataType::Json(Default::default()),
1096                true,
1097            )]))),
1098        )));
1099        let inputs = [
1100            serde_json::Value::Null,
1101            serde_json::json!([2, 4, 6]),
1102            serde_json::json!({"active": true}),
1103            serde_json::json!({"a\"b": "\u{0001}"}),
1104            serde_json::json!([]),
1105            serde_json::json!({}),
1106            serde_json::Value::Null,
1107        ];
1108        let records = inputs
1109            .iter()
1110            .enumerate()
1111            .map(|(i, payload)| serde_json::json!({"ts": i, "payload": payload}))
1112            .collect::<Vec<_>>();
1113        let output = pipeline
1114            .exec_mut(
1115                serde_json::Value::Array(records).into(),
1116                &pipeline_context,
1117                &mut schema_info,
1118            )?
1119            .into_transformed()
1120            .ok_or("expected transformed rows")?;
1121        assert_eq!(output.len(), inputs.len());
1122        let index = schema_info.index["payload"];
1123        assert_eq!(
1124            schema_info.schema[index].column_schema.data_type,
1125            ConcreteDataType::Json(Default::default())
1126        );
1127        for ((row, _), expected) in output.iter().zip(&inputs) {
1128            if expected.is_null() {
1129                assert_eq!(
1130                    row.values.get(index).and_then(|v| v.value_data.as_ref()),
1131                    None
1132                );
1133            } else {
1134                let Some(ValueData::BinaryValue(value)) = &row.values[index].value_data else {
1135                    return Err("expected JSONB-encoded value".into());
1136                };
1137                assert_eq!(datatypes::types::jsonb_to_serde_json(value)?, *expected);
1138            }
1139        }
1140        let error = pipeline
1141            .exec_mut(
1142                serde_json::json!({"ts": 10, "payload": 7}).into(),
1143                &pipeline_context,
1144                &mut schema_info,
1145            )
1146            .err()
1147            .ok_or("expected a datatype mismatch")?;
1148        assert!(
1149            error
1150                .to_string()
1151                .contains("expected datatype: Json, actual datatype: Int64"),
1152            "{error:?}"
1153        );
1154        Ok(())
1155    }
1156
1157    #[test]
1158    fn test_column_metadata_caches_json_settings() -> Result<()> {
1159        let column = ColumnMetadata {
1160            column_schema: datatypes::schema::ColumnSchema::new(
1161                "data",
1162                ConcreteDataType::json2(Default::default()),
1163                true,
1164            ),
1165            semantic_type: SemanticType::Field,
1166            json_settings: OnceCell::new(),
1167        };
1168
1169        let first = column.json_settings()?;
1170        let second = column.json_settings()?;
1171        assert!(std::ptr::eq(first, second));
1172        Ok(())
1173    }
1174
1175    #[test]
1176    fn test_transform_json2_uses_destination_table_settings() {
1177        let table = |name: &str, settings: JsonSettings, sample: serde_json::Value| {
1178            let data_type = settings.encode(sample).unwrap().data_type();
1179            let mut column_schema = DatatypeColumnSchema::new("payload", data_type, true);
1180            column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(
1181                JsonMetadata::new(settings),
1182            )));
1183            MemTable::table(
1184                name,
1185                RecordBatch::new_empty(Arc::new(Schema::new(vec![column_schema]))),
1186            )
1187        };
1188        let int_settings = JsonSettings::try_new(
1189            vec![JsonTypeHint {
1190                path: vec!["age".to_string()],
1191                data_type: ConcreteDataType::int64_datatype(),
1192                inverted_index: false,
1193                skipping_index: None,
1194            }],
1195            None,
1196        )
1197        .unwrap();
1198        let mut schema_info = SchemaInfo::default();
1199        schema_info.set_table(Some(table(
1200            "events",
1201            int_settings,
1202            serde_json::json!({"age": 42}),
1203        )));
1204        let string_settings = JsonSettings::try_new(
1205            vec![JsonTypeHint {
1206                path: vec!["age".to_string()],
1207                data_type: ConcreteDataType::string_datatype(),
1208                inverted_index: false,
1209                skipping_index: None,
1210            }],
1211            None,
1212        )
1213        .unwrap();
1214        schema_info.set_table_for_suffix(
1215            "_mobile".to_string(),
1216            Some(table(
1217                "events_mobile",
1218                string_settings,
1219                serde_json::json!({"age": "42"}),
1220            )),
1221        );
1222
1223        let pipeline = crate::parse(&crate::Content::Yaml(
1224            r#"
1225transform:
1226  - field: source, payload
1227    type: json2
1228table_suffix: _${device}
1229"#,
1230        ))
1231        .unwrap();
1232        let (pipeline, _, pipeline_definition, pipeline_params) = crate::setup_pipeline!(pipeline);
1233        let pipeline_context =
1234            PipelineContext::new(&pipeline_definition, &pipeline_params, Channel::Unknown);
1235
1236        let error = pipeline
1237            .exec_mut(
1238                serde_json::json!({"source": {"age": "42"}}).into(),
1239                &pipeline_context,
1240                &mut schema_info,
1241            )
1242            .unwrap_err();
1243        assert!(
1244            error.to_string().contains("does not match JSON2 type hint"),
1245            "{error:?}"
1246        );
1247
1248        let mut rows = pipeline
1249            .exec_mut(
1250                serde_json::json!({"source": {"age": "42"}, "device": "mobile"}).into(),
1251                &pipeline_context,
1252                &mut schema_info,
1253            )
1254            .unwrap()
1255            .into_transformed()
1256            .unwrap();
1257        let (row, table_suffix) = rows.swap_remove(0);
1258        assert_eq!(table_suffix.as_deref(), Some("_mobile"));
1259        assert!(matches!(
1260            &row.values[0].value_data,
1261            Some(ValueData::JsonValue(_))
1262        ));
1263    }
1264
1265    #[test]
1266    fn test_identify_pipeline() {
1267        let params = GreptimePipelineParams::default();
1268        let pipeline_ctx = PipelineContext::new(
1269            &PipelineDefinition::GreptimeIdentityPipeline(None),
1270            &params,
1271            Channel::Unknown,
1272        );
1273        {
1274            let array = [
1275                serde_json::json!({
1276                    "woshinull": null,
1277                    "name": "Alice",
1278                    "age": 20,
1279                    "is_student": true,
1280                    "score": 99.5,
1281                    "hobbies": "reading",
1282                    "address": "Beijing",
1283                }),
1284                serde_json::json!({
1285                    "name": "Bob",
1286                    "age": 21,
1287                    "is_student": false,
1288                    "score": "88.5",
1289                    "hobbies": "swimming",
1290                    "address": "Shanghai",
1291                    "gaga": "gaga"
1292                }),
1293            ];
1294            let array = array.iter().map(|v| v.into()).collect();
1295            let rows = identity_pipeline(array, None, &pipeline_ctx);
1296            assert!(rows.is_err());
1297            assert_eq!(
1298                rows.err().unwrap().to_string(),
1299                "Column datatype mismatch. For column: score, expected datatype: Float64, actual datatype: String".to_string(),
1300            );
1301        }
1302        {
1303            let array = [
1304                serde_json::json!({
1305                    "woshinull": null,
1306                    "name": "Alice",
1307                    "age": 20,
1308                    "is_student": true,
1309                    "score": 99.5,
1310                    "hobbies": "reading",
1311                    "address": "Beijing",
1312                }),
1313                serde_json::json!({
1314                    "name": "Bob",
1315                    "age": 21,
1316                    "is_student": false,
1317                    "score": 88,
1318                    "hobbies": "swimming",
1319                    "address": "Shanghai",
1320                    "gaga": "gaga"
1321                }),
1322            ];
1323            let array = array.iter().map(|v| v.into()).collect();
1324            let rows = identity_pipeline(array, None, &pipeline_ctx);
1325            assert!(rows.is_err());
1326            assert_eq!(
1327                rows.err().unwrap().to_string(),
1328                "Column datatype mismatch. For column: score, expected datatype: Float64, actual datatype: Int64".to_string(),
1329            );
1330        }
1331        {
1332            let array = [
1333                serde_json::json!({
1334                    "woshinull": null,
1335                    "name": "Alice",
1336                    "age": 20,
1337                    "is_student": true,
1338                    "score": 99.5,
1339                    "hobbies": "reading",
1340                    "address": "Beijing",
1341                }),
1342                serde_json::json!({
1343                    "name": "Bob",
1344                    "age": 21,
1345                    "is_student": false,
1346                    "score": 88.5,
1347                    "hobbies": "swimming",
1348                    "address": "Shanghai",
1349                    "gaga": "gaga"
1350                }),
1351            ];
1352            let array = array.iter().map(|v| v.into()).collect();
1353            let rows = identity_pipeline(array, None, &pipeline_ctx);
1354            assert!(rows.is_ok());
1355            let mut rows = rows.unwrap();
1356            assert!(rows.len() == 1);
1357            let rows = rows.remove(&ContextOpt::default()).unwrap();
1358            assert_eq!(rows.schema.len(), 8);
1359            assert_eq!(rows.rows.len(), 2);
1360            assert_eq!(8, rows.rows[0].values.len());
1361            assert_eq!(8, rows.rows[1].values.len());
1362        }
1363        {
1364            let array = [
1365                serde_json::json!({
1366                    "woshinull": null,
1367                    "name": "Alice",
1368                    "age": 20,
1369                    "is_student": true,
1370                    "score": 99.5,
1371                    "hobbies": "reading",
1372                    "address": "Beijing",
1373                }),
1374                serde_json::json!({
1375                    "name": "Bob",
1376                    "age": 21,
1377                    "is_student": false,
1378                    "score": 88.5,
1379                    "hobbies": "swimming",
1380                    "address": "Shanghai",
1381                    "gaga": "gaga"
1382                }),
1383            ];
1384            let tag_column_names = ["name".to_string(), "address".to_string()];
1385
1386            let rows = identity_pipeline_inner(
1387                array.iter().map(|v| v.into()).collect(),
1388                &pipeline_ctx,
1389                pipeline_ctx.pipeline_param.max_nested_levels(),
1390            )
1391            .map(|(mut schema, mut rows)| {
1392                for name in tag_column_names {
1393                    if let Some(index) = schema.index.get(&name) {
1394                        schema.schema[*index].semantic_type = SemanticType::Tag;
1395                    }
1396                }
1397
1398                assert!(rows.len() == 1);
1399                let rows = rows.remove(&ContextOpt::default()).unwrap();
1400
1401                Rows {
1402                    schema: schema.column_schemas().unwrap(),
1403                    rows,
1404                }
1405            });
1406
1407            assert!(rows.is_ok());
1408            let rows = rows.unwrap();
1409            assert_eq!(rows.schema.len(), 8);
1410            assert_eq!(rows.rows.len(), 2);
1411            assert_eq!(8, rows.rows[0].values.len());
1412            assert_eq!(8, rows.rows[1].values.len());
1413            assert_eq!(
1414                rows.schema
1415                    .iter()
1416                    .find(|x| x.column_name == "name")
1417                    .unwrap()
1418                    .semantic_type,
1419                SemanticType::Tag as i32
1420            );
1421            assert_eq!(
1422                rows.schema
1423                    .iter()
1424                    .find(|x| x.column_name == "address")
1425                    .unwrap()
1426                    .semantic_type,
1427                SemanticType::Tag as i32
1428            );
1429            assert_eq!(
1430                rows.schema
1431                    .iter()
1432                    .filter(|x| x.semantic_type == SemanticType::Tag as i32)
1433                    .count(),
1434                2
1435            );
1436        }
1437    }
1438
1439    #[test]
1440    fn test_flatten() {
1441        let test_cases = vec![
1442            // Basic case.
1443            (
1444                serde_json::json!(
1445                    {
1446                        "a": {
1447                            "b": {
1448                                "c": [1, 2, 3]
1449                            }
1450                        },
1451                        "d": [
1452                            "foo",
1453                            "bar"
1454                        ],
1455                        "e": {
1456                            "f": [7, 8, 9],
1457                            "g": {
1458                                "h": 123,
1459                                "i": "hello",
1460                                "j": {
1461                                    "k": true
1462                                }
1463                            }
1464                        }
1465                    }
1466                ),
1467                10,
1468                Some(serde_json::json!(
1469                    {
1470                        "a.b.c": "[1,2,3]",
1471                        "d": "[\"foo\",\"bar\"]",
1472                        "e.f": "[7,8,9]",
1473                        "e.g.h": 123,
1474                        "e.g.i": "hello",
1475                        "e.g.j.k": true
1476                    }
1477                )),
1478            ),
1479            // Test the case where the object has more than 3 nested levels.
1480            (
1481                serde_json::json!(
1482                    {
1483                        "a": {
1484                            "b": {
1485                                "c": {
1486                                    "d": [1, 2, 3]
1487                                }
1488                            }
1489                        },
1490                        "e": [
1491                            "foo",
1492                            "bar"
1493                        ]
1494                    }
1495                ),
1496                3,
1497                Some(serde_json::json!(
1498                    {
1499                        "a.b.c": "{\"d\":[1,2,3]}",
1500                        "e": "[\"foo\",\"bar\"]"
1501                    }
1502                )),
1503            ),
1504        ];
1505
1506        for (input, max_depth, expected) in test_cases {
1507            let input = input.into();
1508            let expected = expected.map(|e| e.into());
1509
1510            let flattened_object = flatten_object(input, max_depth).ok();
1511            assert_eq!(flattened_object, expected);
1512        }
1513    }
1514
1515    #[test]
1516    fn test_identity_pipeline_skip_error_flattens_valid_rows() {
1517        let params = GreptimePipelineParams::from_map(ahash::HashMap::from_iter([(
1518            "skip_error".to_string(),
1519            "true".to_string(),
1520        )]));
1521        let pipeline_def = PipelineDefinition::GreptimeIdentityPipeline(None);
1522        let pipeline_ctx = PipelineContext::new(&pipeline_def, &params, Channel::Unknown);
1523        let array = vec![
1524            serde_json::json!({
1525                "service": "frontend",
1526                "nested": {
1527                    "status": 200,
1528                    "path": "/v1/ingest"
1529                },
1530                "labels": ["pipeline", "identity"]
1531            })
1532            .into(),
1533            VrlValue::Bytes("invalid_string".into()),
1534            serde_json::json!({
1535                "service": "frontend",
1536                "nested": {
1537                    "status": 201,
1538                    "path": "/v1/ingest"
1539                },
1540                "labels": ["pipeline", "identity"]
1541            })
1542            .into(),
1543        ];
1544
1545        let mut rows_by_opt = identity_pipeline(array, None, &pipeline_ctx).unwrap();
1546        let rows = rows_by_opt.remove(&ContextOpt::default()).unwrap();
1547
1548        assert_eq!(rows.rows.len(), 2);
1549        assert_eq!(rows.schema.len(), rows.rows[0].values.len());
1550        assert!(rows.schema.iter().any(|s| s.column_name == "nested.status"));
1551        assert!(rows.schema.iter().any(|s| s.column_name == "nested.path"));
1552        assert!(rows.schema.iter().any(|s| s.column_name == "labels"));
1553    }
1554
1555    use ahash::HashMap as AHashMap;
1556    #[test]
1557    fn test_values_to_rows_skip_error_handling() {
1558        let table_suffix_template: Option<crate::tablesuffix::TableSuffixTemplate> = None;
1559
1560        // Case 1: skip_error=true, mixed valid/invalid elements
1561        {
1562            let schema_info = &mut SchemaInfo::default();
1563            let input_array = vec![
1564                // Valid object
1565                serde_json::json!({"name": "Alice", "age": 25}).into(),
1566                // Invalid element (string)
1567                VrlValue::Bytes("invalid_string".into()),
1568                // Valid object
1569                serde_json::json!({"name": "Bob", "age": 30}).into(),
1570                // Invalid element (number)
1571                VrlValue::Integer(42),
1572                // Valid object
1573                serde_json::json!({"name": "Charlie", "age": 35}).into(),
1574            ];
1575
1576            let params = GreptimePipelineParams::from_map(AHashMap::from_iter([(
1577                "skip_error".to_string(),
1578                "true".to_string(),
1579            )]));
1580
1581            let pipeline_ctx = PipelineContext::new(
1582                &PipelineDefinition::GreptimeIdentityPipeline(None),
1583                &params,
1584                Channel::Unknown,
1585            );
1586
1587            let result = values_to_rows(
1588                schema_info,
1589                VrlValue::Array(input_array),
1590                &pipeline_ctx,
1591                None,
1592                true,
1593                table_suffix_template.as_ref(),
1594            );
1595
1596            // Should succeed and only process valid objects
1597            assert!(result.is_ok());
1598            let rows_by_context = result.unwrap();
1599            // Count total rows across all ContextOpt groups
1600            let total_rows: usize = rows_by_context.values().map(|v| v.len()).sum();
1601            assert_eq!(total_rows, 3); // Only 3 valid objects
1602        }
1603
1604        // Case 2: skip_error=false, invalid elements present
1605        {
1606            let schema_info = &mut SchemaInfo::default();
1607            let input_array = vec![
1608                serde_json::json!({"name": "Alice", "age": 25}).into(),
1609                VrlValue::Bytes("invalid_string".into()), // This should cause error
1610            ];
1611
1612            let params = GreptimePipelineParams::default(); // skip_error = false
1613
1614            let pipeline_ctx = PipelineContext::new(
1615                &PipelineDefinition::GreptimeIdentityPipeline(None),
1616                &params,
1617                Channel::Unknown,
1618            );
1619
1620            let result = values_to_rows(
1621                schema_info,
1622                VrlValue::Array(input_array),
1623                &pipeline_ctx,
1624                None,
1625                true,
1626                table_suffix_template.as_ref(),
1627            );
1628
1629            // Should fail with ArrayElementMustBeObject error
1630            assert!(result.is_err());
1631            let error_msg = result.unwrap_err().to_string();
1632            assert!(error_msg.contains("Array element at index 1 must be an object for one-to-many transformation, got string"));
1633        }
1634    }
1635
1636    /// Test that values_to_rows correctly groups rows by per-element ContextOpt
1637    #[test]
1638    fn test_values_to_rows_per_element_context_opt() {
1639        let table_suffix_template: Option<crate::tablesuffix::TableSuffixTemplate> = None;
1640        let schema_info = &mut SchemaInfo::default();
1641
1642        // Create array with elements having different TTL values (ContextOpt)
1643        let input_array = vec![
1644            serde_json::json!({"name": "Alice", "greptime_ttl": "1h"}).into(),
1645            serde_json::json!({"name": "Bob", "greptime_ttl": "1h"}).into(),
1646            serde_json::json!({"name": "Charlie", "greptime_ttl": "24h"}).into(),
1647        ];
1648
1649        let params = GreptimePipelineParams::default();
1650        let pipeline_ctx = PipelineContext::new(
1651            &PipelineDefinition::GreptimeIdentityPipeline(None),
1652            &params,
1653            Channel::Unknown,
1654        );
1655
1656        let result = values_to_rows(
1657            schema_info,
1658            VrlValue::Array(input_array),
1659            &pipeline_ctx,
1660            None,
1661            true,
1662            table_suffix_template.as_ref(),
1663        );
1664
1665        assert!(result.is_ok());
1666        let rows_by_context = result.unwrap();
1667
1668        // Should have 2 different ContextOpt groups (1h TTL and 24h TTL)
1669        assert_eq!(rows_by_context.len(), 2);
1670
1671        // Count rows per group
1672        let total_rows: usize = rows_by_context.values().map(|v| v.len()).sum();
1673        assert_eq!(total_rows, 3);
1674
1675        // Verify that rows are correctly grouped by TTL
1676        let mut ttl_1h_count = 0;
1677        let mut ttl_24h_count = 0;
1678        for rows in rows_by_context.values() {
1679            // ContextOpt doesn't expose ttl directly, but we can count by group size
1680            if rows.len() == 2 {
1681                ttl_1h_count = rows.len();
1682            } else if rows.len() == 1 {
1683                ttl_24h_count = rows.len();
1684            }
1685        }
1686        assert_eq!(ttl_1h_count, 2); // Alice and Bob with 1h TTL
1687        assert_eq!(ttl_24h_count, 1); // Charlie with 24h TTL
1688    }
1689}