operator/req_convert/insert/
row_to_batch.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
15use std::collections::HashMap;
16
17use api::helper::{
18    ColumnDataTypeWrapper, pb_value_to_value_ref, proto_value_type, proto_value_type_match,
19};
20use api::v1::column_data_type_extension::TypeExt;
21use api::v1::value::ValueData;
22use api::v1::{ColumnDataType, ColumnDataTypeExtension, Rows, SemanticType, Value};
23use arrow::record_batch::RecordBatch;
24use common_error::ext::BoxedError;
25use datatypes::data_type::ConcreteDataType;
26use datatypes::extension::json::align_schema_with_json_array;
27use snafu::{OptionExt, ResultExt, ensure};
28use table::metadata::TableInfo;
29
30use crate::error::{self, Result};
31
32/// Converts prepared rows in target-schema order without Prom-specific renaming.
33///
34/// Missing columns evaluate their defaults once per conversion. Explicit nulls
35/// remain nulls. Semantic types are checked against the table time index and
36/// primary-key indices.
37///
38/// # Panics
39///
40/// Panics if `rows.rows` is empty. Callers must skip empty writes before conversion.
41pub fn rows_to_record_batch(rows: &Rows, table_info: &TableInfo) -> Result<RecordBatch> {
42    assert!(!rows.rows.is_empty(), "prepared rows must not be empty");
43    let schema = &table_info.meta.schema;
44    let mut source_columns = HashMap::with_capacity(rows.schema.len());
45    for (index, source) in rows.schema.iter().enumerate() {
46        ensure!(
47            source_columns
48                .insert(source.column_name.as_str(), index)
49                .is_none(),
50            error::InvalidInsertRequestSnafu {
51                reason: format!("Duplicate input column {}", source.column_name),
52            }
53        );
54        let target = schema
55            .column_schema_by_name(&source.column_name)
56            .with_context(|| error::InvalidInsertRequestSnafu {
57                reason: format!("Unknown input column {}", source.column_name),
58            })?;
59        let data_type =
60            ColumnDataTypeWrapper::try_new(source.datatype, source.datatype_extension.clone())
61                .map_err(BoxedError::new)
62                .context(error::ExternalSnafu)?;
63        ensure!(
64            ConcreteDataType::from(data_type) == target.data_type,
65            error::InvalidInsertRequestSnafu {
66                reason: format!("Input datatype differs for column {}", source.column_name),
67            }
68        );
69        let semantic = SemanticType::try_from(source.semantic_type)
70            .ok()
71            .with_context(|| error::InvalidInsertRequestSnafu {
72                reason: format!("Invalid semantic type for column {}", source.column_name),
73            })?;
74        let is_tag = table_info.meta.primary_key_indices.iter().any(|&index| {
75            schema
76                .column_schemas()
77                .get(index)
78                .is_some_and(|column| column.name == source.column_name)
79        });
80        let expected_semantic = if target.is_time_index() {
81            SemanticType::Timestamp
82        } else if is_tag {
83            SemanticType::Tag
84        } else {
85            SemanticType::Field
86        };
87        ensure!(
88            semantic == expected_semantic,
89            error::InvalidInsertRequestSnafu {
90                reason: format!("Input semantics differ for column {}", source.column_name),
91            }
92        );
93    }
94    for row in &rows.rows {
95        ensure!(
96            row.values.len() == rows.schema.len(),
97            error::InvalidInsertRequestSnafu {
98                reason: format!(
99                    "Expected {} values, got {}",
100                    rows.schema.len(),
101                    row.values.len()
102                ),
103            }
104        );
105    }
106
107    let mut arrays = Vec::with_capacity(schema.num_columns());
108    for column in schema.column_schemas() {
109        let vector = if let Some(&index) = source_columns.get(column.name.as_str()) {
110            let mut builder = column.create_mutable_vector(rows.rows.len());
111            for row in &rows.rows {
112                let value = &row.values[index];
113                ensure!(
114                    value.value_data.is_some() || column.is_nullable(),
115                    error::InvalidInsertRequestSnafu {
116                        reason: format!("Null supplied for non-nullable column {}", column.name),
117                    }
118                );
119                ensure!(
120                    value_matches_type(
121                        value,
122                        rows.schema[index].datatype,
123                        rows.schema[index].datatype_extension.as_ref(),
124                    ),
125                    error::InvalidInsertRequestSnafu {
126                        reason: format!("Value datatype differs for column {}", column.name),
127                    }
128                );
129                builder
130                    .try_push_value_ref(&pb_value_to_value_ref(
131                        value,
132                        rows.schema[index].datatype_extension.as_ref(),
133                    ))
134                    .map_err(BoxedError::new)
135                    .context(error::ExternalSnafu)?;
136            }
137            builder.to_vector()
138        } else {
139            column
140                .create_default_vector(rows.rows.len())
141                .map_err(BoxedError::new)
142                .context(error::ExternalSnafu)?
143                .with_context(|| error::InvalidInsertRequestSnafu {
144                    reason: format!("Missing required column {}", column.name),
145                })?
146        };
147        arrays.push(vector.to_arrow_array());
148    }
149    let arrow_schema = align_schema_with_json_array(schema.arrow_schema().clone(), &arrays);
150    RecordBatch::try_new(arrow_schema, arrays).context(error::ComputeArrowSnafu)
151}
152
153// Validate nested values before the infallible protobuf conversion reads type extensions.
154fn value_matches_type(
155    value: &Value,
156    datatype: i32,
157    extension: Option<&ColumnDataTypeExtension>,
158) -> bool {
159    let Some(value_type) = proto_value_type(value) else {
160        return true;
161    };
162    let Ok(column_type) = ColumnDataType::try_from(datatype) else {
163        return false;
164    };
165    if !proto_value_type_match(column_type, value_type) {
166        return false;
167    }
168    match value.value_data.as_ref() {
169        Some(ValueData::ListValue(list)) => {
170            let Some(TypeExt::ListType(item)) = extension.and_then(|ext| ext.type_ext.as_ref())
171            else {
172                return false;
173            };
174            list.items.iter().all(|value| {
175                value_matches_type(value, item.datatype, item.datatype_extension.as_deref())
176            })
177        }
178        Some(ValueData::StructValue(value)) => {
179            let Some(TypeExt::StructType(schema)) = extension.and_then(|ext| ext.type_ext.as_ref())
180            else {
181                return false;
182            };
183            value.items.len() == schema.fields.len()
184                && value
185                    .items
186                    .iter()
187                    .zip(&schema.fields)
188                    .all(|(value, field)| {
189                        value_matches_type(value, field.datatype, field.datatype_extension.as_ref())
190                    })
191        }
192        _ => true,
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use std::sync::Arc;
199    use std::thread::yield_now;
200
201    use api::v1::value::ValueData;
202    use api::v1::{ColumnDataType, ColumnSchema as ProtoColumnSchema, Row, Value};
203    use arrow::array::{Array, Int32Array, StringArray};
204    use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema};
205    use datatypes::value::Value as DtValue;
206    use table::metadata::{TableInfoBuilder, TableMetaBuilder};
207
208    use crate::req_convert::insert::row_to_batch::*;
209
210    fn source(name: &str, datatype: ColumnDataType) -> ProtoColumnSchema {
211        ProtoColumnSchema {
212            column_name: name.to_string(),
213            datatype: datatype as i32,
214            semantic_type: SemanticType::Field as i32,
215            ..Default::default()
216        }
217    }
218
219    fn table_info(schema: Schema) -> TableInfo {
220        let next_column_id = schema.num_columns() as u32;
221        TableInfoBuilder::default()
222            .table_id(1)
223            .table_version(0)
224            .name("test")
225            .meta(
226                TableMetaBuilder::empty()
227                    .schema(Arc::new(schema))
228                    .primary_key_indices(vec![])
229                    .next_column_id(next_column_id)
230                    .engine("mito")
231                    .build()
232                    .unwrap(),
233            )
234            .build()
235            .unwrap()
236    }
237
238    fn fixture() -> (Rows, Schema) {
239        let schema = Schema::new(vec![
240            ColumnSchema::new("count", ConcreteDataType::int32_datatype(), false),
241            ColumnSchema::new("label", ConcreteDataType::string_datatype(), true),
242            ColumnSchema::new("fallback", ConcreteDataType::int32_datatype(), true)
243                .with_default_constraint(Some(ColumnDefaultConstraint::Value(DtValue::Int32(7))))
244                .unwrap(),
245        ]);
246        let rows = Rows {
247            schema: vec![
248                source("label", ColumnDataType::String),
249                source("count", ColumnDataType::Int32),
250            ],
251            rows: vec![Row {
252                values: vec![
253                    Value::default(),
254                    Value {
255                        value_data: Some(ValueData::I32Value(3)),
256                    },
257                ],
258            }],
259        };
260        (rows, schema)
261    }
262
263    #[test]
264    fn test_reorder_multifield_and_defaults() {
265        let (rows, schema) = fixture();
266        let batch = rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap();
267        assert_eq!(batch.schema(), schema.arrow_schema().clone());
268        assert_eq!(
269            batch
270                .column(0)
271                .as_any()
272                .downcast_ref::<Int32Array>()
273                .unwrap()
274                .value(0),
275            3
276        );
277        assert!(
278            batch
279                .column(1)
280                .as_any()
281                .downcast_ref::<StringArray>()
282                .unwrap()
283                .is_null(0)
284        );
285        assert_eq!(
286            batch
287                .column(2)
288                .as_any()
289                .downcast_ref::<Int32Array>()
290                .unwrap()
291                .value(0),
292            7
293        );
294    }
295
296    #[test]
297    fn test_explicit_null_does_not_use_default() {
298        let (mut rows, schema) = fixture();
299        rows.schema.push(source("fallback", ColumnDataType::Int32));
300        rows.rows[0].values.push(Value::default());
301        assert!(
302            rows_to_record_batch(&rows, &table_info(schema.clone()))
303                .unwrap()
304                .column(2)
305                .is_null(0)
306        );
307    }
308
309    #[test]
310    fn test_timestamp_precisions() {
311        for (datatype, value) in [
312            (
313                ColumnDataType::TimestampSecond,
314                ValueData::TimestampSecondValue(42),
315            ),
316            (
317                ColumnDataType::TimestampMillisecond,
318                ValueData::TimestampMillisecondValue(42),
319            ),
320            (
321                ColumnDataType::TimestampMicrosecond,
322                ValueData::TimestampMicrosecondValue(42),
323            ),
324            (
325                ColumnDataType::TimestampNanosecond,
326                ValueData::TimestampNanosecondValue(42),
327            ),
328        ] {
329            let target_type = ConcreteDataType::from(ColumnDataTypeWrapper::new(datatype, None));
330            let schema = Schema::new(vec![
331                ColumnSchema::new("ts", target_type, false).with_time_index(true),
332            ]);
333            let mut column = source("ts", datatype);
334            column.semantic_type = SemanticType::Timestamp as i32;
335            let rows = Rows {
336                schema: vec![column],
337                rows: vec![Row {
338                    values: vec![Value {
339                        value_data: Some(value),
340                    }],
341                }],
342            };
343            let batch = rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap();
344            assert_eq!(batch.schema(), schema.arrow_schema().clone());
345            assert_eq!(batch.num_rows(), 1);
346            assert_eq!(batch.column(0).null_count(), 0);
347        }
348    }
349
350    #[test]
351    fn test_json2_expanded_schema() {
352        use api::helper::to_grpc_value;
353        use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
354        use datatypes::json::JsonSettings;
355        use datatypes::schema::SchemaBuilder;
356        use datatypes::types::json_type::JsonNativeType;
357
358        let settings = JsonSettings::default();
359        let mut column_schema = ColumnSchema::new(
360            "data",
361            ConcreteDataType::json2(JsonNativeType::object()),
362            true,
363        );
364        column_schema.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
365            settings.clone(),
366        ))));
367        let schema = SchemaBuilder::try_from(vec![column_schema])
368            .unwrap()
369            .add_metadata("test", "metadata")
370            .build()
371            .unwrap();
372        let arrow_schema = schema.arrow_schema().clone();
373        let field = arrow_schema.field(0);
374        let datatype =
375            ColumnDataTypeWrapper::try_from(schema.column_schemas()[0].data_type.clone()).unwrap();
376        let (kind, extension) = datatype.to_parts();
377        let mut column = source("data", kind);
378        column.datatype_extension = extension;
379        let rows = Rows {
380            schema: vec![column],
381            rows: vec![Row {
382                values: vec![to_grpc_value(
383                    settings.encode(serde_json::json!({"id": 3})).unwrap(),
384                )],
385            }],
386        };
387        let batch = rows_to_record_batch(&rows, &table_info(schema)).unwrap();
388        // The same arrays fail with the static table schema used before alignment.
389        assert!(RecordBatch::try_new(arrow_schema.clone(), batch.columns().to_vec()).is_err());
390        let actual_schema = batch.schema();
391        assert_eq!(actual_schema.metadata(), arrow_schema.metadata());
392        assert_eq!(actual_schema.field(0).metadata(), field.metadata());
393        assert_eq!(actual_schema.field(0).name(), field.name());
394        assert_eq!(actual_schema.field(0).is_nullable(), field.is_nullable());
395        assert_eq!(
396            actual_schema.field(0).data_type(),
397            batch.column(0).data_type()
398        );
399        let array = batch
400            .column(0)
401            .as_any()
402            .downcast_ref::<arrow::array::StructArray>()
403            .unwrap();
404        assert!(array.column_by_name("id").is_some());
405        assert!(array.column_by_name("!__remainder__!").is_some());
406        assert_eq!(batch.num_rows(), 1);
407    }
408
409    #[test]
410    fn test_dynamic_default_is_evaluated_per_conversion() {
411        use std::time::{Duration, Instant};
412
413        use arrow::array::TimestampMillisecondArray;
414
415        let (mut rows, base_schema) = fixture();
416        let mut columns = base_schema.column_schemas().to_vec();
417        columns.push(
418            ColumnSchema::new(
419                "created",
420                ConcreteDataType::timestamp_millisecond_datatype(),
421                true,
422            )
423            .with_default_constraint(Some(ColumnDefaultConstraint::Function(
424                "current_timestamp()".to_string(),
425            )))
426            .unwrap(),
427        );
428        let schema = Schema::new(columns);
429        let timestamp = |batch: &RecordBatch| {
430            batch
431                .column(3)
432                .as_any()
433                .downcast_ref::<TimestampMillisecondArray>()
434                .unwrap()
435                .value(0)
436        };
437        let first = timestamp(&rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap());
438        // The default uses wall-clock milliseconds, not Tokio's controllable
439        // clock. Poll its observable result with a bound instead of sleeping.
440        let deadline = Instant::now() + Duration::from_secs(1);
441        loop {
442            let next =
443                timestamp(&rows_to_record_batch(&rows, &table_info(schema.clone())).unwrap());
444            if next != first {
445                break;
446            }
447            assert!(
448                Instant::now() < deadline,
449                "dynamic default was not refreshed"
450            );
451            yield_now();
452        }
453        rows.schema
454            .push(source("created", ColumnDataType::TimestampMillisecond));
455        rows.rows[0].values.push(Value::default());
456        assert!(
457            rows_to_record_batch(&rows, &table_info(schema.clone()))
458                .unwrap()
459                .column(3)
460                .is_null(0)
461        );
462    }
463
464    #[test]
465    fn test_tag_and_field_semantics() {
466        let (mut rows, schema) = fixture();
467        let mut table = table_info(schema);
468        table.meta.primary_key_indices = vec![1];
469        // The "label" column is a tag in table metadata, not a field.
470        assert!(rows_to_record_batch(&rows, &table).is_err());
471        rows.schema[0].semantic_type = SemanticType::Tag as i32;
472        assert!(rows_to_record_batch(&rows, &table).is_ok());
473        // Conversely the count field cannot be submitted as a tag.
474        rows.schema[1].semantic_type = SemanticType::Tag as i32;
475        assert!(rows_to_record_batch(&rows, &table).is_err());
476    }
477
478    #[test]
479    fn test_invalid_input() {
480        for case in 0..7 {
481            let (mut rows, schema) = fixture();
482            match case {
483                0 => rows.schema[0].column_name = "unknown".to_string(),
484                1 => rows.schema[0].column_name = "count".to_string(),
485                2 => {
486                    rows.rows[0].values.pop();
487                }
488                3 => rows.schema[1].datatype = ColumnDataType::Float64 as i32,
489                4 => rows.rows[0].values[1] = Value::default(),
490                5 => rows.schema[1].semantic_type = SemanticType::Timestamp as i32,
491                _ => {
492                    rows.rows[0].values[1].value_data =
493                        Some(ValueData::StringValue("wrong".to_string()))
494                }
495            }
496            assert!(
497                rows_to_record_batch(&rows, &table_info(schema.clone())).is_err(),
498                "case {case}"
499            );
500        }
501    }
502
503    #[test]
504    #[should_panic(expected = "prepared rows must not be empty")]
505    fn test_empty_rows_contract() {
506        let (mut rows, schema) = fixture();
507        rows.rows.clear();
508        let _ = rows_to_record_batch(&rows, &table_info(schema));
509    }
510
511    #[test]
512    fn test_reject_nested_values_for_scalar_column() {
513        for value in [
514            ValueData::ListValue(api::v1::ListValue { items: vec![] }),
515            ValueData::StructValue(api::v1::StructValue { items: vec![] }),
516        ] {
517            let (mut rows, schema) = fixture();
518            rows.rows[0].values[1].value_data = Some(value);
519            assert!(matches!(
520                rows_to_record_batch(&rows, &table_info(schema)),
521                Err(error::Error::InvalidInsertRequest { .. })
522            ));
523        }
524    }
525
526    #[test]
527    fn test_nested_value_validation() {
528        let int = Value {
529            value_data: Some(ValueData::I32Value(1)),
530        };
531        let list = Value {
532            value_data: Some(ValueData::ListValue(api::v1::ListValue {
533                items: vec![int.clone()],
534            })),
535        };
536        let structure = Value {
537            value_data: Some(ValueData::StructValue(api::v1::StructValue {
538                items: vec![int.clone()],
539            })),
540        };
541        let types = [
542            (
543                ColumnDataTypeWrapper::list_datatype(ColumnDataTypeWrapper::int32_datatype()),
544                list.clone(),
545            ),
546            (
547                ColumnDataTypeWrapper::struct_datatype(vec![(
548                    "count".to_string(),
549                    ColumnDataTypeWrapper::int32_datatype(),
550                )]),
551                structure.clone(),
552            ),
553        ];
554        for (datatype, valid) in types {
555            let (kind, extension) = datatype.to_parts();
556            let schema = Schema::new(vec![ColumnSchema::new("nested", datatype.into(), true)]);
557            let mut column = source("nested", kind);
558            column.datatype_extension = extension;
559            let mut rows = Rows {
560                schema: vec![column],
561                rows: vec![Row {
562                    values: vec![valid],
563                }],
564            };
565            let table = table_info(schema);
566            assert_eq!(rows_to_record_batch(&rows, &table).unwrap().num_rows(), 1);
567            for invalid in [list.clone(), structure.clone()] {
568                match rows.rows[0].values[0].value_data.as_mut().unwrap() {
569                    ValueData::ListValue(v) => v.items = vec![invalid],
570                    ValueData::StructValue(v) => v.items = vec![invalid],
571                    _ => unreachable!(),
572                }
573                assert!(matches!(
574                    rows_to_record_batch(&rows, &table),
575                    Err(error::Error::InvalidInsertRequest { .. })
576                ));
577            }
578        }
579        let datatype = ColumnDataTypeWrapper::list_datatype(ColumnDataTypeWrapper::list_datatype(
580            ColumnDataTypeWrapper::int32_datatype(),
581        ));
582        let (kind, extension) = datatype.to_parts();
583        let value = Value {
584            value_data: Some(ValueData::ListValue(api::v1::ListValue {
585                items: vec![list],
586            })),
587        };
588        assert!(value_matches_type(&value, kind as i32, extension.as_ref()));
589        assert!(!value_matches_type(&value, kind as i32, None));
590    }
591}
operator/req_convert/insert/row_to_batch.rs

operator/req_convert/insert/
row_to_batch.rs
