[deps] bump arrow2 to d14ae86

common/arrow/Cargo.toml

@@ -20,7 +20,7 @@ simd = ["arrow/simd"]
 # Workspace dependencies
 
 # Github dependencies
-arrow = { package = "arrow2", git = "https://github.com/datafuse-extras/arrow2", default-features = false, rev = "f07cc2c"}
+arrow = { package = "arrow2", git = "https://github.com/datafuse-extras/arrow2", default-features = false, rev = "d14ae86"}
 arrow-format = { version = "0.3.0", features = ["flight-data", "flight-service"] }
 parquet2 = { version = "0.8.1", default_features = false }
 # Crates.io dependencies

query/src/api/rpc/flight_client.rs

@@ -19,7 +19,6 @@ use common_arrow::arrow_format::flight::data::FlightData;
 use common_arrow::arrow_format::flight::data::Ticket;
 use common_arrow::arrow_format::flight::service::flight_service_client::FlightServiceClient;
 use common_base::tokio::time::Duration;
-use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_streams::SendableDataBlockStream;
@@ -45,12 +44,11 @@ impl FlightClient {
     pub async fn fetch_stream(
         &mut self,
         ticket: FlightTicket,
-        schema: DataSchemaRef,
         timeout: u64,
     ) -> Result<SendableDataBlockStream> {
         let ticket = ticket.try_into()?;
         let inner = self.do_get(ticket, timeout).await?;
-        Ok(Box::pin(FlightDataStream::from_remote(schema, inner)))
+        Ok(Box::pin(FlightDataStream::from_remote(inner)))
     }
 
     pub async fn execute_action(&mut self, action: FlightAction, timeout: u64) -> Result<()> {

query/src/api/rpc/flight_client_stream.rs

@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use common_arrow::arrow::io::flight::deserialize_batch;
+use common_arrow::arrow::io::flight::deserialize_schemas;
 use common_arrow::arrow::record_batch::RecordBatch;
 use common_arrow::arrow_format::flight::data::FlightData;
 use common_base::tokio::sync::mpsc::Receiver;
@@ -32,7 +33,6 @@ pub struct FlightDataStream();
 impl FlightDataStream {
     #[inline]
     pub fn from_remote(
-        schema: DataSchemaRef,
         inner: Streaming<FlightData>,
     ) -> impl Stream<Item = Result<DataBlock, ErrorCode>> {
         inner.map(move |flight_data| -> Result<DataBlock, ErrorCode> {
@@ -52,11 +52,17 @@ impl FlightDataStream {
                         )
                     }
 
-                    let arrow_schema = Arc::new(schema.to_arrow());
-                    Ok(
-                        deserialize_batch(&flight_data, arrow_schema, true, &Default::default())
-                            .map(create_data_block)?,
+                    let (arrow_schema, ipc_schema) =
+                        deserialize_schemas(flight_data.data_body.as_slice()).unwrap();
+                    let arrow_schema = Arc::new(arrow_schema);
+
+                    Ok(deserialize_batch(
+                        &flight_data,
+                        arrow_schema,
+                        &ipc_schema,
+                        &Default::default(),
                     )
+                    .map(create_data_block)?)
                 }
             }
         })
@@ -66,7 +72,6 @@ impl FlightDataStream {
     #[inline]
     #[allow(dead_code)]
     pub fn from_receiver(
-        schema_ref: DataSchemaRef,
         inner: Receiver<Result<FlightData, ErrorCode>>,
     ) -> impl Stream<Item = Result<DataBlock, ErrorCode>> {
         ReceiverStream::new(inner).map(move |flight_data| match flight_data {
@@ -83,13 +88,14 @@ impl FlightDataStream {
                     DataBlock::create(Arc::new(schema), columns)
                 }
 
-                Ok(deserialize_batch(
-                    &flight_data,
-                    Arc::new(schema_ref.to_arrow()),
-                    true,
-                    &Default::default(),
+                let (arrow_schema, ipc_schema) =
+                    deserialize_schemas(flight_data.data_body.as_slice()).unwrap();
+                let arrow_schema = Arc::new(arrow_schema);
+
+                Ok(
+                    deserialize_batch(&flight_data, arrow_schema, &ipc_schema, &Default::default())
+                        .map(create_data_block)?,
                 )
-                .map(create_data_block)?)
             }
         })
     }

query/src/api/rpc/flight_service_stream.rs

@@ -49,7 +49,7 @@ impl Stream for FlightDataStream {
             Some(Ok(block)) => match block.try_into() {
                 Err(error) => Some(Err(Status::from(error))),
                 Ok(record_batch) => {
-                    let (dicts, values) = serialize_batch(&record_batch, &self.options);
+                    let (dicts, values) = serialize_batch(&record_batch, &[], &self.options);
 
                     match dicts.is_empty() {
                         true => Some(Ok(values)),

query/src/pipelines/processors/pipeline_builder.rs

@@ -131,7 +131,6 @@ impl PipelineBuilder {
                 flight_ticket,
                 self.ctx.clone(),
                 /* fetch_node_name */ fetch_node.clone(),
-                /* fetch_stream_schema */ plan.schema.clone(),
             )?))?;
         }
 

query/src/pipelines/transforms/transform_remote.rs

@@ -15,7 +15,6 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_streams::SendableDataBlockStream;
@@ -30,7 +29,6 @@ use crate::sessions::QueryContext;
 pub struct RemoteTransform {
     ticket: FlightTicket,
     fetch_node_name: String,
-    schema: DataSchemaRef,
     pub ctx: Arc<QueryContext>,
 }
 
@@ -39,12 +37,10 @@ impl RemoteTransform {
         ticket: FlightTicket,
         context: Arc<QueryContext>,
         fetch_node_name: String,
-        schema: DataSchemaRef,
     ) -> Result<RemoteTransform> {
         Ok(RemoteTransform {
             ticket,
             fetch_node_name,
-            schema,
             ctx: context,
         })
     }
@@ -88,14 +84,11 @@ impl Processor for RemoteTransform {
             self.fetch_node_name
         );
 
-        let data_schema = self.schema.clone();
         let timeout = self.ctx.get_settings().get_flight_client_timeout()?;
 
         let fetch_ticket = self.ticket.clone();
         let mut flight_client = self.flight_client().await?;
-        let fetch_stream = flight_client
-            .fetch_stream(fetch_ticket, data_schema, timeout)
-            .await?;
+        let fetch_stream = flight_client.fetch_stream(fetch_ticket, timeout).await?;
         Ok(Box::pin(self.ctx.try_create_abortable(fetch_stream)?))
     }
 }

query/src/servers/mysql/reject_connection.rs

@@ -38,7 +38,7 @@ impl RejectConnection {
 
         let size = buffer.len().to_le_bytes();
         buffer.splice(0..0, [size[0], size[1], size[2], 2].iter().cloned());
-        stream.write(&buffer).await?;
+        stream.write_all(&buffer).await?;
         stream.flush().await?;
 
         Ok(())
@@ -47,7 +47,7 @@ impl RejectConnection {
     async fn send_handshake(stream: &mut TcpStream) -> Result<()> {
         // Send handshake, packet from msql-srv. Packet[seq = 0]
         stream
-            .write(&[
+            .write_all(&[
                 69, 00, 00, 00, 10, 53, 46, 49, 46, 49, 48, 45, 97, 108, 112, 104, 97, 45, 109,
                 115, 113, 108, 45, 112, 114, 111, 120, 121, 0, 8, 0, 0, 0, 59, 88, 44, 112, 111,
                 95, 107, 125, 0, 0, 66, 33, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 62, 111,
@@ -62,12 +62,12 @@ impl RejectConnection {
 
     async fn receive_handshake_response(stream: &mut TcpStream) -> Result<()> {
         let mut buffer = vec![0; 4];
-        stream.read(&mut buffer).await?;
+        stream.read_exact(&mut buffer).await?;
 
         // Ignore handshake response. Packet[seq = 1]
         let len = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], 0]);
         buffer.resize(len as usize, 0);
-        stream.read(&mut buffer).await?;
+        stream.read_exact(&mut buffer).await?;
 
         Ok(())
     }