Add arrow flight reads

Author: gabotechs

Cargo.lock

@@ -16,6 +16,7 @@ serde = "1.0.219"
 serde_json = "1.0.141"
 env_logger = "0.11.8"
 async-trait = "0.1.88"
+async-stream = "0.3.6"
 chrono = "0.4.41"
 futures = "0.3.31"
 dashmap = "6.1.0"

benchmarks/Cargo.toml

@@ -7,7 +7,7 @@ export const ROOT = path.join(__dirname, '../../..')
 // Simple data structures
 export type QueryResult = {
   query: string;
-  iterations: { elapsed: number; row_count: number }[];
+  iterations: { elapsed: number; row_count?: number }[];
 }
 
 export type BenchmarkResults = {
@@ -19,7 +19,7 @@ export const BenchmarkResults = z.object({
     query: z.string(),
     iterations: z.array(z.object({
       elapsed: z.number(),
-      row_count: z.number()
+      row_count: z.number().optional()
     }))
   }))
 })
@@ -69,7 +69,7 @@ export async function compareWithPrevious(results: BenchmarkResults, outputPath:
 export interface BenchmarkRunner {
   createTables(sf: number): Promise<void>;
 
-  executeQuery(query: string): Promise<{ rowCount: number }>;
+  executeQuery(query: string): Promise<{ rowCount?: number }>;
 }
 
 export async function runBenchmark(
@@ -103,6 +103,9 @@ export async function runBenchmark(
       iterations: []
     };
 
+    console.log(`Warming up query ${id}...`)
+    await runner.executeQuery(queryToExecute);
+   
     for (let i = 0; i < iterations; i++) {
       const start = new Date()
       const response = await runner.executeQuery(queryToExecute);

benchmarks/cdk/bin/@bench-common.ts

@@ -1,6 +1,5 @@
 import path from "path";
 import {Command} from "commander";
-import {z} from 'zod';
 import {BenchmarkRunner, ROOT, runBenchmark} from "./@bench-common";
 
 // Remember to port-forward a worker with
@@ -40,12 +39,6 @@ async function main() {
     });
 }
 
-const QueryResponse = z.object({
-    count: z.number(),
-    plan: z.string()
-})
-type QueryResponse = z.infer<typeof QueryResponse>
-
 class DataFusionRunner implements BenchmarkRunner {
     private url = 'http://localhost:9000';
 
@@ -55,7 +48,7 @@ class DataFusionRunner implements BenchmarkRunner {
     }) {
     }
 
-    async executeQuery(sql: string): Promise<{ rowCount: number }> {
+    async executeQuery(sql: string): Promise<{ rowCount?: number }> {
         let response
         if (sql.includes("create view")) {
             // This is query 15
@@ -67,10 +60,10 @@ class DataFusionRunner implements BenchmarkRunner {
             response = await this.query(sql)
         }
 
-        return {rowCount: response.count};
+        return response
     }
 
-    private async query(sql: string): Promise<QueryResponse> {
+    private async query(sql: string): Promise<{ rowCount?: number }> {
         const url = new URL(this.url);
         url.searchParams.set('sql', sql);
 
@@ -80,11 +73,13 @@ class DataFusionRunner implements BenchmarkRunner {
             const msg = await response.text();
             throw new Error(`Query failed: ${response.status} ${msg}`);
         }
-
-        const unparsed = await response.json();
-        return QueryResponse.parse(unparsed);
+        await response.arrayBuffer()
+        // TODO: the library still cannot decode Utf8Views, and some queries fail because of that
+        // const table = await tableFromIPC(ipcData)
+        return {}
     }
 
+
     async createTables(sf: number): Promise<void> {
         let stmt = '';
         for (const tbl of [

benchmarks/cdk/bin/datafusion-bench.ts

@@ -2,8 +2,11 @@ use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use aws_config::BehaviorVersion;
 use aws_sdk_ec2::Client as Ec2Client;
+use axum::body::Body;
+use axum::http::Response;
 use axum::{Json, Router, extract::Query, http::StatusCode, routing::get};
 use dashmap::{DashMap, Entry};
+use datafusion::arrow::ipc::writer::{IpcWriteOptions, StreamWriter};
 use datafusion::common::DataFusionError;
 use datafusion::common::instant::Instant;
 use datafusion::common::runtime::SpawnedTask;
@@ -15,7 +18,7 @@ use datafusion_distributed::{
     DistributedPhysicalOptimizerRule, DistributedSessionBuilder, DistributedSessionBuilderContext,
     create_flight_client, display_plan_ascii,
 };
-use futures::{StreamExt, TryFutureExt};
+use futures::{StreamExt, TryFutureExt, TryStreamExt};
 use log::{error, info, warn};
 use object_store::ObjectStore;
 use object_store::aws::AmazonS3Builder;
@@ -27,15 +30,10 @@ use std::sync::atomic::AtomicBool;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use structopt::StructOpt;
+use tonic::Status;
 use tonic::transport::{Channel, Server};
 use url::Url;
 
-#[derive(Serialize)]
-struct QueryResult {
-    plan: String,
-    count: usize,
-}
-
 #[derive(Debug, StructOpt, Clone)]
 #[structopt(about = "worker spawn command")]
 struct Cmd {
@@ -115,7 +113,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
                 let ctx = ctx.clone();
 
                 async move {
-                    let sql = params.get("sql").ok_or(err("Missing 'sql' parameter"))?;
+                    let sql = params
+                        .get("sql")
+                        .ok_or(err("Missing 'sql' parameter"))?
+                        .clone();
 
                     let mut df_opt = None;
                     for sql in sql.split(";") {
@@ -129,12 +130,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
                         return Err(err("Empty 'sql' parameter"));
                     };
 
-                    let start = Instant::now();
-
                     info!("Executing query...");
                     let abort_notifier = AbortNotifier::new("Query aborted");
                     let abort_notifier_clone = abort_notifier.clone();
-                    let task = SpawnedTask::spawn(async move {
+                    let still_running_log_task = SpawnedTask::spawn(async move {
                         let _ = abort_notifier_clone;
                         loop {
                             tokio::time::sleep(Duration::from_secs(5)).await;
@@ -144,25 +143,42 @@ async fn main() -> Result<(), Box<dyn Error>> {
                     let physical = df.create_physical_plan().await.map_err(err)?;
                     let mut stream =
                         execute_stream(physical.clone(), ctx.task_ctx()).map_err(err)?;
-                    let mut count = 0;
-                    while let Some(batch) = stream.next().await {
-                        count += batch.map_err(err)?.num_rows();
-                        info!("Gathered {count} rows, query still in progress..")
-                    }
-                    let plan = display_plan_ascii(physical.as_ref(), true);
-                    drop(task);
 
-                    let elapsed = start.elapsed();
-                    let ms = elapsed.as_secs_f64() * 1000.0;
-                    info!("Finished executing query:\n{sql}\n\n{plan}");
-                    info!("Returned {count} rows in {ms} ms");
-                    abort_notifier.finished();
+                    let start = Instant::now();
+                    let mut count: usize = 0;
+
+                    let stream = async_stream::stream! {
+                        // Stream the data
+                        while let Some(batch) = stream.next().await {
+                            let batch = batch?;
+                            count += batch.num_rows();
+                            info!("Gathered {count} rows, query still in progress..");
+
+                            let mut writer = StreamWriter::try_new(vec![], batch.schema().as_ref())?;
+                            writer.write(&batch)?;
+                            yield writer.into_inner()
+                        }
+
+                        // After stream completes gracefully - all cleanup code runs here
+                        let elapsed = start.elapsed();
+                        let ms = elapsed.as_secs_f64() * 1000.0;
+                        info!("Finished executing query:\n{sql}");
+                        info!("Returned {count} rows in {ms} ms");
+                        abort_notifier.finished();
+                        // keep the task alive
+                        drop(still_running_log_task);
+                    };
 
-                    Ok::<_, (StatusCode, String)>(Json(QueryResult { count, plan }))
+                    Ok(Response::builder()
+                        .header("content-type", "application/octet-stream")
+                        .body(Body::from_stream(
+                            stream.map_err(|err| Status::internal(err.to_string())),
+                        ))
+                        .expect("building a Response from a body should never fail"))
                 }
-                .inspect_err(|(_, msg)| {
-                    error!("Error executing query: {msg}");
-                })
+                    .inspect_err(|(_, msg)| {
+                        error!("Error executing query: {msg}");
+                    })
             }),
         ),
     );