feat: frontend instance (#238)

* feat: frontend instance * no need to carry column length in `Column` proto * add more tests * rebase develop * create a new variant with already provisioned RecordBatches in Output * resolve code review comments * new frontend instance does not connect datanode grpc * add more tests * add more tests * rebase develop Co-authored-by: luofucong <luofucong@greptime.com>
2026-05-20 15:00:40 +00:00 · 2022-09-13 17:10:22 +08:00
parent bdd5bdd917
commit ec99eb0cd0
71 changed files with 2324 additions and 362 deletions
--- a/src/client/Cargo.toml
+++ b/src/client/Cargo.toml
@@ -7,13 +7,21 @@ edition = "2021"

 [dependencies]
 api = { path = "../api" }
+async-stream = "0.3"
+catalog = { path = "../catalog" }
+common-base = { path = "../common/base" }
 common-error = { path = "../common/error" }
 common-grpc = { path = "../common/grpc" }
+common-recordbatch = { path = "../common/recordbatch" }
+common-time = { path = "../common/time" }
 datafusion = { git = "https://github.com/apache/arrow-datafusion.git", branch = "arrow2", features = ["simd"] }
+datatypes = { path = "../datatypes" }
+query = { path = "../query" }
 snafu = { version = "0.7", features = ["backtraces"] }
 tonic = "0.8"

 [dev-dependencies]
+datanode = { path = "../datanode" }
 tokio = { version = "1.0", features = ["full"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
--- a/src/client/examples/insert.rs
+++ b/src/client/examples/insert.rs
@@ -13,7 +13,13 @@ async fn run() {
    let client = Client::connect("http://127.0.0.1:3001").await.unwrap();
    let db = Database::new("greptime", client);

-    db.insert("demo", insert_batches()).await.unwrap();
+    let expr = InsertExpr {
+        table_name: "demo".to_string(),
+        expr: Some(insert_expr::Expr::Values(insert_expr::Values {
+            values: insert_batches(),
+        })),
+    };
+    db.insert(expr).await.unwrap();
 }

 fn insert_batches() -> Vec<Vec<u8>> {
@@ -37,6 +43,7 @@ fn insert_batches() -> Vec<Vec<u8>> {
        semantic_type: SEMANTIC_TAG,
        values: Some(host_vals),
        null_mask: vec![0],
+        ..Default::default()
    };

    let cpu_vals = column::Values {
@@ -48,6 +55,7 @@ fn insert_batches() -> Vec<Vec<u8>> {
        semantic_type: SEMANTIC_FEILD,
        values: Some(cpu_vals),
        null_mask: vec![2],
+        ..Default::default()
    };

    let mem_vals = column::Values {
@@ -59,6 +67,7 @@ fn insert_batches() -> Vec<Vec<u8>> {
        semantic_type: SEMANTIC_FEILD,
        values: Some(mem_vals),
        null_mask: vec![4],
+        ..Default::default()
    };

    let ts_vals = column::Values {
@@ -70,6 +79,7 @@ fn insert_batches() -> Vec<Vec<u8>> {
        semantic_type: SEMANTIC_TS,
        values: Some(ts_vals),
        null_mask: vec![0],
+        ..Default::default()
    };

    let insert_batch = InsertBatch {
--- a/src/client/src/admin.rs
+++ b/src/client/src/admin.rs
@@ -1,4 +1,6 @@
 use api::v1::*;
+use common_error::prelude::StatusCode;
+use query::Output;
 use snafu::prelude::*;

 use crate::database::PROTOCOL_VERSION;
@@ -20,6 +22,10 @@ impl Admin {
        }
    }

+    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
+        self.client.start(url).await
+    }
+
    pub async fn create(&self, expr: CreateExpr) -> Result<AdminResult> {
        let header = ExprHeader {
            version: PROTOCOL_VERSION,
@@ -28,8 +34,12 @@ impl Admin {
            header: Some(header),
            expr: Some(admin_expr::Expr::Create(expr)),
        };
-        // `remove(0)` is safe because of `do_request`'s invariants.
-        Ok(self.do_request(vec![expr]).await?.remove(0))
+        self.do_request(expr).await
+    }
+
+    pub async fn do_request(&self, expr: AdminExpr) -> Result<AdminResult> {
+        // `remove(0)` is safe because of `do_requests`'s invariants.
+        Ok(self.do_requests(vec![expr]).await?.remove(0))
    }

    pub async fn alter(&self, expr: AlterExpr) -> Result<AdminResult> {
@@ -40,11 +50,11 @@ impl Admin {
            header: Some(header),
            expr: Some(admin_expr::Expr::Alter(expr)),
        };
-        Ok(self.do_request(vec![expr]).await?.remove(0))
+        Ok(self.do_requests(vec![expr]).await?.remove(0))
    }

    /// Invariants: the lengths of input vec (`Vec<AdminExpr>`) and output vec (`Vec<AdminResult>`) are equal.
-    async fn do_request(&self, exprs: Vec<AdminExpr>) -> Result<Vec<AdminResult>> {
+    async fn do_requests(&self, exprs: Vec<AdminExpr>) -> Result<Vec<AdminResult>> {
        let expr_count = exprs.len();
        let req = AdminRequest {
            name: self.name.clone(),
@@ -65,3 +75,32 @@ impl Admin {
        Ok(results)
    }
 }
+
+pub fn admin_result_to_output(admin_result: AdminResult) -> Result<Output> {
+    let header = admin_result.header.context(error::MissingHeaderSnafu)?;
+    if !StatusCode::is_success(header.code) {
+        return error::DatanodeSnafu {
+            code: header.code,
+            msg: header.err_msg,
+        }
+        .fail();
+    }
+
+    let result = admin_result.result.context(error::MissingResultSnafu {
+        name: "result".to_string(),
+        expected: 1_usize,
+        actual: 0_usize,
+    })?;
+    let output = match result {
+        admin_result::Result::Mutate(mutate) => {
+            if mutate.failure != 0 {
+                return error::MutateFailureSnafu {
+                    failure: mutate.failure,
+                }
+                .fail();
+            }
+            Output::AffectedRows(mutate.success as usize)
+        }
+    };
+    Ok(output)
+}
--- a/src/client/src/client.rs
+++ b/src/client/src/client.rs
@@ -5,18 +5,43 @@ use tonic::transport::Channel;
 use crate::error;
 use crate::Result;

-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct Client {
-    client: GreptimeClient<Channel>,
+    client: Option<GreptimeClient<Channel>>,
 }

 impl Client {
+    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
+        match self.client.as_ref() {
+            None => {
+                let url = url.into();
+                let client = GreptimeClient::connect(url.clone())
+                    .await
+                    .context(error::ConnectFailedSnafu { url })?;
+                self.client = Some(client);
+                Ok(())
+            }
+            Some(_) => error::IllegalGrpcClientStateSnafu {
+                err_msg: "already started",
+            }
+            .fail(),
+        }
+    }
+
+    pub fn with_client(client: GreptimeClient<Channel>) -> Self {
+        Self {
+            client: Some(client),
+        }
+    }
+
    pub async fn connect(url: impl Into<String>) -> Result<Self> {
        let url = url.into();
        let client = GreptimeClient::connect(url.clone())
            .await
            .context(error::ConnectFailedSnafu { url })?;
-        Ok(Self { client })
+        Ok(Self {
+            client: Some(client),
+        })
    }

    pub async fn admin(&self, req: AdminRequest) -> Result<AdminResponse> {
@@ -48,12 +73,18 @@ impl Client {
    }

    pub async fn batch(&self, req: BatchRequest) -> Result<BatchResponse> {
-        let res = self
-            .client
-            .clone()
-            .batch(req)
-            .await
-            .context(error::TonicStatusSnafu)?;
-        Ok(res.into_inner())
+        if let Some(client) = self.client.as_ref() {
+            let res = client
+                .clone()
+                .batch(req)
+                .await
+                .context(error::TonicStatusSnafu)?;
+            Ok(res.into_inner())
+        } else {
+            error::IllegalGrpcClientStateSnafu {
+                err_msg: "not started",
+            }
+            .fail()
+        }
    }
 }
--- a/src/client/src/database.rs
+++ b/src/client/src/database.rs
@@ -1,15 +1,24 @@
 use std::sync::Arc;

+use api::helper::ColumnDataTypeWrapper;
 use api::v1::codec::SelectResult as GrpcSelectResult;
 use api::v1::{
-    object_expr, object_result, select_expr, DatabaseRequest, ExprHeader, InsertExpr,
-    MutateResult as GrpcMutateResult, ObjectExpr, ObjectResult as GrpcObjectResult, PhysicalPlan,
-    SelectExpr,
+    column::Values, object_expr, object_result, select_expr, Column, ColumnDataType,
+    DatabaseRequest, ExprHeader, InsertExpr, MutateResult as GrpcMutateResult, ObjectExpr,
+    ObjectResult as GrpcObjectResult, PhysicalPlan, SelectExpr,
 };
+use common_base::BitVec;
 use common_error::status_code::StatusCode;
 use common_grpc::AsExcutionPlan;
 use common_grpc::DefaultAsPlanImpl;
+use common_recordbatch::{RecordBatch, RecordBatches};
+use common_time::date::Date;
+use common_time::datetime::DateTime;
+use common_time::timestamp::Timestamp;
 use datafusion::physical_plan::ExecutionPlan;
+use datatypes::prelude::*;
+use datatypes::schema::{ColumnSchema, Schema};
+use query::Output;
 use snafu::{ensure, OptionExt, ResultExt};

 use crate::error;
@@ -19,8 +28,6 @@ use crate::{

 pub const PROTOCOL_VERSION: u32 = 1;

-pub type Bytes = Vec<u8>;
-
 #[derive(Clone, Debug)]
 pub struct Database {
    name: String,
@@ -35,26 +42,23 @@ impl Database {
        }
    }

+    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
+        self.client.start(url).await
+    }
+
    pub fn name(&self) -> &str {
        &self.name
    }

-    pub async fn insert(&self, table: impl Into<String>, values: Vec<Bytes>) -> Result<()> {
+    pub async fn insert(&self, insert: InsertExpr) -> Result<ObjectResult> {
        let header = ExprHeader {
            version: PROTOCOL_VERSION,
        };
-        let insert = InsertExpr {
-            table_name: table.into(),
-            values,
-        };
        let expr = ObjectExpr {
            header: Some(header),
            expr: Some(object_expr::Expr::Insert(insert)),
        };
-
-        self.object(expr).await?;
-
-        Ok(())
+        self.object(expr).await?.try_into()
    }

    pub async fn select(&self, expr: Select) -> Result<ObjectResult> {
@@ -100,7 +104,7 @@ impl Database {

    // TODO(jiachun) update/delete

-    async fn object(&self, expr: ObjectExpr) -> Result<GrpcObjectResult> {
+    pub async fn object(&self, expr: ObjectExpr) -> Result<GrpcObjectResult> {
        let res = self.objects(vec![expr]).await?.pop().unwrap();
        Ok(res)
    }
@@ -165,3 +169,234 @@ impl TryFrom<api::v1::ObjectResult> for ObjectResult {
 pub enum Select {
    Sql(String),
 }
+
+impl TryFrom<ObjectResult> for Output {
+    type Error = error::Error;
+
+    fn try_from(value: ObjectResult) -> Result<Self> {
+        let output = match value {
+            ObjectResult::Select(select) => {
+                let vectors = select
+                    .columns
+                    .iter()
+                    .map(|column| column_to_vector(column, select.row_count))
+                    .collect::<Result<Vec<VectorRef>>>()?;
+
+                let column_schemas = select
+                    .columns
+                    .iter()
+                    .zip(vectors.iter())
+                    .map(|(column, vector)| {
+                        let datatype = vector.data_type();
+                        // nullable or not, does not affect the output
+                        ColumnSchema::new(&column.column_name, datatype, true)
+                    })
+                    .collect::<Vec<ColumnSchema>>();
+
+                let schema = Arc::new(Schema::new(column_schemas));
+                let recordbatches = RecordBatch::new(schema, vectors)
+                    .and_then(|batch| RecordBatches::try_new(batch.schema.clone(), vec![batch]))
+                    .context(error::CreateRecordBatchesSnafu)?;
+                Output::RecordBatches(recordbatches)
+            }
+            ObjectResult::Mutate(mutate) => {
+                if mutate.failure != 0 {
+                    return error::MutateFailureSnafu {
+                        failure: mutate.failure,
+                    }
+                    .fail();
+                }
+                Output::AffectedRows(mutate.success as usize)
+            }
+        };
+        Ok(output)
+    }
+}
+
+fn column_to_vector(column: &Column, rows: u32) -> Result<VectorRef> {
+    let wrapper =
+        ColumnDataTypeWrapper::try_new(column.datatype).context(error::ColumnDataTypeSnafu)?;
+    let column_datatype = wrapper.datatype();
+
+    let rows = rows as usize;
+    let mut vector = VectorBuilder::with_capacity(wrapper.into(), rows);
+
+    if let Some(values) = &column.values {
+        let values = collect_column_values(column_datatype, values);
+        let mut values_iter = values.into_iter();
+
+        let null_mask = BitVec::from_slice(&column.null_mask);
+        let mut nulls_iter = null_mask.iter().by_vals().fuse();
+
+        for i in 0..rows {
+            if let Some(true) = nulls_iter.next() {
+                vector.push_null();
+            } else {
+                let value_ref = values_iter.next().context(error::InvalidColumnProtoSnafu {
+                    err_msg: format!(
+                        "value not found at position {} of column {}",
+                        i, &column.column_name
+                    ),
+                })?;
+                vector
+                    .try_push_ref(value_ref)
+                    .context(error::CreateVectorSnafu)?;
+            }
+        }
+    } else {
+        (0..rows).for_each(|_| vector.push_null());
+    }
+    Ok(vector.finish())
+}
+
+fn collect_column_values(column_datatype: ColumnDataType, values: &Values) -> Vec<ValueRef> {
+    macro_rules! collect_values {
+        ($value: expr, $mapper: expr) => {
+            $value.iter().map($mapper).collect::<Vec<ValueRef>>()
+        };
+    }
+
+    match column_datatype {
+        ColumnDataType::Boolean => collect_values!(values.bool_values, |v| ValueRef::from(*v)),
+        ColumnDataType::Int8 => collect_values!(values.i8_values, |v| ValueRef::from(*v as i8)),
+        ColumnDataType::Int16 => {
+            collect_values!(values.i16_values, |v| ValueRef::from(*v as i16))
+        }
+        ColumnDataType::Int32 => {
+            collect_values!(values.i32_values, |v| ValueRef::from(*v))
+        }
+        ColumnDataType::Int64 => {
+            collect_values!(values.i64_values, |v| ValueRef::from(*v as i64))
+        }
+        ColumnDataType::Uint8 => {
+            collect_values!(values.u8_values, |v| ValueRef::from(*v as u8))
+        }
+        ColumnDataType::Uint16 => {
+            collect_values!(values.u16_values, |v| ValueRef::from(*v as u16))
+        }
+        ColumnDataType::Uint32 => {
+            collect_values!(values.u32_values, |v| ValueRef::from(*v))
+        }
+        ColumnDataType::Uint64 => {
+            collect_values!(values.u64_values, |v| ValueRef::from(*v as u64))
+        }
+        ColumnDataType::Float32 => collect_values!(values.f32_values, |v| ValueRef::from(*v)),
+        ColumnDataType::Float64 => collect_values!(values.f64_values, |v| ValueRef::from(*v)),
+        ColumnDataType::Binary => {
+            collect_values!(values.binary_values, |v| ValueRef::from(v.as_slice()))
+        }
+        ColumnDataType::String => {
+            collect_values!(values.string_values, |v| ValueRef::from(v.as_str()))
+        }
+        ColumnDataType::Date => {
+            collect_values!(values.date_values, |v| ValueRef::Date(Date::new(*v)))
+        }
+        ColumnDataType::Datetime => {
+            collect_values!(values.datetime_values, |v| ValueRef::DateTime(
+                DateTime::new(*v)
+            ))
+        }
+        ColumnDataType::Timestamp => {
+            collect_values!(values.ts_millis_values, |v| ValueRef::Timestamp(
+                Timestamp::from_millis(*v)
+            ))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use datanode::server::grpc::select::{null_mask, values};
+    use datatypes::vectors::{
+        BinaryVector, BooleanVector, DateTimeVector, DateVector, Float32Vector, Float64Vector,
+        Int16Vector, Int32Vector, Int64Vector, Int8Vector, StringVector, UInt16Vector,
+        UInt32Vector, UInt64Vector, UInt8Vector,
+    };
+
+    use super::*;
+
+    #[test]
+    fn test_column_to_vector() {
+        let mut column = create_test_column(Arc::new(BooleanVector::from(vec![true])));
+        column.datatype = -100;
+        let result = column_to_vector(&column, 1);
+        assert!(result.is_err());
+        assert_eq!(
+            result.unwrap_err().to_string(),
+            "Column datatype error, source: Unknown proto column datatype: -100"
+        );
+
+        macro_rules! test_with_vector {
+            ($vector: expr) => {
+                let vector = Arc::new($vector);
+                let column = create_test_column(vector.clone());
+                let result = column_to_vector(&column, vector.len() as u32).unwrap();
+                assert_eq!(result, vector as VectorRef);
+            };
+        }
+
+        test_with_vector!(BooleanVector::from(vec![Some(true), None, Some(false)]));
+        test_with_vector!(Int8Vector::from(vec![Some(i8::MIN), None, Some(i8::MAX)]));
+        test_with_vector!(Int16Vector::from(vec![
+            Some(i16::MIN),
+            None,
+            Some(i16::MAX)
+        ]));
+        test_with_vector!(Int32Vector::from(vec![
+            Some(i32::MIN),
+            None,
+            Some(i32::MAX)
+        ]));
+        test_with_vector!(Int64Vector::from(vec![
+            Some(i64::MIN),
+            None,
+            Some(i64::MAX)
+        ]));
+        test_with_vector!(UInt8Vector::from(vec![Some(u8::MIN), None, Some(u8::MAX)]));
+        test_with_vector!(UInt16Vector::from(vec![
+            Some(u16::MIN),
+            None,
+            Some(u16::MAX)
+        ]));
+        test_with_vector!(UInt32Vector::from(vec![
+            Some(u32::MIN),
+            None,
+            Some(u32::MAX)
+        ]));
+        test_with_vector!(UInt64Vector::from(vec![
+            Some(u64::MIN),
+            None,
+            Some(u64::MAX)
+        ]));
+        test_with_vector!(Float32Vector::from(vec![
+            Some(f32::MIN),
+            None,
+            Some(f32::MAX)
+        ]));
+        test_with_vector!(Float64Vector::from(vec![
+            Some(f64::MIN),
+            None,
+            Some(f64::MAX)
+        ]));
+        test_with_vector!(BinaryVector::from(vec![
+            Some(b"".to_vec()),
+            None,
+            Some(b"hello".to_vec())
+        ]));
+        test_with_vector!(StringVector::from(vec![Some(""), None, Some("foo"),]));
+        test_with_vector!(DateVector::from(vec![Some(1), None, Some(3)]));
+        test_with_vector!(DateTimeVector::from(vec![Some(4), None, Some(6)]));
+    }
+
+    fn create_test_column(vector: VectorRef) -> Column {
+        let wrapper: ColumnDataTypeWrapper = vector.data_type().try_into().unwrap();
+        let array = vector.to_arrow_array();
+        Column {
+            column_name: "test".to_string(),
+            semantic_type: 1,
+            values: Some(values(&[array.clone()]).unwrap()),
+            null_mask: null_mask(&vec![array], vector.len()),
+            datatype: wrapper.datatype() as i32,
+        }
+    }
+}
--- a/src/client/src/error.rs
+++ b/src/client/src/error.rs
@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::sync::Arc;

 use api::serde::DecodeError;
@@ -42,6 +43,67 @@ pub enum Error {
        #[snafu(backtrace)]
        source: common_grpc::Error,
    },
+
+    #[snafu(display("Mutate result has failure {}", failure))]
+    MutateFailure { failure: u32, backtrace: Backtrace },
+
+    #[snafu(display("Invalid column proto: {}", err_msg))]
+    InvalidColumnProto {
+        err_msg: String,
+        backtrace: Backtrace,
+    },
+
+    #[snafu(display("Column datatype error, source: {}", source))]
+    ColumnDataType {
+        #[snafu(backtrace)]
+        source: api::error::Error,
+    },
+
+    #[snafu(display("Failed to create vector, source: {}", source))]
+    CreateVector {
+        #[snafu(backtrace)]
+        source: datatypes::error::Error,
+    },
+
+    #[snafu(display("Failed to create RecordBatches, source: {}", source))]
+    CreateRecordBatches {
+        #[snafu(backtrace)]
+        source: common_recordbatch::error::Error,
+    },
+
+    #[snafu(display("Illegal GRPC client state: {}", err_msg))]
+    IllegalGrpcClientState {
+        err_msg: String,
+        backtrace: Backtrace,
+    },
 }

 pub type Result<T> = std::result::Result<T, Error>;
+
+impl ErrorExt for Error {
+    fn status_code(&self) -> StatusCode {
+        match self {
+            Error::ConnectFailed { .. }
+            | Error::MissingResult { .. }
+            | Error::MissingHeader { .. }
+            | Error::TonicStatus { .. }
+            | Error::DecodeSelect { .. }
+            | Error::Datanode { .. }
+            | Error::EncodePhysical { .. }
+            | Error::MutateFailure { .. }
+            | Error::InvalidColumnProto { .. }
+            | Error::ColumnDataType { .. } => StatusCode::Internal,
+            Error::CreateVector { source } => source.status_code(),
+            Error::CreateRecordBatches { source } => source.status_code(),
+            Error::IllegalGrpcClientState { .. } => StatusCode::Unexpected,
+        }
+    }
+
+    fn backtrace_opt(&self) -> Option<&Backtrace> {
+        ErrorCompat::backtrace(self)
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}