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

servers/
row_writer.rs
