fix(cluster): fix precommit deser and ser

Author: zhang2014

src/query/service/src/api/mod.rs

@@ -30,6 +30,7 @@ pub use rpc::FragmentPayload;
 pub use rpc::FragmentPlanPacket;
 pub use rpc::InitNodesChannelPacket;
 pub use rpc::MergeExchange;
+pub use rpc::PrecommitBlock;
 pub use rpc::QueryFragmentsPlanPacket;
 pub use rpc::ServerFlightExchange;
 pub use rpc::ShuffleDataExchange;

src/query/service/src/api/rpc/mod.rs

@@ -43,4 +43,5 @@ pub use packets::FragmentPayload;
 pub use packets::FragmentPlanPacket;
 pub use packets::InitNodesChannelPacket;
 pub use packets::Packet;
+pub use packets::PrecommitBlock;
 pub use packets::QueryFragmentsPlanPacket;

src/query/service/src/api/rpc/packets/packet_data.rs

@@ -26,8 +26,6 @@ use byteorder::ReadBytesExt;
 use byteorder::WriteBytesExt;
 use common_arrow::arrow::io::flight::deserialize_batch;
 use common_arrow::arrow::io::flight::serialize_batch;
-use common_arrow::arrow::io::flight::serialize_schema_to_info;
-use common_arrow::arrow::io::ipc::read::deserialize_schema;
 use common_arrow::arrow::io::ipc::write::default_ipc_fields;
 use common_arrow::arrow::io::ipc::write::WriteOptions;
 use common_arrow::arrow::io::ipc::IpcSchema;
@@ -67,7 +65,7 @@ pub enum ProgressInfo {
     ResultProgress(ProgressValues),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Eq, PartialEq, Clone)]
 pub struct PrecommitBlock(pub DataBlock);
 
 #[derive(Debug)]
@@ -187,20 +185,21 @@ impl From<FragmentData> for FlightData {
 impl PrecommitBlock {
     pub fn write<T: Write>(self, bytes: &mut T) -> Result<()> {
         let data_block = self.0;
-        let arrow_schema = data_block.schema().to_arrow();
-        let schema = serialize_schema_to_info(&arrow_schema, None)?;
+        let data_schema = data_block.schema();
+        let serialized_schema = serde_json::to_vec(data_schema)?;
+        let arrow_schema = data_schema.to_arrow();
 
         // schema_flight
         let options = WriteOptions { compression: None };
         let ipc_fields = default_ipc_fields(&arrow_schema.fields);
         let chunks = data_block.try_into()?;
         let (_dicts, data_flight) = serialize_batch(&chunks, &ipc_fields, &options)?;
 
-        bytes.write_u64::<BigEndian>(schema.len() as u64)?;
+        bytes.write_u64::<BigEndian>(serialized_schema.len() as u64)?;
         bytes.write_u64::<BigEndian>(data_flight.data_header.len() as u64)?;
         bytes.write_u64::<BigEndian>(data_flight.data_body.len() as u64)?;
 
-        bytes.write_all(&schema)?;
+        bytes.write_all(&serialized_schema)?;
         bytes.write_all(&data_flight.data_header)?;
         bytes.write_all(&data_flight.data_body)?;
         Ok(())
@@ -218,7 +217,8 @@ impl PrecommitBlock {
         bytes.read_exact(&mut schema)?;
         bytes.read_exact(&mut flight_header)?;
         bytes.read_exact(&mut flight_body)?;
-        let (arrow_schema, _) = deserialize_schema(&schema)?;
+        let data_schema = serde_json::from_slice::<DataSchema>(&schema)?;
+        let arrow_schema = data_schema.to_arrow();
 
         let ipc_fields = default_ipc_fields(&arrow_schema.fields);
         let ipc_schema = IpcSchema {

src/query/service/tests/it/api/mod.rs

@@ -14,4 +14,5 @@
 
 mod http;
 mod http_service;
+mod rpc;
 mod rpc_service;

src/query/service/tests/it/api/rpc/mod.rs

@@ -0,0 +1,15 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod packets;

src/query/service/tests/it/api/rpc/packets/mod.rs

@@ -0,0 +1,15 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod packet_data;

src/query/service/tests/it/api/rpc/packets/packet_data.rs

@@ -0,0 +1,39 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use common_base::base::tokio;
+use common_datablocks::DataBlock;
+use common_datavalues::prelude::*;
+use common_datavalues::DataField;
+use common_datavalues::DataSchema;
+use common_datavalues::UInt8Column;
+use common_exception::Result;
+use databend_query::api::PrecommitBlock;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_precommit_ser_and_deser() -> Result<()> {
+    let data_schema = DataSchema::new(vec![DataField::new("test", UInt8Type::new_impl())]);
+
+    let test_precommit = PrecommitBlock(DataBlock::create(Arc::new(data_schema), vec![Arc::new(
+        UInt8Column::new_from_vec(vec![1, 2, 3]),
+    )]));
+
+    let mut bytes = vec![];
+    PrecommitBlock::write(test_precommit.clone(), &mut bytes)?;
+    let mut read = bytes.as_slice();
+    assert_eq!(test_precommit, PrecommitBlock::read(&mut read)?);
+    Ok(())
+}