clean up python tpch client a bit

Author: robtandy

scripts/launch_python_arrowflightsql_client.sh

@@ -185,7 +185,7 @@ def run_sql(sql_query):
             rich_table.add_row(*row_data)
 
         console = Console()
-        console.print(rich_table)
+        console.print(rich_table, markup=False)
         
     except Exception as e:
         print(f"Error executing SQL query: {str(e)}")

src/analyze.rs

@@ -115,7 +115,7 @@ pub struct DistributedAnalyzeRootExec {
 impl DistributedAnalyzeRootExec {
     pub fn new(input: Arc<dyn ExecutionPlan>, verbose: bool, show_statistics: bool) -> Self {
         let field_a = Field::new("Task", DataType::Utf8, false);
-        let field_b = Field::new("Plan", DataType::Utf8, false);
+        let field_b = Field::new("Plan with Metrics", DataType::Utf8, false);
         let schema = Arc::new(Schema::new(vec![field_a, field_b]));
 
         let properties = PlanProperties::new(
@@ -226,7 +226,7 @@ impl ExecutionPlan for DistributedAnalyzeRootExec {
             .clone();
 
         // we want to gather all partitions
-        let coalesce = CoalescePartitionsExec::new(self.input.clone());
+        let coalesce = Arc::new(CoalescePartitionsExec::new(self.input.clone()));
 
         let mut input_stream = coalesce.execute(partition, context)?;
 
@@ -272,10 +272,10 @@ impl ExecutionPlan for DistributedAnalyzeRootExec {
 
             tasks.sort_by_key(|t| (t.stage_id, t.partition_group.clone()));
 
+            trace!("sorted tasks: {:?}", tasks);
+
             let mut task_builder = StringBuilder::with_capacity(1, 1024);
             let mut plan_builder = StringBuilder::with_capacity(1, 1024);
-            task_builder.append_value("Task");
-            plan_builder.append_value("Plan with Metrics");
 
             for task_output in tasks.iter() {
                 task_builder.append_value(format!(
@@ -285,7 +285,7 @@ impl ExecutionPlan for DistributedAnalyzeRootExec {
                     task_output
                         .host
                         .as_ref()
-                        .map(|h| h.to_string())
+                        .map(|h| format!("{} {}", h.name, h.addr))
                         .unwrap_or("Unknown".to_string())
                 ));
                 plan_builder.append_value(&task_output.plan);

src/query_planner.rs

@@ -1,12 +1,15 @@
 use std::sync::Arc;
 
 use anyhow::anyhow;
-use arrow::datatypes::SchemaRef;
+use arrow::{compute::concat_batches, datatypes::SchemaRef};
 use datafusion::{
-    logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext,
+    logical_expr::LogicalPlan,
+    physical_plan::{coalesce_partitions::CoalescePartitionsExec, ExecutionPlan},
+    prelude::SessionContext,
 };
 
 use datafusion_substrait::{logical_plan::consumer::from_substrait_plan, substrait::proto::Plan};
+use tokio_stream::StreamExt;
 
 use crate::{
     explain::build_explain_batch,
@@ -77,6 +80,7 @@ impl QueryPlanner {
         match logical_plan {
             p @ LogicalPlan::Explain(_) => self.prepare_explain(p, ctx).await,
             // add other logical plans for local execution here following the pattern for explain
+            p @ LogicalPlan::DescribeTable(_) => self.prepare_local(p, ctx).await,
             p => self.prepare_query(p, ctx).await,
         }
     }
@@ -89,6 +93,7 @@ impl QueryPlanner {
         match logical_plan {
             p @ LogicalPlan::Explain(_) => self.prepare_explain(p, ctx).await,
             // add other logical plans for local execution here following the pattern for explain
+            p @ LogicalPlan::DescribeTable(_) => self.prepare_local(p, ctx).await,
             p => self.prepare_query(p, ctx).await,
         }
     }
@@ -103,44 +108,30 @@ impl QueryPlanner {
         self.send_it(logical_plan, physical_plan, ctx).await
     }
 
-    async fn send_it(
+    async fn prepare_local(
         &self,
         logical_plan: LogicalPlan,
-        physical_plan: Arc<dyn ExecutionPlan>,
         ctx: SessionContext,
     ) -> Result<QueryPlan> {
-        let query_id = uuid::Uuid::new_v4().to_string();
-
-        // divide the physical plan into chunks (tasks) that we can distribute to workers
-        let (distributed_plan, distributed_stages) =
-            execution_planning(physical_plan.clone(), 8192, Some(2)).await?;
+        let physical_plan = physical_planning(&logical_plan, &ctx).await?;
 
-        let worker_addrs = get_worker_addresses()?;
+        // execute it locally
+        let mut stream =
+            Arc::new(CoalescePartitionsExec::new(physical_plan)).execute(0, ctx.task_ctx())?;
+        let mut batches = vec![];
 
-        // gather some information we need to send back such that
-        // we can send a ticket to the client
-        let final_stage = &distributed_stages[distributed_stages.len() - 1];
-        let schema = Arc::clone(&final_stage.plan.schema());
-        let final_stage_id = final_stage.stage_id;
+        while let Some(batch) = stream.next().await {
+            batches.push(batch?);
+        }
 
-        // distribute the stages to workers, further dividing them up
-        // into chunks of partitions (partition_groups)
-        let (final_workers, tasks) =
-            distribute_stages(&query_id, distributed_stages, worker_addrs).await?;
+        if batches.is_empty() {
+            return Err(anyhow!("No data returned from local execution").into());
+        }
 
-        let qp = QueryPlan {
-            query_id,
-            session_context: ctx,
-            worker_addresses: final_workers,
-            final_stage_id,
-            schema,
-            logical_plan,
-            physical_plan,
-            distributed_plan,
-            distributed_tasks: tasks,
-        };
+        let combined_batch = concat_batches(&batches[0].schema(), &batches)?;
+        let physical_plan = Arc::new(RecordBatchExec::new(combined_batch));
 
-        Ok(qp)
+        self.send_it(logical_plan, physical_plan, ctx).await
     }
 
     async fn prepare_explain(
@@ -172,4 +163,43 @@ impl QueryPlanner {
         )
         .await
     }
+    async fn send_it(
+        &self,
+        logical_plan: LogicalPlan,
+        physical_plan: Arc<dyn ExecutionPlan>,
+        ctx: SessionContext,
+    ) -> Result<QueryPlan> {
+        let query_id = uuid::Uuid::new_v4().to_string();
+
+        // divide the physical plan into chunks (tasks) that we can distribute to workers
+        let (distributed_plan, distributed_stages) =
+            execution_planning(physical_plan.clone(), 8192, Some(2)).await?;
+
+        let worker_addrs = get_worker_addresses()?;
+
+        // gather some information we need to send back such that
+        // we can send a ticket to the client
+        let final_stage = &distributed_stages[distributed_stages.len() - 1];
+        let schema = Arc::clone(&final_stage.plan.schema());
+        let final_stage_id = final_stage.stage_id;
+
+        // distribute the stages to workers, further dividing them up
+        // into chunks of partitions (partition_groups)
+        let (final_workers, tasks) =
+            distribute_stages(&query_id, distributed_stages, worker_addrs).await?;
+
+        let qp = QueryPlan {
+            query_id,
+            session_context: ctx,
+            worker_addresses: final_workers,
+            final_stage_id,
+            schema,
+            logical_plan,
+            physical_plan,
+            distributed_plan,
+            distributed_tasks: tasks,
+        };
+
+        Ok(qp)
+    }
 }

tests/common/mod.rs

@@ -220,7 +220,7 @@ except Exception as e:
     print(f"Error executing distributed query: {{str(e)}}", file=sys.stderr)
     sys.exit(1)
 "#,
-            self.get_proxy_address().split(':').last().unwrap()
+            self.get_proxy_address().split(':').next_back().unwrap()
         );
 
         let script_path = Self::write_temp_file(
@@ -271,15 +271,13 @@ except Exception as e:
         for port in &ports {
             // Find and kill processes using lsof
             if let Ok(output) = Command::new("lsof")
-                .args(&["-ti", &format!(":{}", port)])
+                .args(["-ti", &format!(":{}", port)])
                 .output()
             {
                 let stdout = String::from_utf8_lossy(&output.stdout);
                 for line in stdout.lines() {
                     if let Ok(pid) = line.trim().parse::<u32>() {
-                        let _ = Command::new("kill")
-                            .args(&["-9", &pid.to_string()])
-                            .output();
+                        let _ = Command::new("kill").args(["-9", &pid.to_string()]).output();
                         println!("  🔥 Killed process {} on port {}", pid, port);
                     }
                 }
@@ -362,7 +360,7 @@ except Exception as e:
 
         // Check if packages are already installed
         let check_cmd = Command::new("python3")
-            .args(&["-c", "import adbc_driver_manager; import adbc_driver_flightsql; import duckdb; import pyarrow; print('OK')"])
+            .args(["-c", "import adbc_driver_manager; import adbc_driver_flightsql; import duckdb; import pyarrow; print('OK')"])
             .output();
 
         if let Ok(output) = check_cmd {
@@ -449,7 +447,7 @@ except Exception as e:
             let worker = Self::spawn_process(
                 binary_path_str,
                 &["--mode", "worker", "--port", &port.to_string()],
-                &[("DFRAY_TABLES", &tpch_tables), ("DFRAY_VIEWS", &tpch_views)],
+                &[("DFRAY_TABLES", &tpch_tables), ("DFRAY_VIEWS", tpch_views)],
                 &format!("start worker {}", i + 1),
             )?;
             self.worker_processes.push(worker);
@@ -473,7 +471,7 @@ except Exception as e:
             &[
                 ("DFRAY_WORKER_ADDRESSES", &worker_addresses),
                 ("DFRAY_TABLES", &tpch_tables),
-                ("DFRAY_VIEWS", &tpch_views),
+                ("DFRAY_VIEWS", tpch_views),
             ],
             "start proxy",
         )?;
@@ -819,7 +817,7 @@ pub fn batches_to_sorted_strings(
     let mut in_table = false;
     let mut header_found = false;
 
-    for (_line_idx, line) in lines.iter().enumerate() {
+    for line in lines.iter() {
         let trimmed = line.trim();
 
         // Detect table boundaries (+---+)