refactor: find new columns to improve write performance (#918)

## Related Issues
Closes #

## Detailed Changes
Refactor find new columns to improve write performance.
* Skip building column schema when columns already exists.

## Test Plan 
TSBS.

## Benchmark
| | auto_create_table = false | auto_create_table = true| performance
loss |

|:--------:|:-------------------------:|:------------------------:|:-----------------:|
| Before | 171273.03 rows/s | 152436.35 rows/s | 11% |
| After | 171699.82 rows/s | 164791.73 rows/s | 4.1% |

Author: chunshao90

integration_tests/sdk/go/issue-779.go

@@ -7,7 +7,7 @@ import (
 )
 
 func checkAutoAddColumnsWithCreateTable(ctx context.Context, client ceresdb.Client) error {
-	timestampName := "t"
+	timestampName := "timestamp"
 
 	err := dropTable(ctx, client)
 	if err != nil {

integration_tests/sdk/go/util.go

@@ -10,7 +10,7 @@ import (
 const table = "godemo"
 
 func createTable(ctx context.Context, client ceresdb.Client, timestampName string) error {
-	_, err := ddl(ctx, client, fmt.Sprintf("create table %s (%s timestamp not null, name string tag, value int64,TIMESTAMP KEY(t))", table, timestampName))
+	_, err := ddl(ctx, client, fmt.Sprintf("create table %s (`%s` timestamp not null, name string tag, value int64,TIMESTAMP KEY(%s))", table, timestampName, timestampName))
 	return err
 }
 
@@ -60,9 +60,13 @@ func ensureRow(expectedVals []ceresdb.Value, actualRow []ceresdb.Column) error {
 }
 
 func query(ctx context.Context, client ceresdb.Client, ts int64, timestampName string, addNewColumn bool) error {
+	sql := fmt.Sprintf("select timestamp, name, value from %s where %s = %d order by name", table, timestampName, ts)
+	if addNewColumn {
+		sql = fmt.Sprintf("select timestamp, name, value, new_tag, new_field from %s where %s = %d order by name", table, timestampName, ts)
+	}
 	resp, err := client.SQLQuery(ctx, ceresdb.SQLQueryRequest{
 		Tables: []string{table},
-		SQL:    fmt.Sprintf("select * from %s where %s = %d order by name", table, timestampName, ts),
+		SQL:    sql,
 	})
 	if err != nil {
 		return err
@@ -73,23 +77,19 @@ func query(ctx context.Context, client ceresdb.Client, ts int64, timestampName s
 	}
 
 	row0 := []ceresdb.Value{
-		ceresdb.NewUint64Value(4024844655630594205),
 		ceresdb.NewInt64Value(ts),
 		ceresdb.NewStringValue("tag-0"),
 		ceresdb.NewInt64Value(0)}
 
 	row1 := []ceresdb.Value{
-		ceresdb.NewUint64Value(14230010170561829440),
 		ceresdb.NewInt64Value(ts),
 		ceresdb.NewStringValue("tag-1"),
 		ceresdb.NewInt64Value(1),
 	}
 
 	if addNewColumn {
-		row0[0] = ceresdb.NewUint64Value(8341999341185504339)
-		row1[0] = ceresdb.NewUint64Value(4452331151453582498)
-		row0 = append(row0, ceresdb.NewInt64Value(0), ceresdb.NewStringValue("new-tag-0"))
-		row1 = append(row1, ceresdb.NewInt64Value(1), ceresdb.NewStringValue("new-tag-1"))
+		row0 = append(row0, ceresdb.NewStringValue("new-tag-0"), ceresdb.NewInt64Value(0))
+		row1 = append(row1, ceresdb.NewStringValue("new-tag-1"), ceresdb.NewInt64Value(1))
 	}
 
 	if err := ensureRow(row0,

proxy/src/write.rs

@@ -10,7 +10,7 @@ use std::{
 
 use bytes::Bytes;
 use ceresdbproto::storage::{
-    storage_service_client::StorageServiceClient, value, RouteRequest, WriteRequest,
+    storage_service_client::StorageServiceClient, value, RouteRequest, Value, WriteRequest,
     WriteResponse as WriteResponsePB, WriteSeriesEntry, WriteTableRequest,
 };
 use cluster::config::SchemaConfig;
@@ -34,7 +34,7 @@ use query_engine::executor::Executor as QueryExecutor;
 use query_frontend::{
     frontend::{Context as FrontendContext, Frontend},
     plan::{AlterTableOperation, AlterTablePlan, InsertPlan, Plan},
-    planner::build_schema_from_write_table_request,
+    planner::{build_column_schema, try_get_data_type_from_value},
     provider::CatalogMetaProvider,
 };
 use router::endpoint::Endpoint;
@@ -43,7 +43,7 @@ use table_engine::table::TableRef;
 use tonic::transport::Channel;
 
 use crate::{
-    error::{ErrNoCause, ErrWithCause, InternalNoCause, Result},
+    error::{ErrNoCause, ErrWithCause, Internal, InternalNoCause, Result},
     forward::{ForwardResult, ForwarderRef},
     Context, Proxy,
 };
@@ -477,14 +477,6 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
             code: StatusCode::BAD_REQUEST,
         })?;
         let schema = req_ctx.database;
-        let schema_config = self
-            .schema_config_provider
-            .schema_config(&schema)
-            .box_err()
-            .with_context(|| ErrWithCause {
-                code: StatusCode::INTERNAL_SERVER_ERROR,
-                msg: format!("Fail to fetch schema config, schema_name:{schema}"),
-            })?;
 
         debug!(
             "Local write begin, catalog:{catalog}, schema:{schema}, request_id:{request_id}, first_table:{:?}, num_tables:{}",
@@ -503,7 +495,7 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
         };
 
         let plan_vec = self
-            .write_request_to_insert_plan(req.table_requests, schema_config, write_context)
+            .write_request_to_insert_plan(req.table_requests, write_context)
             .await?;
 
         let mut success = 0;
@@ -522,7 +514,6 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
     async fn write_request_to_insert_plan(
         &self,
         table_requests: Vec<WriteTableRequest>,
-        schema_config: Option<&SchemaConfig>,
         write_context: WriteContext,
     ) -> Result<Vec<InsertPlan>> {
         let mut plan_vec = Vec::with_capacity(table_requests.len());
@@ -534,7 +525,6 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
             deadline,
             auto_create_table,
         } = write_context;
-        let schema_config = schema_config.cloned().unwrap_or_default();
         for write_table_req in table_requests {
             let table_name = &write_table_req.table;
             self.maybe_open_partition_table_if_not_exist(&catalog, &schema, table_name)
@@ -555,7 +545,7 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
                 // * Currently, the decision to add columns is made at the request level, not at
                 //   the row level, so the cost is relatively small.
                 let table_schema = table.schema();
-                let columns = find_new_columns(&table_schema, &schema_config, &write_table_req)?;
+                let columns = find_new_columns(&table_schema, &write_table_req)?;
                 if !columns.is_empty() {
                     self.execute_add_columns_plan(
                         request_id,
@@ -668,32 +658,93 @@ impl<Q: QueryExecutor + 'static> Proxy<Q> {
 
 fn find_new_columns(
     schema: &Schema,
-    schema_config: &SchemaConfig,
-    write_req: &WriteTableRequest,
+    write_table_req: &WriteTableRequest,
 ) -> Result<Vec<ColumnSchema>> {
-    let new_schema = build_schema_from_write_table_request(schema_config, write_req)
+    let WriteTableRequest {
+        table,
+        field_names,
+        tag_names,
+        entries: write_entries,
+    } = write_table_req;
+
+    let mut columns: HashMap<_, ColumnSchema> = HashMap::new();
+    for write_entry in write_entries {
+        // Parse tags.
+        for tag in &write_entry.tags {
+            let name_index = tag.name_index as usize;
+            ensure!(
+                name_index < tag_names.len(),
+                InternalNoCause {
+                    msg: format!(
+                        "Tag {tag:?} is not found in tag_names:{tag_names:?}, table:{table}",
+                    ),
+                }
+            );
+
+            let tag_name = &tag_names[name_index];
+
+            build_column(&mut columns, schema, tag_name, &tag.value, true)?;
+        }
+
+        // Parse fields.
+        for field_group in &write_entry.field_groups {
+            for field in &field_group.fields {
+                let field_index = field.name_index as usize;
+                ensure!(
+                    field_index < field_names.len(),
+                    InternalNoCause {
+                        msg: format!(
+                        "Field {field:?} is not found in field_names:{field_names:?}, table:{table}",
+                    ),
+                    }
+                );
+                let field_name = &field_names[field.name_index as usize];
+                build_column(&mut columns, schema, field_name, &field.value, false)?;
+            }
+        }
+    }
+
+    Ok(columns.into_iter().map(|v| v.1).collect())
+}
+
+fn build_column<'a>(
+    columns: &mut HashMap<&'a str, ColumnSchema>,
+    schema: &Schema,
+    name: &'a str,
+    value: &Option<Value>,
+    is_tag: bool,
+) -> Result<()> {
+    // Skip adding columns, the following cases:
+    // 1. Column already exists.
+    // 2. The new column has been added.
+    if schema.index_of(name).is_some() || columns.get(name).is_some() {
+        return Ok(());
+    }
+
+    let column_value = value
+        .as_ref()
+        .with_context(|| InternalNoCause {
+            msg: format!("Column value is needed, column:{name}"),
+        })?
+        .value
+        .as_ref()
+        .with_context(|| InternalNoCause {
+            msg: format!("Column value type is not supported, column:{name}"),
+        })?;
+
+    let data_type = try_get_data_type_from_value(column_value)
         .box_err()
-        .context(ErrWithCause {
-            code: StatusCode::INTERNAL_SERVER_ERROR,
-            msg: "Build schema from write table request failed",
+        .context(Internal {
+            msg: "Failed to get data type",
         })?;
 
-    let columns = new_schema.columns();
-    let old_columns = schema.columns();
-
-    // find new columns:
-    // 1. timestamp column can't be a new column;
-    // 2. column not in old schema is a new column.
-    let new_columns = columns
-        .iter()
-        .enumerate()
-        .filter(|(idx, column)| {
-            *idx != new_schema.timestamp_index()
-                && !old_columns.iter().any(|c| c.name == column.name)
-        })
-        .map(|(_, column)| column.clone())
-        .collect();
-    Ok(new_columns)
+    let column_schema = build_column_schema(name, data_type, is_tag)
+        .box_err()
+        .context(Internal {
+            msg: "Failed to build column schema",
+        })?;
+    columns.insert(name, column_schema);
+    Ok(())
 }
 
 fn write_table_request_to_insert_plan(
@@ -898,7 +949,7 @@ fn convert_proto_value_to_datum(
 mod test {
     use ceresdbproto::storage::{value, Field, FieldGroup, Tag, Value, WriteSeriesEntry};
     use common_types::{
-        column_schema::{self, ColumnSchema},
+        column_schema::{self},
         datum::{Datum, DatumKind},
         row::Row,
         schema::Builder,
@@ -908,45 +959,136 @@ mod test {
 
     use super::*;
 
-    const TAG_K: &str = "tagk";
-    const TAG_V: &str = "tagv";
-    const TAG_K1: &str = "tagk1";
-    const TAG_V1: &str = "tagv1";
-    const FIELD_NAME: &str = "field";
-    const FIELD_NAME1: &str = "field1";
-    const FIELD_VALUE_STRING: &str = "stringValue";
+    const NAME_COL1: &str = "col1";
+    const NAME_NEW_COL1: &str = "new_col1";
+    const NAME_COL2: &str = "col2";
+    const NAME_COL3: &str = "col3";
+    const NAME_COL4: &str = "col4";
+    const NAME_COL5: &str = "col5";
 
-    // tag_names field_names write_entry
-    fn generate_write_entry() -> (Schema, Vec<String>, Vec<String>, WriteSeriesEntry) {
-        let tag_names = vec![TAG_K.to_string(), TAG_K1.to_string()];
-        let field_names = vec![FIELD_NAME.to_string(), FIELD_NAME1.to_string()];
+    #[test]
+    fn test_write_entry_to_row_group() {
+        let (schema, tag_names, field_names, write_entry) = generate_write_entry();
+        let rows =
+            write_entry_to_rows("test_table", &schema, &tag_names, &field_names, write_entry)
+                .unwrap();
+        let row0 = vec![
+            Datum::Timestamp(Timestamp::new(1000)),
+            Datum::String(NAME_COL1.into()),
+            Datum::String(NAME_COL2.into()),
+            Datum::Int64(100),
+            Datum::Null,
+        ];
+        let row1 = vec![
+            Datum::Timestamp(Timestamp::new(2000)),
+            Datum::String(NAME_COL1.into()),
+            Datum::String(NAME_COL2.into()),
+            Datum::Null,
+            Datum::Int64(10),
+        ];
+        let row2 = vec![
+            Datum::Timestamp(Timestamp::new(3000)),
+            Datum::String(NAME_COL1.into()),
+            Datum::String(NAME_COL2.into()),
+            Datum::Null,
+            Datum::Int64(10),
+        ];
 
-        let tag = Tag {
-            name_index: 0,
-            value: Some(Value {
-                value: Some(value::Value::StringValue(TAG_V.to_string())),
-            }),
-        };
-        let tag1 = Tag {
-            name_index: 1,
+        let expect_rows = vec![
+            Row::from_datums(row0),
+            Row::from_datums(row1),
+            Row::from_datums(row2),
+        ];
+        assert_eq!(rows, expect_rows);
+    }
+
+    #[test]
+    fn test_find_new_columns() {
+        let write_table_request = generate_write_table_request();
+        let schema = build_schema();
+        let new_columns = find_new_columns(&schema, &write_table_request)
+            .unwrap()
+            .into_iter()
+            .map(|v| (v.name.clone(), v))
+            .collect::<HashMap<_, _>>();
+
+        assert_eq!(new_columns.len(), 2);
+        assert!(new_columns.get(NAME_NEW_COL1).is_some());
+        assert!(new_columns.get(NAME_NEW_COL1).unwrap().is_tag);
+        assert!(new_columns.get(NAME_COL5).is_some());
+        assert!(!new_columns.get(NAME_COL5).unwrap().is_tag);
+    }
+
+    fn build_schema() -> Schema {
+        Builder::new()
+            .auto_increment_column_id(true)
+            .add_key_column(
+                column_schema::Builder::new(
+                    TIMESTAMP_COLUMN_NAME.to_string(),
+                    DatumKind::Timestamp,
+                )
+                .build()
+                .unwrap(),
+            )
+            .unwrap()
+            .add_key_column(
+                column_schema::Builder::new(NAME_COL1.to_string(), DatumKind::String)
+                    .is_tag(true)
+                    .build()
+                    .unwrap(),
+            )
+            .unwrap()
+            .add_key_column(
+                column_schema::Builder::new(NAME_COL2.to_string(), DatumKind::String)
+                    .is_tag(true)
+                    .build()
+                    .unwrap(),
+            )
+            .unwrap()
+            .add_normal_column(
+                column_schema::Builder::new(NAME_COL3.to_string(), DatumKind::Int64)
+                    .build()
+                    .unwrap(),
+            )
+            .unwrap()
+            .add_normal_column(
+                column_schema::Builder::new(NAME_COL4.to_string(), DatumKind::Int64)
+                    .build()
+                    .unwrap(),
+            )
+            .unwrap()
+            .build()
+            .unwrap()
+    }
+
+    fn make_tag(name_index: u32, val: &str) -> Tag {
+        Tag {
+            name_index,
             value: Some(Value {
-                value: Some(value::Value::StringValue(TAG_V1.to_string())),
+                value: Some(value::Value::StringValue(val.to_string())),
             }),
-        };
+        }
+    }
+
+    fn make_field(name_index: u32, val: value::Value) -> Field {
+        Field {
+            name_index,
+            value: Some(Value { value: Some(val) }),
+        }
+    }
+
+    // tag_names field_names write_entry
+    fn generate_write_entry() -> (Schema, Vec<String>, Vec<String>, WriteSeriesEntry) {
+        let tag_names = vec![NAME_COL1.to_string(), NAME_COL2.to_string()];
+        let field_names = vec![NAME_COL3.to_string(), NAME_COL4.to_string()];
+
+        let tag = make_tag(0, NAME_COL1);
+        let tag1 = make_tag(1, NAME_COL2);
         let tags = vec![tag, tag1];
 
-        let field = Field {
-            name_index: 0,
-            value: Some(Value {
-                value: Some(value::Value::Float64Value(100.0)),
-            }),
-        };
-        let field1 = Field {
-            name_index: 1,
-            value: Some(Value {
-                value: Some(value::Value::StringValue(FIELD_VALUE_STRING.to_string())),
-            }),
-        };
+        let field = make_field(0, value::Value::Int64Value(100));
+        let field1 = make_field(1, value::Value::Int64Value(10));
+
         let field_group = FieldGroup {
             timestamp: 1000,
             fields: vec![field],
@@ -965,102 +1107,44 @@ mod test {
             field_groups: vec![field_group, field_group1, field_group2],
         };
 
-        let schema_builder = Builder::new();
-        let schema = schema_builder
-            .auto_increment_column_id(true)
-            .add_key_column(ColumnSchema {
-                id: column_schema::COLUMN_ID_UNINIT,
-                name: TIMESTAMP_COLUMN_NAME.to_string(),
-                data_type: DatumKind::Timestamp,
-                is_nullable: false,
-                is_tag: false,
-                comment: String::new(),
-                escaped_name: TIMESTAMP_COLUMN_NAME.escape_debug().to_string(),
-                default_value: None,
-            })
-            .unwrap()
-            .add_key_column(ColumnSchema {
-                id: column_schema::COLUMN_ID_UNINIT,
-                name: TAG_K.to_string(),
-                data_type: DatumKind::String,
-                is_nullable: false,
-                is_tag: true,
-                comment: String::new(),
-                escaped_name: TAG_K.escape_debug().to_string(),
-                default_value: None,
-            })
-            .unwrap()
-            .add_normal_column(ColumnSchema {
-                id: column_schema::COLUMN_ID_UNINIT,
-                name: TAG_K1.to_string(),
-                data_type: DatumKind::String,
-                is_nullable: false,
-                is_tag: true,
-                comment: String::new(),
-                escaped_name: TAG_K1.escape_debug().to_string(),
-                default_value: None,
-            })
-            .unwrap()
-            .add_normal_column(ColumnSchema {
-                id: column_schema::COLUMN_ID_UNINIT,
-                name: FIELD_NAME.to_string(),
-                data_type: DatumKind::Double,
-                is_nullable: true,
-                is_tag: false,
-                comment: String::new(),
-                escaped_name: FIELD_NAME.escape_debug().to_string(),
-                default_value: None,
-            })
-            .unwrap()
-            .add_normal_column(ColumnSchema {
-                id: column_schema::COLUMN_ID_UNINIT,
-                name: FIELD_NAME1.to_string(),
-                data_type: DatumKind::String,
-                is_nullable: true,
-                is_tag: false,
-                comment: String::new(),
-                escaped_name: FIELD_NAME1.escape_debug().to_string(),
-                default_value: None,
-            })
-            .unwrap()
-            .build()
-            .unwrap();
+        let schema = build_schema();
         (schema, tag_names, field_names, write_entry)
     }
 
-    #[test]
-    fn test_write_entry_to_row_group() {
-        let (schema, tag_names, field_names, write_entry) = generate_write_entry();
-        let rows =
-            write_entry_to_rows("test_table", &schema, &tag_names, &field_names, write_entry)
-                .unwrap();
-        let row0 = vec![
-            Datum::Timestamp(Timestamp::new(1000)),
-            Datum::String(TAG_V.into()),
-            Datum::String(TAG_V1.into()),
-            Datum::Double(100.0),
-            Datum::Null,
-        ];
-        let row1 = vec![
-            Datum::Timestamp(Timestamp::new(2000)),
-            Datum::String(TAG_V.into()),
-            Datum::String(TAG_V1.into()),
-            Datum::Null,
-            Datum::String(FIELD_VALUE_STRING.into()),
-        ];
-        let row2 = vec![
-            Datum::Timestamp(Timestamp::new(3000)),
-            Datum::String(TAG_V.into()),
-            Datum::String(TAG_V1.into()),
-            Datum::Null,
-            Datum::String(FIELD_VALUE_STRING.into()),
-        ];
+    fn generate_write_table_request() -> WriteTableRequest {
+        let tag1 = make_tag(0, NAME_NEW_COL1);
+        let tag2 = make_tag(1, NAME_COL1);
+        let tags = vec![tag1, tag2];
 
-        let expect_rows = vec![
-            Row::from_datums(row0),
-            Row::from_datums(row1),
-            Row::from_datums(row2),
-        ];
-        assert_eq!(rows, expect_rows);
+        let field1 = make_field(0, value::Value::Int64Value(100));
+        let field2 = make_field(1, value::Value::Int64Value(10));
+
+        let field_group1 = FieldGroup {
+            timestamp: 1000,
+            fields: vec![field1.clone(), field2.clone()],
+        };
+        let field_group2 = FieldGroup {
+            timestamp: 2000,
+            fields: vec![field1],
+        };
+        let field_group3 = FieldGroup {
+            timestamp: 3000,
+            fields: vec![field2],
+        };
+
+        let write_entry = WriteSeriesEntry {
+            tags,
+            field_groups: vec![field_group1, field_group2, field_group3],
+        };
+
+        let tag_names = vec![NAME_NEW_COL1.to_string(), NAME_COL1.to_string()];
+        let field_names = vec![NAME_COL3.to_string(), NAME_COL5.to_string()];
+
+        WriteTableRequest {
+            table: "test".to_string(),
+            tag_names,
+            field_names,
+            entries: vec![write_entry],
+        }
     }
 }

query_frontend/src/planner.rs

@@ -363,7 +363,7 @@ impl<'a, P: MetaProvider> Planner<'a, P> {
     }
 }
 
-fn build_column_schema(
+pub fn build_column_schema(
     column_name: &str,
     data_type: DatumKind,
     is_tag: bool,
@@ -537,7 +537,7 @@ fn ensure_data_type_compatible(
     Ok(())
 }
 
-fn try_get_data_type_from_value(value: &PbValue) -> Result<DatumKind> {
+pub fn try_get_data_type_from_value(value: &PbValue) -> Result<DatumKind> {
     match value {
         PbValue::Float64Value(_) => Ok(DatumKind::Double),
         PbValue::StringValue(_) => Ok(DatumKind::String),
@@ -555,6 +555,7 @@ fn try_get_data_type_from_value(value: &PbValue) -> Result<DatumKind> {
         PbValue::VarbinaryValue(_) => Ok(DatumKind::Varbinary),
     }
 }
+
 /// A planner wraps the datafusion's logical planner, and delegate sql like
 /// select/explain to datafusion's planner.
 pub(crate) struct PlannerDelegate<'a, P: MetaProvider> {