send schema before batches

Author: nevi-me

rust/arrow/src/flight/mod.rs

@@ -19,10 +19,11 @@
 
 use flight::FlightData;
 
+use crate::datatypes::Schema;
 use crate::ipc::writer;
 use crate::record_batch::RecordBatch;
 
-/// Convert a `RecordBatch` to `FlightData by getting the header and body as bytes
+/// Convert a `RecordBatch` to `FlightData` by getting the header and body as bytes
 impl From<&RecordBatch> for FlightData {
     fn from(batch: &RecordBatch) -> Self {
         let (header, body) = writer::record_batch_to_bytes(batch);
@@ -35,4 +36,17 @@ impl From<&RecordBatch> for FlightData {
     }
 }
 
+/// Convert a `Schema` to `FlightData` by converting to an IPC message
+impl From<&Schema> for FlightData {
+    fn from(schema: &Schema) -> Self {
+        let schema = writer::schema_to_bytes(schema);
+        Self {
+            flight_descriptor: None,
+            app_metadata: vec![],
+            data_header: schema,
+            data_body: vec![],
+        }
+    }
+}
+
 // TODO: add more explicit conversion that expoess flight descriptor and metadata options

rust/arrow/src/ipc/writer.rs

@@ -209,8 +209,7 @@ impl<W: Write> Drop for StreamWriter<W> {
     }
 }
 
-/// Convert the schema to its IPC representation, and write it to the `writer`
-fn write_schema<R: Write>(writer: &mut BufWriter<R>, schema: &Schema) -> Result<usize> {
+pub(crate) fn schema_to_bytes(schema: &Schema) -> Vec<u8> {
     let mut fbb = FlatBufferBuilder::new();
     let schema = {
         let fb = ipc::convert::schema_to_fb_offset(&mut fbb, schema);
@@ -227,9 +226,13 @@ fn write_schema<R: Write>(writer: &mut BufWriter<R>, schema: &Schema) -> Result<
     fbb.finish(data, None);
 
     let data = fbb.finished_data();
-    let written = write_padded_data(writer, data, WriteDataType::Header);
+    data.to_vec()
+}
 
-    written
+/// Convert the schema to its IPC representation, and write it to the `writer`
+fn write_schema<R: Write>(writer: &mut BufWriter<R>, schema: &Schema) -> Result<usize> {
+    let data = schema_to_bytes(schema);
+    write_padded_data(writer, &data[..], WriteDataType::Header)
 }
 
 /// The message type being written. This determines whether to write the data length or not.

rust/datafusion/examples/flight-server.rs

@@ -77,12 +77,24 @@ impl FlightService for FlightServiceImpl {
 
                 // execute the query
                 let results = ctx.collect(plan.as_ref()).map_err(|e| to_tonic_err(&e))?;
+                if results.is_empty() {
+                    return Err(Status::internal("There were no results from ticket"));
+                }
 
-                let flights: Vec<Result<FlightData, Status>> = results
+                // add an initial FlightData message that sends schema
+                // TODO: find a more ergonomic way of doing this
+                let schema = results[0].schema();
+                let mut flights: Vec<Result<FlightData, Status>> =
+                    vec![Ok(FlightData::from(schema.as_ref()))];
+
+                let mut batches: Vec<Result<FlightData, Status>> = results
                     .iter()
                     .map(|batch| Ok(FlightData::from(batch)))
                     .collect();
 
+                // append batch vector to schema vector, so that the first message sent is the schema
+                flights.append(&mut batches);
+
                 let output = futures::stream::iter(flights);
 
                 Ok(Response::new(Box::pin(output) as Self::DoGetStream))