address some concerns

Author: ShreyeshArangath

python/datafusion/plan.py

@@ -157,7 +157,13 @@ def to_proto(self) -> bytes:
         return self._raw_plan.to_proto()
 
     def metrics(self) -> MetricsSet | None:
-        """Return metrics for this plan node after execution, or None if unavailable."""
+        """Return metrics for this plan node, or None if this node has no MetricsSet.
+
+        Some operators (e.g. DataSourceExec) eagerly initialize a MetricsSet
+        when the plan is created, so this may return a set even before
+        execution.  Metric *values* (such as ``output_rows``) are only
+        meaningful after the DataFrame has been executed.
+        """
         raw = self._raw_plan.metrics()
         if raw is None:
             return None

python/tests/test_plans.py

@@ -95,14 +95,21 @@ def test_metric_properties() -> None:
     assert found_any_metric, "Expected at least one metric after execution"
 
 
-def test_no_metrics_before_execution() -> None:
+def test_no_meaningful_metrics_before_execution() -> None:
     ctx = SessionContext()
     ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)")
     df = ctx.sql("SELECT * FROM t")
     plan = df.execution_plan()
-    ms = plan.metrics()
-    # Before execution, the root plan node has no MetricsSet
-    assert ms is None
+
+    # Some plan nodes (e.g. DataSourceExec) eagerly initialize a MetricsSet,
+    # so metrics() may return a set even before execution.  However, no rows
+    # should have been processed yet.
+    output_rows_before = [
+        ms.output_rows
+        for _, ms in plan.collect_metrics()
+        if ms.output_rows is not None and ms.output_rows > 0
+    ]
+    assert len(output_rows_before) == 0, "Expected no output rows before execution"
 
 
 def test_collect_partitioned_metrics() -> None:
@@ -174,7 +181,7 @@ def test_value_as_datetime() -> None:
                 assert metric.value_as_datetime is None
 
 
-def test_collect_twice_reuses_plan() -> None:
+def test_collect_twice_has_metrics() -> None:
     ctx = SessionContext()
     ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
     df = ctx.sql("SELECT * FROM t WHERE column1 > 1")

src/dataframe.rs

@@ -387,21 +387,11 @@ impl PyDataFrame {
         &self,
         py: Python,
     ) -> PyDataFusionResult<(Arc<dyn DFExecutionPlan>, Arc<TaskContext>)> {
-        let plan = {
-            let cached = self.last_plan.lock();
-            cached.as_ref().map(Arc::clone)
-        };
-        let plan = match plan {
-            Some(p) => p,
-            None => {
-                let df = self.df.as_ref().clone();
-                let new_plan = wait_for_future(py, df.create_physical_plan())??;
-                *self.last_plan.lock() = Some(Arc::clone(&new_plan));
-                new_plan
-            }
-        };
+        let df = self.df.as_ref().clone();
+        let new_plan = wait_for_future(py, df.create_physical_plan())??;
+        *self.last_plan.lock() = Some(Arc::clone(&new_plan));
         let task_ctx = Arc::new(self.df.as_ref().task_ctx());
-        Ok((plan, task_ctx))
+        Ok((new_plan, task_ctx))
     }
 
     async fn collect_column_inner(&self, column: &str) -> Result<ArrayRef, DataFusionError> {