plan caching

Author: ShreyeshArangath

python/datafusion/plan.py

@@ -19,6 +19,8 @@
 
 from __future__ import annotations
 
+import datetime
+
 from typing import TYPE_CHECKING, Any
 
 import datafusion._internal as df_internal
@@ -290,6 +292,11 @@ def value(self) -> int | None:
         """
         return self._raw.value
 
+    @property
+    def value_as_datetime(self) -> datetime.datetime | None:
+        """The value as a UTC datetime for timestamp metrics, or ``None``."""
+        return self._raw.value_as_datetime()
+
     @property
     def partition(self) -> int | None:
         """The 0-based partition index this metric applies to.

python/tests/test_plans.py

@@ -15,6 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import datetime
+
+import pytest
+
 from datafusion import (
     ExecutionPlan,
     LogicalPlan,
@@ -150,3 +154,38 @@ def test_execute_stream_partitioned_metrics() -> None:
         if ms.output_rows is not None
     ]
     assert 2 in output_rows_values
+
+
+def test_value_as_datetime() -> 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")
+    df.collect()
+    plan = df.execution_plan()
+
+    for _, ms in plan.collect_metrics():
+        for metric in ms.metrics():
+            if metric.name in ("start_timestamp", "end_timestamp"):
+                dt = metric.value_as_datetime
+                assert dt is None or isinstance(dt, datetime.datetime)
+                if dt is not None:
+                    assert dt.tzinfo is not None
+            else:
+                assert metric.value_as_datetime is None
+
+
+def test_collect_twice_reuses_plan() -> 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")
+
+    df.collect()
+    df.collect()
+
+    plan = df.execution_plan()
+    output_rows_values = [
+        ms.output_rows
+        for _, ms in plan.collect_metrics()
+        if ms.output_rows is not None
+    ]
+    assert len(output_rows_values) > 0

src/dataframe.rs

@@ -387,9 +387,19 @@ impl PyDataFrame {
         &self,
         py: Python,
     ) -> PyDataFusionResult<(Arc<dyn DFExecutionPlan>, Arc<TaskContext>)> {
-        let df = self.df.as_ref().clone();
-        let plan = wait_for_future(py, df.create_physical_plan())??;
-        *self.last_plan.lock() = Some(Arc::clone(&plan));
+        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 task_ctx = Arc::new(self.df.as_ref().task_ctx());
         Ok((plan, task_ctx))
     }

src/metrics.rs

@@ -122,6 +122,35 @@ impl PyMetric {
         }
     }
 
+    /// Returns the value as a Python `datetime` for `StartTimestamp` / `EndTimestamp`
+    /// metrics, or `None` for all other metric types.
+    fn value_as_datetime<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
+        match self.metric.value() {
+            MetricValue::StartTimestamp(ts) | MetricValue::EndTimestamp(ts) => {
+                match ts.value() {
+                    Some(dt) => {
+                        let nanos = dt.timestamp_nanos_opt()
+                            .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyOverflowError, _>(
+                                "timestamp out of range"
+                            ))?;
+                        let datetime_mod = py.import("datetime")?;
+                        let datetime_cls = datetime_mod.getattr("datetime")?;
+                        let tz_utc = datetime_mod.getattr("timezone")?.getattr("utc")?;
+                        let secs = nanos / 1_000_000_000;
+                        let micros = (nanos % 1_000_000_000) / 1_000;
+                        let result = datetime_cls.call_method1(
+                            "fromtimestamp",
+                            (secs as f64 + micros as f64 / 1_000_000.0, tz_utc),
+                        )?;
+                        Ok(Some(result))
+                    }
+                    None => Ok(None),
+                }
+            }
+            _ => Ok(None),
+        }
+    }
+
     #[getter]
     fn partition(&self) -> Option<usize> {
         self.metric.partition()