Improve the datafusion explain output

Author: jackkleeman

crates/storage-query-datafusion/src/invocation_state/table.rs

@@ -16,6 +16,7 @@ use anyhow::anyhow;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::common::DataFusionError;
+use datafusion::physical_plan::metrics::Time;
 use datafusion::physical_plan::stream::RecordBatchReceiverStream;
 use datafusion::physical_plan::{PhysicalExpr, SendableRecordBatchStream};
 use tokio::sync::mpsc::Sender;
@@ -100,6 +101,7 @@ impl<S: StatusHandle + Send + Sync + Debug + Clone + 'static> ScanPartition for
         _predicate: Option<Arc<dyn PhysicalExpr>>,
         batch_size: usize,
         limit: Option<usize>,
+        _elapsed_compute: Time,
     ) -> anyhow::Result<SendableRecordBatchStream> {
         let status = self.status_handle.clone();
         let partition_store_manager = self.partition_store_manager.clone();

crates/storage-query-datafusion/src/partition_store_scanner.rs

@@ -10,13 +10,15 @@
 
 use std::ops::RangeInclusive;
 use std::sync::Arc;
+use std::time::Instant;
 use std::{fmt::Debug, ops::ControlFlow};
 
 use anyhow::anyhow;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::error::DataFusionError;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::PhysicalExpr;
+use datafusion::physical_plan::metrics::Time;
 use datafusion::physical_plan::stream::RecordBatchReceiverStream;
 
 use restate_partition_store::{PartitionStore, PartitionStoreManager};
@@ -72,6 +74,7 @@ where
     S: ScanLocalPartition<Builder = RB>,
     RB: crate::table_util::Builder + Send + Sync + 'static,
 {
+    #[allow(clippy::too_many_arguments)]
     fn scan_partition(
         &self,
         partition_id: PartitionId,
@@ -80,6 +83,7 @@ where
         predicate: Option<Arc<dyn PhysicalExpr>>,
         batch_size: usize,
         limit: Option<usize>,
+        elapsed_compute: Time,
     ) -> anyhow::Result<SendableRecordBatchStream> {
         let partition_store_manager = self.partition_store_manager.clone();
         let mut stream_builder = RecordBatchReceiverStream::builder(projection.clone(), 1);
@@ -94,9 +98,13 @@ where
                 DataFusionError::External(err.into())
             })?;
 
+            // timer starts on first row, stops on scanner drop
+            let mut elapsed_compute = ElapsedCompute::new(elapsed_compute);
+
             let mut batch_sender = BatchSender::new(projection, tx, predicate, batch_size, limit);
 
             S::for_each_row(&partition_store, range, move |row| {
+                elapsed_compute.start();
                 match S::append_row(batch_sender.builder_mut(), row) {
                     Ok(()) => {}
                     err => return ControlFlow::Break(err),
@@ -127,6 +135,7 @@ where
         predicate: Option<Arc<dyn PhysicalExpr>>,
         batch_size: usize,
         limit: Option<usize>,
+        elapsed_compute: Time,
     ) -> anyhow::Result<SendableRecordBatchStream> {
         self.scan_partition(
             partition_id,
@@ -135,6 +144,30 @@ where
             predicate,
             batch_size,
             limit,
+            elapsed_compute,
         )
     }
 }
+
+struct ElapsedCompute {
+    time: Time,
+    start: Option<Instant>,
+}
+
+impl ElapsedCompute {
+    fn new(time: Time) -> Self {
+        Self { time, start: None }
+    }
+
+    fn start(&mut self) {
+        self.start.get_or_insert_with(Instant::now);
+    }
+}
+
+impl Drop for ElapsedCompute {
+    fn drop(&mut self) {
+        if let Some(start) = &self.start {
+            self.time.add_elapsed(*start)
+        }
+    }
+}

crates/storage-query-datafusion/src/remote_query_scanner_manager.rs

@@ -18,6 +18,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::execution::SendableRecordBatchStream;
 
 use datafusion::physical_plan::PhysicalExpr;
+use datafusion::physical_plan::metrics::Time;
 use restate_core::Metadata;
 use restate_core::partitions::PartitionRouting;
 use restate_types::NodeId;
@@ -184,6 +185,7 @@ impl ScanPartition for RemotePartitionsScanner {
         predicate: Option<Arc<dyn PhysicalExpr>>,
         batch_size: usize,
         limit: Option<usize>,
+        elapsed_compute: Time,
     ) -> anyhow::Result<SendableRecordBatchStream> {
         match self.manager.get_partition_target_node(partition_id)? {
             PartitionLocation::Local => {
@@ -197,6 +199,7 @@ impl ScanPartition for RemotePartitionsScanner {
                     predicate,
                     batch_size,
                     limit,
+                    elapsed_compute,
                 )?)
             }
             PartitionLocation::Remote { node_id } => Ok(remote_scan_as_datafusion_stream(

crates/storage-query-datafusion/src/scanner_task.rs

@@ -16,6 +16,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr;
 use datafusion::physical_plan::PhysicalExpr;
+use datafusion::physical_plan::metrics::Time;
 use tokio::sync::mpsc;
 use tokio_stream::StreamExt as TokioStreamExt;
 use tracing::{debug, warn};
@@ -92,6 +93,7 @@ impl ScannerTask {
             request
                 .limit
                 .map(|limit| usize::try_from(limit).expect("limit to fit in a usize")),
+            Time::new(),
         )?;
 
         let (tx, rx) = mpsc::unbounded_channel();