feat: FlightSQL support update (DDL and insert) (#10765)

* refactor: make has_result_set a method of Plan.

* refactor: Column::into_arrow_rs support Boolean.

* feat: DataBlock::from_record_batch.

* refactor: polish log for do_action_close_prepared_statement.

* refactor: do_action_create_prepared_statement return empty schema for plan without result set.

* feat: support FlightSQL support update.

* feat: covert int arrays from arrow_rs.

* feat: FlightSQL impl do_put_statement_update.

* ci: test FlightSQL update.

* feat: change FlightSQL port to 8900.

* chore: clear println.

Author: youngsofun

scripts/ci/deploy/config/databend-query-node-1.toml

@@ -29,7 +29,7 @@ http_handler_port = 8000
 
 # Databend Query FlightSQL Handler.
 flight_sql_handler_host = "0.0.0.0"
-flight_sql_handler_port = 50050
+flight_sql_handler_port = 8900
 
 tenant_id = "test_tenant"
 cluster_id = "test_cluster"

scripts/ci/deploy/config/databend-query-node-2.toml

@@ -29,7 +29,7 @@ http_handler_port = 8002
 
 # Databend Query FlightSQL Handler.
 flight_sql_handler_host = "0.0.0.0"
-flight_sql_handler_port = 7002
+flight_sql_handler_port = 8902
 
 tenant_id = "test_tenant"
 cluster_id = "test_cluster"

scripts/ci/deploy/config/databend-query-node-3.toml

@@ -30,7 +30,7 @@ http_handler_port = 8003
 
 # Databend Query FlightSQL Handler.
 flight_sql_handler_host = "0.0.0.0"
-flight_sql_handler_port = 7003
+flight_sql_handler_port = 8903
 
 tenant_id = "test_tenant"
 cluster_id = "test_cluster"

scripts/ci/deploy/config/databend-query-node-shared.toml

@@ -27,6 +27,9 @@ clickhouse_http_handler_port = 58124
 http_handler_host = "0.0.0.0"
 http_handler_port = 58000
 
+flight_sql_handler_host = "0.0.0.0"
+flight_sql_handler_port = 58900
+
 tenant_id = "shared_tenant"
 cluster_id = "test_cluster"
 

src/query/config/src/config.rs

@@ -1154,7 +1154,7 @@ pub struct QueryConfig {
     #[clap(long, default_value = "127.0.0.1")]
     pub flight_sql_handler_host: String,
 
-    #[clap(long, default_value = "7000")]
+    #[clap(long, default_value = "8900")]
     pub flight_sql_handler_port: u16,
 
     #[clap(long, default_value = "127.0.0.1:9090")]

src/query/config/src/inner.rs

@@ -196,7 +196,7 @@ impl Default for QueryConfig {
             http_handler_result_timeout_secs: 60,
             flight_api_address: "127.0.0.1:9090".to_string(),
             flight_sql_handler_host: "127.0.0.1".to_string(),
-            flight_sql_handler_port: 7000,
+            flight_sql_handler_port: 8900,
             admin_api_address: "127.0.0.1:8080".to_string(),
             metric_api_address: "127.0.0.1:7070".to_string(),
             api_tls_server_cert: "".to_string(),

src/query/expression/src/convert_arrow_rs/array.rs

@@ -17,9 +17,19 @@ use std::sync::Arc;
 use arrow_array::make_array;
 use arrow_array::Array;
 use arrow_array::ArrowPrimitiveType;
+use arrow_array::BooleanArray;
+use arrow_array::Int16Array;
+use arrow_array::Int32Array;
+use arrow_array::Int64Array;
+use arrow_array::Int8Array;
 use arrow_array::LargeStringArray;
 use arrow_array::NullArray;
 use arrow_array::PrimitiveArray;
+use arrow_array::StringArray;
+use arrow_array::UInt16Array;
+use arrow_array::UInt32Array;
+use arrow_array::UInt64Array;
+use arrow_array::UInt8Array;
 use arrow_buffer::i256;
 use arrow_buffer::ArrowNativeType;
 use arrow_buffer::Buffer;
@@ -28,12 +38,15 @@ use arrow_data::ArrayDataBuilder;
 use arrow_schema::ArrowError;
 use arrow_schema::DataType;
 use arrow_schema::TimeUnit;
+use common_arrow::arrow::bitmap::Bitmap;
 use common_arrow::arrow::buffer::Buffer as Buffer2;
 use common_arrow::arrow::Either;
 use ordered_float::OrderedFloat;
 
 use crate::types::decimal::DecimalColumn;
+use crate::types::nullable::NullableColumn;
 use crate::types::number::NumberColumn;
+use crate::types::string::StringColumn;
 use crate::Column;
 
 fn try_take_buffer<T: Clone>(buffer: Buffer2<T>) -> Vec<T> {
@@ -65,6 +78,14 @@ impl Column {
             Column::Null { len } => Arc::new(NullArray::new(len)),
             Column::EmptyArray { len } => Arc::new(NullArray::new(len)),
             Column::EmptyMap { len } => Arc::new(NullArray::new(len)),
+            Column::Boolean(bitmap) => {
+                let buf = bitmap.as_slice().0;
+                let array_data = ArrayData::builder(DataType::Boolean)
+                    .len(bitmap.len())
+                    .add_buffer(buf.into())
+                    .build()?;
+                Arc::new(BooleanArray::from(array_data))
+            }
             Column::String(col) => {
                 let len = col.len();
                 let values: Vec<u8> = try_take_buffer(col.data);
@@ -160,10 +181,108 @@ impl Column {
             _ => {
                 let data_type = self.data_type();
                 Err(ArrowError::NotYetImplemented(format!(
-                    "Column::into_arrow_rs()  for {data_type} not implemented yet"
+                    "Column::into_arrow_rs() for {data_type} not implemented yet"
                 )))?
             }
         };
         Ok(array)
     }
+
+    pub fn from_arrow_rs(array: Arc<dyn Array>) -> Result<Self, ArrowError> {
+        let data_type = array.data_type();
+        let column = match data_type {
+            DataType::Null => Column::Null { len: array.len() },
+            DataType::LargeUtf8 => {
+                let array = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
+                let offsets = array.value_offsets().to_vec();
+                let offsets = unsafe { std::mem::transmute::<Vec<i64>, Vec<u64>>(offsets) };
+                Column::String(StringColumn {
+                    offsets: offsets.into(),
+                    data: array.value_data().to_vec().into(),
+                })
+            }
+            DataType::Utf8 => {
+                let array = array.as_any().downcast_ref::<StringArray>().unwrap();
+                let offsets = array
+                    .value_offsets()
+                    .iter()
+                    .map(|x| *x as u64)
+                    .collect::<Vec<_>>();
+                Column::String(StringColumn {
+                    offsets: offsets.into(),
+                    data: array.value_data().to_vec().into(),
+                })
+            }
+            DataType::Boolean => {
+                let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
+                let bytes = array.values().clone().into_vec().map_err(|_| {
+                    ArrowError::CastError(
+                        "can not covert Buffer of BooleanArray to Vec".to_string(),
+                    )
+                })?;
+                let bitmap = Bitmap::try_new(bytes, array.len()).map_err(|e| {
+                    ArrowError::CastError(format!(
+                        "can not covert  BooleanArray to Column::Boolean: {e:?}"
+                    ))
+                })?;
+                Column::Boolean(bitmap)
+            }
+            DataType::Int8 => {
+                let array = array.as_any().downcast_ref::<Int8Array>().unwrap();
+                let buffer2: Buffer2<i8> = array.values().to_vec().into();
+                Column::Number(NumberColumn::Int8(buffer2))
+            }
+            DataType::UInt8 => {
+                let array = array.as_any().downcast_ref::<UInt8Array>().unwrap();
+                let buffer2: Buffer2<u8> = array.values().to_vec().into();
+                Column::Number(NumberColumn::UInt8(buffer2))
+            }
+            DataType::Int16 => {
+                let array = array.as_any().downcast_ref::<Int16Array>().unwrap();
+                let buffer2: Buffer2<i16> = array.values().to_vec().into();
+                Column::Number(NumberColumn::Int16(buffer2))
+            }
+            DataType::UInt16 => {
+                let array = array.as_any().downcast_ref::<UInt16Array>().unwrap();
+                let buffer2: Buffer2<u16> = array.values().to_vec().into();
+                Column::Number(NumberColumn::UInt16(buffer2))
+            }
+            DataType::Int32 => {
+                let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
+                let buffer2: Buffer2<i32> = array.values().to_vec().into();
+                Column::Number(NumberColumn::Int32(buffer2))
+            }
+            DataType::UInt32 => {
+                let array = array.as_any().downcast_ref::<UInt32Array>().unwrap();
+                let buffer2: Buffer2<u32> = array.values().to_vec().into();
+                Column::Number(NumberColumn::UInt32(buffer2))
+            }
+            DataType::Int64 => {
+                let array = array.as_any().downcast_ref::<Int64Array>().unwrap();
+                let buffer2: Buffer2<i64> = array.values().to_vec().into();
+                Column::Number(NumberColumn::Int64(buffer2))
+            }
+            DataType::UInt64 => {
+                let array = array.as_any().downcast_ref::<UInt64Array>().unwrap();
+                let buffer2: Buffer2<u64> = array.values().to_vec().into();
+                Column::Number(NumberColumn::UInt64(buffer2))
+            }
+
+            _ => Err(ArrowError::NotYetImplemented(format!(
+                "Column::from_arrow_rs() for {data_type} not implemented yet"
+            )))?,
+        };
+        if let Some(nulls) = array.into_data().nulls() {
+            let validity =
+                Bitmap::try_new(nulls.buffer().to_vec(), nulls.offset()).map_err(|e| {
+                    ArrowError::CastError(format!(
+                        "fail to cast arrow_rs::NullBuffer to arrow2::Bitmap: {e}"
+                    ))
+                })?;
+            let column = NullableColumn { column, validity };
+            Ok(Column::Nullable(Box::new(column)))
+        } else {
+            Ok(column)
+        }
+    }
 }

src/query/expression/src/convert_arrow_rs/record_batch.rs

@@ -17,6 +17,7 @@ use std::sync::Arc;
 use arrow_array::RecordBatch;
 use arrow_schema::ArrowError;
 
+use crate::Column;
 use crate::DataBlock;
 use crate::DataSchema;
 
@@ -30,4 +31,13 @@ impl DataBlock {
         let schema = Arc::new(data_schema.into());
         RecordBatch::try_new(schema, arrays)
     }
+
+    pub fn from_record_batch(batch: &RecordBatch) -> Result<(Self, DataSchema), ArrowError> {
+        let mut columns = Vec::with_capacity(batch.columns().len());
+        for array in batch.columns() {
+            columns.push(Column::from_arrow_rs(array.clone())?)
+        }
+        let schema = DataSchema::try_from(&(*batch.schema()))?;
+        Ok((DataBlock::new_from_columns(columns), schema))
+    }
 }

src/query/expression/src/convert_arrow_rs/schema.rs

@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use arrow_schema::ArrowError;
 use arrow_schema::DataType as ArrowDataType;
 use arrow_schema::Field as ArrowField;
 use arrow_schema::Schema as ArrowSchema;
 use arrow_schema::TimeUnit;
 
+use crate::types::decimal::DecimalSize;
 use crate::types::DataType;
 use crate::types::DecimalDataType;
 use crate::types::NumberDataType;
@@ -127,3 +129,73 @@ impl From<&DataSchema> for ArrowSchema {
         }
     }
 }
+
+impl TryFrom<&ArrowField> for DataField {
+    type Error = ArrowError;
+
+    fn try_from(f: &ArrowField) -> Result<Self, ArrowError> {
+        let ty = f.data_type().try_into()?;
+        if f.is_nullable() {
+            Ok(DataField::new_nullable(f.name().as_str(), ty))
+        } else {
+            Ok(DataField::new(f.name(), ty))
+        }
+    }
+}
+
+impl TryFrom<&ArrowSchema> for DataSchema {
+    type Error = ArrowError;
+
+    fn try_from(schema: &ArrowSchema) -> Result<Self, ArrowError> {
+        let mut fields = vec![];
+        for field in &schema.fields {
+            fields.push(DataField::try_from(field)?)
+        }
+        Ok(DataSchema {
+            fields,
+            metadata: Default::default(),
+        })
+    }
+}
+
+impl TryFrom<&ArrowDataType> for DataType {
+    type Error = ArrowError;
+
+    fn try_from(ty: &ArrowDataType) -> Result<Self, ArrowError> {
+        let data_type = match ty {
+            ArrowDataType::Null => DataType::Null,
+            ArrowDataType::Boolean => DataType::Boolean,
+            ArrowDataType::Int8 => DataType::Number(NumberDataType::Int8),
+            ArrowDataType::Int16 => DataType::Number(NumberDataType::Int16),
+            ArrowDataType::Int32 => DataType::Number(NumberDataType::Int32),
+            ArrowDataType::Int64 => DataType::Number(NumberDataType::Int64),
+            ArrowDataType::UInt8 => DataType::Number(NumberDataType::UInt8),
+            ArrowDataType::UInt16 => DataType::Number(NumberDataType::UInt16),
+            ArrowDataType::UInt32 => DataType::Number(NumberDataType::UInt32),
+            ArrowDataType::UInt64 => DataType::Number(NumberDataType::UInt64),
+            ArrowDataType::Float32 | ArrowDataType::Float16 => {
+                DataType::Number(NumberDataType::Float32)
+            }
+            ArrowDataType::Float64 => DataType::Number(NumberDataType::Float64),
+            ArrowDataType::Timestamp(_unit, _tz) => DataType::Timestamp,
+            ArrowDataType::Date32 | ArrowDataType::Date64 => DataType::Date,
+            ArrowDataType::Utf8 | ArrowDataType::LargeUtf8 => DataType::String,
+            ArrowDataType::Decimal128(p, s) => {
+                DataType::Decimal(DecimalDataType::Decimal128(DecimalSize {
+                    precision: *p,
+                    scale: (*s) as u8,
+                }))
+            }
+            ArrowDataType::Decimal256(p, s) => {
+                DataType::Decimal(DecimalDataType::Decimal256(DecimalSize {
+                    precision: *p,
+                    scale: (*s) as u8,
+                }))
+            }
+            _ => Err(ArrowError::CastError(format!(
+                "cast {ty} to DataType not not implemented yet"
+            )))?,
+        };
+        Ok(data_type)
+    }
+}

src/query/service/src/servers/flight_sql/flight_sql_service/query.rs

@@ -61,7 +61,30 @@ impl FlightSqlServiceImpl {
         planner.plan_sql(query).await
     }
 
-    pub(super) async fn execute_plan(
+    pub(super) async fn execute_update(
+        &self,
+        session: Arc<Session>,
+        plan: &Plan,
+        plan_extras: &PlanExtras,
+    ) -> Result<i64> {
+        let context = session
+            .create_query_context()
+            .await
+            .map_err(|e| status!("Could not create_query_context", e))?;
+
+        context.attach_query_str(plan.to_string(), plan_extras.stament.to_mask_sql());
+        let interpreter = InterpreterFactory::get(context.clone(), plan).await?;
+
+        let mut blocks = interpreter.execute(context.clone()).await?;
+        while let Some(block) = blocks.next().await {
+            block?;
+        }
+
+        let affected_rows = context.get_write_progress_value().rows;
+        Ok(affected_rows as i64)
+    }
+
+    pub(super) async fn execute_query(
         &self,
         session: Arc<Session>,
         plan: &Plan,