feat: add support for explain analyze (#3484)

This adds runtime execution metrics to all of our exec nodes. These
metrics can be accessed by calling plan.analyze_plan().

Author: wkalt

python/python/lance/dataset.py

@@ -3412,6 +3412,21 @@ def explain_plan(self, verbose=False) -> str:
 
         return self._scanner.explain_plan(verbose=verbose)
 
+    def analyze_plan(self) -> str:
+        """Execute the plan for this scanner and display with runtime metrics.
+
+        Parameters
+        ----------
+        verbose : bool, default False
+            Use a verbose output format.
+
+        Returns
+        -------
+        plan : str
+        """
+
+        return self._scanner.analyze_plan()
+
 
 class DatasetOptimizer:
     def __init__(self, dataset: LanceDataset):

python/python/lance/lance/__init__.pyi

@@ -322,6 +322,7 @@ class _Scanner:
     @property
     def schema(self) -> pa.Schema: ...
     def explain_plan(self, verbose: bool) -> str: ...
+    def analyze_plan(self) -> str: ...
     def count_rows(self) -> int: ...
     def to_pyarrow(self) -> pa.RecordBatchReader: ...
 

python/python/tests/test_dataset.py

@@ -788,6 +788,41 @@ def test_select_none(tmp_path: Path):
     ).explain_plan(True)
 
 
+def test_analyze_filtered_scan(tmp_path: Path):
+    table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
+    base_dir = tmp_path / "test"
+    ds = lance.write_dataset(table, base_dir)
+    plan = ds.scanner(columns=[], filter="a < 50", with_row_id=True).analyze_plan()
+    print(plan)
+    assert re.search(r"^\s*LanceScan:.*output_rows=100.*$", plan, re.MULTILINE)
+    assert re.search(r"^\s*FilterExec:.*output_rows=50.*$", plan, re.MULTILINE)
+
+
+def test_analyze_index_scan(tmp_path: Path):
+    table = pa.table({"filter": range(100)})
+    dataset = lance.write_dataset(table, tmp_path)
+    dataset.create_scalar_index("filter", "BTREE")
+    plan = dataset.scanner(filter="filter = 10").analyze_plan()
+    assert "MaterializeIndex: query=filter = 10, metrics=[output_rows=1" in plan
+
+
+def test_analyze_vector_search(tmp_path: Path):
+    table = pa.Table.from_pydict(
+        {
+            "id": [i for i in range(10)],
+            "vector": pa.array(
+                [[1.0, 1.0] for _ in range(10)], pa.list_(pa.float32(), 2)
+            ),
+        }
+    )
+    dataset = lance.write_dataset(table, tmp_path / "dataset", mode="create")
+    dataset.delete("id = 0")
+    plan = dataset.scanner(
+        nearest={"column": "vector", "k": 10, "q": [1.0, 1.0]}
+    ).analyze_plan()
+    assert "KNNVectorDistance: metric=l2, metrics=[output_rows=10" in plan
+
+
 def test_get_fragments(tmp_path: Path):
     table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
     base_dir = tmp_path / "test"

python/src/scanner.rs

@@ -68,6 +68,18 @@ impl Scanner {
         Ok(res)
     }
 
+    #[pyo3(signature = (*))]
+    fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult<String> {
+        let scanner = self_.scanner.clone();
+        let res = RT
+            .spawn(Some(self_.py()), async move {
+                scanner.analyze_plan().await
+            })?
+            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+
+        Ok(res)
+    }
+
     fn count_rows(self_: PyRef<'_, Self>) -> PyResult<u64> {
         let scanner = self_.scanner.clone();
         RT.spawn(Some(self_.py()), async move { scanner.count_rows().await })?

rust/lance/src/dataset/scanner.rs

@@ -13,10 +13,12 @@ use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaR
 use arrow_select::concat::concat_batches;
 use async_recursion::async_recursion;
 use datafusion::common::SchemaExt;
+use datafusion::execution::TaskContext;
 use datafusion::functions_aggregate;
 use datafusion::functions_aggregate::count::count_udaf;
 use datafusion::logical_expr::Expr;
 use datafusion::physical_expr::PhysicalSortExpr;
+use datafusion::physical_plan::analyze::AnalyzeExec;
 use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::expressions;
@@ -2356,6 +2358,26 @@ impl Scanner {
         ))
     }
 
+    #[instrument(level = "info", skip(self))]
+    pub async fn analyze_plan(&self) -> Result<String> {
+        let plan = self.create_plan().await?;
+        let schema = plan.schema();
+        let analyze = Arc::new(AnalyzeExec::new(true, true, plan, schema));
+        let ctx = Arc::new(TaskContext::default());
+        let mut stream = analyze.execute(0, ctx).map_err(|err| {
+            Error::io(
+                format!("Failed to execute analyze plan: {}", err),
+                location!(),
+            )
+        })?;
+
+        // fully execute the plan
+        while (stream.next().await).is_some() {}
+
+        let display = DisplayableExecutionPlan::with_metrics(analyze.as_ref());
+        Ok(format!("{}", display.indent(true)))
+    }
+
     #[instrument(level = "info", skip(self))]
     pub async fn explain_plan(&self, verbose: bool) -> Result<String> {
         let plan = self.create_plan().await?;

rust/lance/src/io/exec/fts.rs

@@ -10,7 +10,7 @@ use datafusion::common::Statistics;
 use datafusion::error::{DataFusionError, Result as DataFusionResult};
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
-use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
 use futures::stream::{self};
@@ -24,7 +24,9 @@ use tracing::instrument;
 use crate::index::prefilter::DatasetPreFilter;
 use crate::{index::DatasetIndexInternalExt, Dataset};
 
-use super::utils::{FilteredRowIdsToPrefilter, SelectionVectorToPrefilter};
+use super::utils::{
+    FilteredRowIdsToPrefilter, InstrumentedRecordBatchStreamAdapter, SelectionVectorToPrefilter,
+};
 use super::PreFilterSource;
 
 /// An execution node that performs full text search
@@ -41,6 +43,8 @@ pub struct FtsExec {
     /// Prefiltering input
     prefilter_source: PreFilterSource,
     properties: PlanProperties,
+
+    metrics: ExecutionPlanMetricsSet,
 }
 
 impl DisplayAs for FtsExec {
@@ -72,6 +76,7 @@ impl FtsExec {
             query,
             prefilter_source,
             properties,
+            metrics: ExecutionPlanMetricsSet::new(),
         }
     }
 }
@@ -108,6 +113,7 @@ impl ExecutionPlan for FtsExec {
                 query: self.query.clone(),
                 prefilter_source: PreFilterSource::None,
                 properties: self.properties.clone(),
+                metrics: ExecutionPlanMetricsSet::new(),
             },
             1 => {
                 let src = children.pop().unwrap();
@@ -130,6 +136,7 @@ impl ExecutionPlan for FtsExec {
                     query: self.query.clone(),
                     prefilter_source,
                     properties: self.properties.clone(),
+                    metrics: ExecutionPlanMetricsSet::new(),
                 }
             }
             _ => {
@@ -150,6 +157,7 @@ impl ExecutionPlan for FtsExec {
         let query = self.query.clone();
         let ds = self.dataset.clone();
         let prefilter_source = self.prefilter_source.clone();
+        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
 
         let indices = self.indices.clone();
         let stream = stream::iter(indices)
@@ -208,16 +216,21 @@ impl ExecutionPlan for FtsExec {
             })
             .buffered(self.indices.len());
         let schema = self.schema();
-        Ok(
-            Box::pin(RecordBatchStreamAdapter::new(schema, stream.boxed()))
-                as SendableRecordBatchStream,
-        )
+        Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new(
+            schema,
+            stream.boxed(),
+            baseline_metrics,
+        )) as SendableRecordBatchStream)
     }
 
     fn statistics(&self) -> DataFusionResult<datafusion::physical_plan::Statistics> {
         Ok(Statistics::new_unknown(&FTS_SCHEMA))
     }
 
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
     fn properties(&self) -> &PlanProperties {
         &self.properties
     }
@@ -235,6 +248,7 @@ pub struct FlatFtsExec {
     column_inputs: Vec<(String, Vec<Index>, Arc<dyn ExecutionPlan>)>,
     query: FullTextSearchQuery,
     properties: PlanProperties,
+    metrics: ExecutionPlanMetricsSet,
 }
 
 impl DisplayAs for FlatFtsExec {
@@ -264,6 +278,7 @@ impl FlatFtsExec {
             column_inputs,
             query,
             properties,
+            metrics: ExecutionPlanMetricsSet::new(),
         }
     }
 }
@@ -309,6 +324,7 @@ impl ExecutionPlan for FlatFtsExec {
             column_inputs,
             query: self.query.clone(),
             properties: self.properties.clone(),
+            metrics: ExecutionPlanMetricsSet::new(),
         }))
     }
 
@@ -321,6 +337,7 @@ impl ExecutionPlan for FlatFtsExec {
         let query = self.query.clone();
         let ds = self.dataset.clone();
         let column_inputs = self.column_inputs.clone();
+        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
 
         let stream = stream::iter(column_inputs)
             .map(move |(column, indices, input)| {
@@ -353,16 +370,21 @@ impl ExecutionPlan for FlatFtsExec {
             .buffered(self.column_inputs.len())
             .try_flatten();
         let schema = self.schema();
-        Ok(
-            Box::pin(RecordBatchStreamAdapter::new(schema, stream.boxed()))
-                as SendableRecordBatchStream,
-        )
+        Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new(
+            schema,
+            stream.boxed(),
+            baseline_metrics,
+        )) as SendableRecordBatchStream)
     }
 
     fn statistics(&self) -> DataFusionResult<datafusion::physical_plan::Statistics> {
         Ok(Statistics::new_unknown(&FTS_SCHEMA))
     }
 
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
+
     fn properties(&self) -> &PlanProperties {
         &self.properties
     }