[FEAT] Enable to_json_string for physical plan

daft/daft/__init__.pyi

@@ -1623,6 +1623,7 @@ class PhysicalPlanScheduler:
     def num_partitions(self) -> int: ...
     def repr_ascii(self, simple: bool) -> str: ...
     def repr_mermaid(self, options: MermaidOptions) -> str: ...
+    def to_json_string(self) -> str: ...
     def to_partition_tasks(self, psets: dict[str, list[PartitionT]]) -> physical_plan.InProgressPhysicalPlan: ...
     def run(self, psets: dict[str, list[PartitionT]]) -> Iterator[PyMicroPartition]: ...
 

daft/logical/builder.py

@@ -62,6 +62,9 @@ def to_physical_plan_scheduler(self, daft_execution_config: PyDaftExecutionConfi
         used to generate executable tasks for the physical plan.
 
         This should be called after triggering optimization with self.optimize().
+
+        **Warning**: This function is not part of the stable API and may change
+        without notice. It is intended for internal or experimental use only.
         """
         from daft.plan_scheduler.physical_plan_scheduler import PhysicalPlanScheduler
 

daft/plan_scheduler/physical_plan_scheduler.py

@@ -53,6 +53,9 @@
     def __repr__(self) -> str:
         return self._scheduler.repr_ascii(simple=False)
 
+    def to_json_string(self) -> str:
+        return self._scheduler.to_json_string()
+
     def to_partition_tasks(
         self, psets: dict[str, list[PartitionT]], results_buffer_size: int | None
     ) -> physical_plan.MaterializedPhysicalPlan:

src/daft-scheduler/Cargo.toml

@@ -11,6 +11,7 @@ daft-plan = {path = "../daft-plan", default-features = false}
 daft-scan = {path = "../daft-scan", default-features = false}
 pyo3 = {workspace = true, optional = true}
 serde = {workspace = true, features = ["rc"]}
+serde_json = {workspace = true}
 
 [dev-dependencies]
 rstest = {workspace = true}

src/daft-scheduler/src/scheduler.rs

@@ -80,6 +80,10 @@
         Ok(self.plan().repr_mermaid(options))
     }
 
+    pub fn to_json_string(&self) -> PyResult<String> {
+        serde_json::to_string(&self.plan())
+            .map_err(|e| pyo3::PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
+    }
     /// Converts the contained physical plan into an iterator of executable partition tasks.
     pub fn to_partition_tasks(
         &self,