servers/
row_writer.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;
16use std::sync::Arc;
17
18use ahash::{HashMap as AHashMap, HashMapExt};
19use api::helper::encode_json_value;
20use api::v1::column_data_type_extension::TypeExt;
21use api::v1::helper::time_index_column_schema;
22use api::v1::value::ValueData;
23use api::v1::{
24    ColumnDataType, ColumnDataTypeExtension, ColumnOptions, ColumnSchema, JsonTypeExtension, Row,
25    RowInsertRequest, RowInsertRequests, Rows, SemanticType, Value,
26};
27use arrow_schema::extension::{
28    EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
29};
30use common_grpc::precision::Precision;
31use common_time::Timestamp;
32use common_time::timestamp::TimeUnit;
33use common_time::timestamp::TimeUnit::Nanosecond;
34use datatypes::extension::json::{Json2ExtensionType, JsonMetadata};
35use datatypes::json::JsonSettings;
36use datatypes::value::Value as DataValue;
37use serde::Serialize;
38use snafu::{OptionExt, ResultExt, ensure};
39
40use crate::error::{
41    ConvertScalarValueSnafu, IncompatibleSchemaSnafu, InternalSnafu, Result, RowWriterSnafu,
42    TimePrecisionSnafu, TimestampOverflowSnafu, ToJsonSnafu,
43};
44
45/// The intermediate data structure for building the write request.
46/// It constructs the `schema` and `rows` as all input data row
47/// parsing is completed.
48pub struct TableData {
49    schema: Vec<ColumnSchema>,
50    rows: Vec<Row>,
51    column_indexes: AHashMap<String, usize>,
52}
53
54impl TableData {
55    pub fn new(num_columns: usize, num_rows: usize) -> Self {
56        Self {
57            schema: Vec::with_capacity(num_columns),
58            rows: Vec::with_capacity(num_rows),
59            column_indexes: AHashMap::with_capacity(num_columns),
60        }
61    }
62
63    #[inline]
64    pub fn num_columns(&self) -> usize {
65        self.schema.len()
66    }
67
68    #[inline]
69    pub fn num_rows(&self) -> usize {
70        self.rows.len()
71    }
72
73    #[inline]
74    pub fn alloc_one_row(&self) -> Vec<Value> {
75        vec![Value { value_data: None }; self.num_columns()]
76    }
77
78    #[inline]
79    pub fn add_row(&mut self, values: Vec<Value>) {
80        self.rows.push(Row { values })
81    }
82
83    #[inline]
84    pub fn reserve_rows(&mut self, additional: usize) {
85        self.rows.reserve(additional);
86    }
87
88    pub(crate) fn ensure_column(&mut self, column_schema: ColumnSchema) -> Result<usize> {
89        if let Some(index) = self.column_indexes.get(&column_schema.column_name).copied() {
90            check_schema_number(
91                column_schema.datatype,
92                column_schema.semantic_type,
93                &self.schema[index],
94            )?;
95            return Ok(index);
96        }
97
98        let index = self.schema.len();
99        let name = column_schema.column_name.clone();
100        self.schema.push(column_schema);
101        self.column_indexes.insert(name, index);
102        Ok(index)
103    }
104
105    /// Ensures a default column schema without allocating its name when it already exists.
106    pub(crate) fn ensure_column_by_name(
107        &mut self,
108        name: &str,
109        datatype: ColumnDataType,
110        semantic_type: SemanticType,
111    ) -> Result<usize> {
112        if let Some(index) = self.column_indexes.get(name).copied() {
113            check_schema(datatype, semantic_type, &self.schema[index])?;
114            return Ok(index);
115        }
116
117        self.ensure_column(ColumnSchema {
118            column_name: name.to_string(),
119            datatype: datatype as i32,
120            semantic_type: semantic_type as i32,
121            ..Default::default()
122        })
123    }
124
125    #[allow(dead_code)]
126    pub fn columns(&self) -> &Vec<ColumnSchema> {
127        &self.schema
128    }
129
130    pub fn into_schema_and_rows(self) -> (Vec<ColumnSchema>, Vec<Row>) {
131        (self.schema, self.rows)
132    }
133
134    pub(crate) fn into_rows(self) -> Rows {
135        let (schema, mut rows) = self.into_schema_and_rows();
136        for row in &mut rows {
137            if row.values.len() < schema.len() {
138                row.values.resize(schema.len(), Value::default());
139            }
140        }
141        Rows { schema, rows }
142    }
143
144    /// Writes a field value without enforcing that later writes use the same datatype
145    /// as the first-seen schema entry.
146    ///
147    /// The OTLP trace v1 path uses this to preserve raw mixed values inside one request
148    /// so the frontend can reconcile them later against both the full batch and the
149    /// existing table schema.
150    pub fn write_field_unchecked(
151        &mut self,
152        name: impl ToString,
153        datatype: ColumnDataType,
154        value: Option<ValueData>,
155        one_row: &mut Vec<Value>,
156    ) {
157        self.write_column_unchecked(
158            ColumnSchema {
159                column_name: name.to_string(),
160                datatype: datatype as i32,
161                semantic_type: SemanticType::Field as i32,
162                ..Default::default()
163            },
164            value,
165            one_row,
166        );
167    }
168
169    pub fn write_column_unchecked(
170        &mut self,
171        column_schema: ColumnSchema,
172        value: Option<ValueData>,
173        one_row: &mut Vec<Value>,
174    ) {
175        if let Some(index) = self.column_indexes.get(&column_schema.column_name).copied() {
176            one_row[index].value_data = value;
177        } else {
178            let index = self.schema.len();
179            let name = column_schema.column_name.clone();
180            self.schema.push(column_schema);
181            self.column_indexes.insert(name, index);
182            one_row.push(Value { value_data: value });
183        }
184    }
185}
186
187pub struct MultiTableData {
188    table_data_map: HashMap<String, TableData>,
189}
190
191impl Default for MultiTableData {
192    fn default() -> Self {
193        Self::new()
194    }
195}
196
197impl MultiTableData {
198    pub fn new() -> Self {
199        Self {
200            table_data_map: HashMap::new(),
201        }
202    }
203
204    pub fn get_or_default_table_data(
205        &mut self,
206        table_name: impl ToString,
207        num_columns: usize,
208        num_rows: usize,
209    ) -> &mut TableData {
210        self.table_data_map
211            .entry(table_name.to_string())
212            .or_insert_with(|| TableData::new(num_columns, num_rows))
213    }
214
215    pub fn add_table_data(&mut self, table_name: impl ToString, table_data: TableData) {
216        self.table_data_map
217            .insert(table_name.to_string(), table_data);
218    }
219
220    pub(crate) fn iter(&self) -> impl Iterator<Item = (&str, &TableData)> {
221        self.table_data_map
222            .iter()
223            .map(|(name, table)| (name.as_str(), table))
224    }
225
226    /// Retains schemas and row capacity for the next point with the same attributes.
227    pub(crate) fn drain_rows(
228        &mut self,
229    ) -> impl Iterator<Item = (&str, &[ColumnSchema], std::vec::Drain<'_, Row>)> {
230        self.table_data_map.iter_mut().filter_map(|(name, table)| {
231            if table.rows.is_empty() {
232                return None;
233            }
234            for row in &mut table.rows {
235                if row.values.len() < table.schema.len() {
236                    row.values.resize(table.schema.len(), Value::default());
237                }
238            }
239            Some((name.as_str(), table.schema.as_slice(), table.rows.drain(..)))
240        })
241    }
242
243    #[allow(dead_code)]
244    pub fn num_tables(&self) -> usize {
245        self.table_data_map.len()
246    }
247
248    /// Returns the request and number of rows in it.
249    pub fn into_row_insert_requests(self) -> (RowInsertRequests, usize) {
250        let mut total_rows = 0;
251        let inserts = self
252            .table_data_map
253            .into_iter()
254            .map(|(table_name, table_data)| {
255                total_rows += table_data.num_rows();
256                RowInsertRequest {
257                    table_name,
258                    rows: Some(table_data.into_rows()),
259                }
260            })
261            .collect::<Vec<_>>();
262        let row_insert_requests = RowInsertRequests { inserts };
263
264        (row_insert_requests, total_rows)
265    }
266}
267
268/// Write data as tags into the table data.
269pub fn write_tags<K>(
270    table_data: &mut TableData,
271    tags: impl Iterator<Item = (K, String)>,
272    one_row: &mut Vec<Value>,
273) -> Result<()>
274where
275    K: AsRef<str> + Into<String>,
276{
277    let ktv_iter = tags.map(|(k, v)| (k, ColumnDataType::String, Some(ValueData::StringValue(v))));
278    write_by_semantic_type(table_data, SemanticType::Tag, ktv_iter, one_row)
279}
280
281/// Write data as fields into the table data.
282pub fn write_fields(
283    table_data: &mut TableData,
284    fields: impl Iterator<Item = (String, ColumnDataType, Option<ValueData>)>,
285    one_row: &mut Vec<Value>,
286) -> Result<()> {
287    write_by_semantic_type(table_data, SemanticType::Field, fields, one_row)
288}
289
290/// Write data as a tag into the table data.
291pub fn write_tag(
292    table_data: &mut TableData,
293    name: impl ToString,
294    value: impl ToString,
295    one_row: &mut Vec<Value>,
296) -> Result<()> {
297    write_by_semantic_type(
298        table_data,
299        SemanticType::Tag,
300        std::iter::once((
301            name.to_string(),
302            ColumnDataType::String,
303            Some(ValueData::StringValue(value.to_string())),
304        )),
305        one_row,
306    )
307}
308
309/// Write float64 data as a field into the table data.
310pub fn write_f64(
311    table_data: &mut TableData,
312    name: impl ToString,
313    value: f64,
314    one_row: &mut Vec<Value>,
315) -> Result<()> {
316    write_fields(
317        table_data,
318        std::iter::once((
319            name.to_string(),
320            ColumnDataType::Float64,
321            Some(ValueData::F64Value(value)),
322        )),
323        one_row,
324    )
325}
326
327pub(crate) fn build_json_column_schema(name: impl ToString) -> ColumnSchema {
328    ColumnSchema {
329        column_name: name.to_string(),
330        datatype: ColumnDataType::Binary as i32,
331        semantic_type: SemanticType::Field as i32,
332        datatype_extension: Some(ColumnDataTypeExtension {
333            type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())),
334        }),
335        ..Default::default()
336    }
337}
338
339pub fn write_json(
340    table_data: &mut TableData,
341    name: impl ToString,
342    value: jsonb::Value,
343    one_row: &mut Vec<Value>,
344) -> Result<()> {
345    write_by_schema(
346        table_data,
347        std::iter::once((
348            build_json_column_schema(name),
349            Some(ValueData::BinaryValue(value.to_vec())),
350        )),
351        one_row,
352    )
353}
354
355pub(crate) fn build_json2_column_schema(name: impl ToString) -> ColumnSchema {
356    let extension = Json2ExtensionType::new(Arc::new(JsonMetadata::new(JsonSettings::new_v2())));
357    let mut options = ColumnOptions::default();
358    options.options.insert(
359        EXTENSION_TYPE_NAME_KEY.to_string(),
360        Json2ExtensionType::NAME.to_string(),
361    );
362    if let Some(metadata) = extension.serialize_metadata() {
363        options
364            .options
365            .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata);
366    }
367
368    ColumnSchema {
369        column_name: name.to_string(),
370        datatype: ColumnDataType::Json as i32,
371        semantic_type: SemanticType::Field as i32,
372        options: Some(options),
373        ..Default::default()
374    }
375}
376
377/// Encodes a JSON2 field without constructing its column schema.
378pub(crate) fn encode_json2(value: impl Serialize) -> Result<ValueData> {
379    let json = serde_json::to_value(value).context(ToJsonSnafu)?;
380    let value = JsonSettings::new_v2()
381        .encode(json)
382        .context(ConvertScalarValueSnafu)?;
383    let DataValue::Json(value) = value else {
384        return InternalSnafu {
385            err_msg: "JSON2 encoding returned a non-JSON value",
386        }
387        .fail();
388    };
389    Ok(ValueData::JsonValue(encode_json_value(*value)))
390}
391
392pub(crate) fn write_by_schema(
393    table_data: &mut TableData,
394    kv_iter: impl Iterator<Item = (ColumnSchema, Option<ValueData>)>,
395    one_row: &mut Vec<Value>,
396) -> Result<()> {
397    let TableData {
398        schema,
399        column_indexes,
400        ..
401    } = table_data;
402
403    for (column_schema, value) in kv_iter {
404        let index = column_indexes.get(&column_schema.column_name);
405        if let Some(index) = index {
406            check_schema_number(
407                column_schema.datatype,
408                column_schema.semantic_type,
409                &schema[*index],
410            )?;
411            one_row[*index].value_data = value;
412        } else {
413            let index = schema.len();
414            let key = column_schema.column_name.clone();
415            schema.push(column_schema);
416            column_indexes.insert(key, index);
417            one_row.push(Value { value_data: value });
418        }
419    }
420
421    Ok(())
422}
423
424fn write_by_semantic_type<K>(
425    table_data: &mut TableData,
426    semantic_type: SemanticType,
427    ktv_iter: impl Iterator<Item = (K, ColumnDataType, Option<ValueData>)>,
428    one_row: &mut Vec<Value>,
429) -> Result<()>
430where
431    K: AsRef<str> + Into<String>,
432{
433    let TableData {
434        schema,
435        column_indexes,
436        ..
437    } = table_data;
438
439    for (name, datatype, value) in ktv_iter {
440        let index = column_indexes.get(name.as_ref()).copied();
441        if let Some(index) = index {
442            check_schema(datatype, semantic_type, &schema[index])?;
443            one_row[index].value_data = value;
444        } else {
445            let index = schema.len();
446            let name = name.into();
447            schema.push(ColumnSchema {
448                column_name: name.clone(),
449                datatype: datatype as i32,
450                semantic_type: semantic_type as i32,
451                ..Default::default()
452            });
453            column_indexes.insert(name, index);
454            one_row.push(Value { value_data: value });
455        }
456    }
457
458    Ok(())
459}
460
461/// Write timestamp data as milliseconds into the table data.
462pub fn write_ts_to_millis(
463    table_data: &mut TableData,
464    name: impl ToString,
465    ts: Option<i64>,
466    precision: Precision,
467    one_row: &mut Vec<Value>,
468) -> Result<()> {
469    write_ts_to(
470        table_data,
471        name,
472        ts,
473        precision,
474        TimestampType::Millis,
475        one_row,
476    )
477}
478
479/// Write timestamp data as nanoseconds into the table data.
480pub fn write_ts_to_nanos(
481    table_data: &mut TableData,
482    name: impl ToString,
483    ts: Option<i64>,
484    precision: Precision,
485    one_row: &mut Vec<Value>,
486) -> Result<()> {
487    write_ts_to(
488        table_data,
489        name,
490        ts,
491        precision,
492        TimestampType::Nanos,
493        one_row,
494    )
495}
496
497enum TimestampType {
498    Millis,
499    Nanos,
500}
501
502fn write_ts_to(
503    table_data: &mut TableData,
504    name: impl ToString,
505    ts: Option<i64>,
506    precision: Precision,
507    ts_type: TimestampType,
508    one_row: &mut Vec<Value>,
509) -> Result<()> {
510    let TableData {
511        schema,
512        column_indexes,
513        ..
514    } = table_data;
515    let name = name.to_string();
516
517    let ts = match ts {
518        Some(timestamp) => match ts_type {
519            TimestampType::Millis => precision.to_millis(timestamp),
520            TimestampType::Nanos => precision.to_nanos(timestamp),
521        }
522        .with_context(|| TimestampOverflowSnafu {
523            error: format!(
524                "timestamp {} overflow with precision {}",
525                timestamp, precision
526            ),
527        })?,
528        None => {
529            let timestamp = Timestamp::current_time(Nanosecond);
530            let unit: TimeUnit = precision.try_into().context(RowWriterSnafu)?;
531            let timestamp = timestamp
532                .convert_to(unit)
533                .with_context(|| TimePrecisionSnafu {
534                    name: precision.to_string(),
535                })?
536                .into();
537            match ts_type {
538                TimestampType::Millis => precision.to_millis(timestamp),
539                TimestampType::Nanos => precision.to_nanos(timestamp),
540            }
541            .with_context(|| TimestampOverflowSnafu {
542                error: format!(
543                    "timestamp {} overflow with precision {}",
544                    timestamp, precision
545                ),
546            })?
547        }
548    };
549
550    let (datatype, ts) = match ts_type {
551        TimestampType::Millis => (
552            ColumnDataType::TimestampMillisecond,
553            ValueData::TimestampMillisecondValue(ts),
554        ),
555        TimestampType::Nanos => (
556            ColumnDataType::TimestampNanosecond,
557            ValueData::TimestampNanosecondValue(ts),
558        ),
559    };
560
561    let index = column_indexes.get(&name);
562    if let Some(index) = index {
563        check_schema(datatype, SemanticType::Timestamp, &schema[*index])?;
564        one_row[*index].value_data = Some(ts);
565    } else {
566        let index = schema.len();
567        schema.push(time_index_column_schema(&name, datatype));
568        column_indexes.insert(name, index);
569        one_row.push(ts.into())
570    }
571
572    Ok(())
573}
574
575fn check_schema(
576    datatype: ColumnDataType,
577    semantic_type: SemanticType,
578    schema: &ColumnSchema,
579) -> Result<()> {
580    check_schema_number(datatype as i32, semantic_type as i32, schema)
581}
582
583fn check_schema_number(datatype: i32, semantic_type: i32, schema: &ColumnSchema) -> Result<()> {
584    ensure!(
585        schema.datatype == datatype,
586        IncompatibleSchemaSnafu {
587            column_name: &schema.column_name,
588            datatype: "datatype",
589            expected: schema.datatype,
590            actual: datatype,
591        }
592    );
593
594    ensure!(
595        schema.semantic_type == semantic_type,
596        IncompatibleSchemaSnafu {
597            column_name: &schema.column_name,
598            datatype: "semantic_type",
599            expected: schema.semantic_type,
600            actual: semantic_type,
601        }
602    );
603
604    Ok(())
605}
servers/row_writer.rs

servers/
row_writer.rs
