feat: use arrow-ipc to communicate between remote server and client

Cargo.toml

@@ -52,6 +52,7 @@ path = "src/bin/ceresdb-server.rs"
 
 [workspace.dependencies]
 arrow = { version = "23.0.0", features = ["prettyprint"] }
+arrow_ipc = { version = "23.0.0" }
 arrow_ext = { path = "components/arrow_ext" }
 analytic_engine = { path = "analytic_engine" }
 arena = { path = "components/arena" }

benchmarks/Cargo.toml

@@ -15,6 +15,7 @@ analytic_engine = { workspace = true }
 arena = { workspace = true }
 arrow = { workspace = true }
 arrow2 = { version = "0.12.0", features = [ "io_parquet" ] }
+arrow_ext = { workspace = true }
 base64 = { workspace = true }
 clap = { workspace = true }
 common_types = { workspace = true }

components/arrow_ext/Cargo.toml

@@ -9,10 +9,7 @@ workspace = true
 
 [package.edition]
 workspace = true
-[dependencies.uncover]
-git = "https://github.com/matklad/uncover.git"
-rev = "1d0770d997e29731b287e9e11e4ffbbea5f456da"
 
 [dependencies]
 arrow = { workspace = true }
-
+snafu = { workspace = true }

components/arrow_ext/src/ipc.rs

@@ -0,0 +1,70 @@
+// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
+
+//! Utilities for `RecordBatch` serialization using Arrow IPC
+
+use std::io::Cursor;
+
+use arrow::{
+    ipc::{reader::StreamReader, writer::StreamWriter},
+    record_batch::RecordBatch,
+};
+use snafu::{Backtrace, OptionExt, ResultExt, Snafu};
+
+#[derive(Snafu, Debug)]
+#[snafu(visibility(pub))]
+pub enum Error {
+    #[snafu(display("Arror error, err:{}.\nBacktrace:\n{}", source, backtrace))]
+    ArrowError {
+        source: arrow::error::ArrowError,
+        backtrace: Backtrace,
+    },
+
+    #[snafu(display("Failed to decode record batch.\nBacktrace:\n{}", backtrace))]
+    Decode { backtrace: Backtrace },
+}
+
+type Result<T> = std::result::Result<T, Error>;
+
+pub fn encode_record_batch(batch: &RecordBatch) -> Result<Vec<u8>> {
+    let buffer: Vec<u8> = Vec::new();
+    let mut stream_writer = StreamWriter::try_new(buffer, &batch.schema()).context(ArrowError)?;
+    stream_writer.write(batch).context(ArrowError)?;
+    stream_writer.into_inner().context(ArrowError)
+}
+
+pub fn decode_record_batch(bytes: Vec<u8>) -> Result<RecordBatch> {
+    let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).context(ArrowError)?;
+    stream_reader.next().context(Decode)?.context(ArrowError)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use arrow::{
+        array::{Int32Array, StringArray},
+        datatypes::{DataType, Field, Schema},
+    };
+
+    use super::*;
+
+    fn create_batch(rows: usize) -> RecordBatch {
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Utf8, false),
+        ]);
+
+        let a = Int32Array::from_iter_values(0..rows as i32);
+        let b = StringArray::from_iter_values((0..rows).map(|i| i.to_string()));
+
+        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
+    }
+
+    #[test]
+    fn test_ipc_encode_decode() {
+        let batch = create_batch(1024);
+        let bytes = encode_record_batch(&batch).unwrap();
+
+        assert_eq!(batch, decode_record_batch(bytes).unwrap());
+    }
+}

components/arrow_ext/src/lib.rs

@@ -1,4 +1,5 @@
 // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
 
 pub mod display;
+pub mod ipc;
 pub mod operation;

proto/protos/remote_engine.proto

@@ -54,7 +54,7 @@ message ReadResponse {
   ResponseHeader header = 1;
   // Version of row encoding method
   uint32 version = 2;
-  repeated bytes rows = 3;
+  bytes rows = 3;
 }
 
 message RowGroup {

remote_engine_client/Cargo.toml

@@ -11,6 +11,7 @@ workspace = true
 workspace = true
 
 [dependencies]
+arrow_ext = { workspace = true }
 async-trait = { workspace = true }
 ceresdbproto = { workspace = true }
 clru = { workspace = true }

remote_engine_client/src/client.rs

@@ -7,11 +7,11 @@ use std::{
     task::{Context, Poll},
 };
 
+use arrow_ext::ipc;
 use ceresdbproto::storage;
 use common_types::{
     projected_schema::ProjectedSchema, record_batch::RecordBatch, schema::RecordSchema,
 };
-use common_util::avro;
 use futures::{Stream, StreamExt};
 use proto::remote_engine::{self, remote_engine_service_client::*};
 use router::{endpoint::Endpoint, RouterRef};
@@ -178,15 +178,19 @@ impl Stream for ClientReadRecordBatchStream {
                     }.fail()));
                 }
 
-                // It's ok, try to convert rows to record batch and return.
-                let record_schema = this.projected_schema.to_record_schema();
-                let record_batch_result =
-                    avro::avro_rows_to_record_batch(response.rows, record_schema)
-                        .map_err(|e| Box::new(e) as _)
-                        .context(ConvertReadResponse {
-                            msg: "build record batch failed",
-                        });
-                Poll::Ready(Some(record_batch_result))
+                let record_batch = ipc::decode_record_batch(response.rows)
+                    .map_err(|e| Box::new(e) as _)
+                    .context(ConvertReadResponse {
+                        msg: "decode record batch failed",
+                    })
+                    .and_then(|v| {
+                        RecordBatch::try_from(v)
+                            .map_err(|e| Box::new(e) as _)
+                            .context(ConvertReadResponse {
+                                msg: "convert record batch failed",
+                            })
+                    });
+                Poll::Ready(Some(record_batch))
             }
 
             Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e).context(Rpc {

server/Cargo.toml

@@ -13,6 +13,7 @@ workspace = true
 [dependencies]
 analytic_engine = { workspace = true }
 arrow = { workspace = true }
+arrow_ext = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 catalog = { workspace = true }

server/src/grpc/remote_engine_service/mod.rs

@@ -4,10 +4,10 @@
 
 use std::sync::Arc;
 
+use arrow_ext::ipc;
 use async_trait::async_trait;
 use catalog::manager::ManagerRef;
 use common_types::record_batch::RecordBatch;
-use common_util::avro;
 use futures::stream::{self, BoxStream, StreamExt};
 use log::error;
 use proto::remote_engine::{
@@ -137,22 +137,23 @@ impl<Q: QueryExecutor + 'static> RemoteEngineService for RemoteEngineServiceImpl
             Ok(stream) => {
                 let new_stream: Self::ReadStream = Box::pin(stream.map(|res| match res {
                     Ok(record_batch) => {
-                        let resp = match avro::record_batch_to_avro_rows(&record_batch)
-                            .map_err(|e| Box::new(e) as _)
-                            .context(ErrWithCause {
-                                code: StatusCode::Internal,
-                                msg: "fail to convert record batch to avro",
-                            }) {
-                            Err(e) => ReadResponse {
-                                header: Some(error::build_err_header(e)),
-                                ..Default::default()
-                            },
-                            Ok(rows) => ReadResponse {
-                                header: Some(build_ok_header()),
-                                version: ENCODE_ROWS_WITH_AVRO,
-                                rows,
-                            },
-                        };
+                        let resp =
+                            match ipc::encode_record_batch(&record_batch.into_arrow_record_batch())
+                                .map_err(|e| Box::new(e) as _)
+                                .context(ErrWithCause {
+                                    code: StatusCode::Internal,
+                                    msg: "fail to convert record batch to avro",
+                                }) {
+                                Err(e) => ReadResponse {
+                                    header: Some(error::build_err_header(e)),
+                                    ..Default::default()
+                                },
+                                Ok(rows) => ReadResponse {
+                                    header: Some(build_ok_header()),
+                                    version: ENCODE_ROWS_WITH_AVRO,
+                                    rows,
+                                },
+                            };
 
                         Ok(resp)
                     }