feat: add support for generating JSON formatted substrait plan (#1376)

* chore: add new dependencies to pyproject.toml

* chore: rename from_json to parse_json

* fix: missin import

* Fix call to internal function. Drive by update to dquality on logical plan. Switch unit test to focus on json parsing and not byte serialization.

* fix: resolve clippy redundant closure lint in substrait.rs

---------

Co-authored-by: Tim Saucer <timsaucer@gmail.com>

Author: Prathamesh9284

Cargo.lock

@@ -62,6 +62,7 @@ datafusion-substrait = { version = "52", optional = true }
 datafusion-proto = { version = "52" }
 datafusion-ffi = { version = "52" }
 prost = "0.14.1" # keep in line with `datafusion-substrait`
+serde_json = "1"
 uuid = { version = "1.18", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = [
   "local_dynamic_tls",

Cargo.toml

@@ -98,6 +98,12 @@ def to_proto(self) -> bytes:
         """
         return self._raw_plan.to_proto()
 
+    def __eq__(self, other: LogicalPlan) -> bool:
+        """Test equality."""
+        if not isinstance(other, LogicalPlan):
+            return False
+        return self._raw_plan.__eq__(other._raw_plan)
+
 
 class ExecutionPlan:
     """Represent nodes in the DataFusion Physical Plan."""

python/datafusion/plan.py

@@ -67,6 +67,26 @@ def encode(self) -> bytes:
         """
         return self.plan_internal.encode()
 
+    def to_json(self) -> str:
+        """Get the JSON representation of the Substrait plan.
+
+        Returns:
+            A JSON representation of the Substrait plan.
+        """
+        return self.plan_internal.to_json()
+
+    @staticmethod
+    def from_json(json: str) -> Plan:
+        """Parse a plan from a JSON string representation.
+
+        Args:
+            json: JSON representation of a Substrait plan.
+
+        Returns:
+            Plan object representing the Substrait plan.
+        """
+        return Plan(substrait_internal.Plan.from_json(json))
+
 
 @deprecated("Use `Plan` instead.")
 class plan(Plan):  # noqa: N801

python/datafusion/substrait.py

@@ -74,3 +74,76 @@ def test_substrait_file_serialization(ctx, tmp_path, path_to_str):
     expected_actual_plan = ss.Consumer.from_substrait_plan(ctx, actual_plan)
 
     assert str(expected_logical_plan) == str(expected_actual_plan)
+
+
+def test_json_processing_round_trip(ctx: SessionContext):
+    ctx.register_record_batches("t", [[pa.record_batch({"a": [1]})]])
+    original_logical_plan = ctx.sql("SELECT * FROM t").logical_plan()
+
+    substrait_plan = ss.Producer.to_substrait_plan(original_logical_plan, ctx)
+    json_plan = substrait_plan.to_json()
+
+    expected = """\
+  "relations": [
+    {
+      "root": {
+        "input": {
+          "project": {
+            "common": {
+              "emit": {
+                "outputMapping": [
+                  1
+                ]
+              }
+            },
+            "input": {
+              "read": {
+                "baseSchema": {
+                  "names": [
+                    "a"
+                  ],
+                  "struct": {
+                    "types": [
+                      {
+                        "i64": {
+                          "nullability": "NULLABILITY_NULLABLE"
+                        }
+                      }
+                    ],
+                    "nullability": "NULLABILITY_REQUIRED"
+                  }
+                },
+                "namedTable": {
+                  "names": [
+                    "t"
+                  ]
+                }
+              }
+            },
+            "expressions": [
+              {
+                "selection": {
+                  "directReference": {
+                    "structField": {}
+                  },
+                  "rootReference": {}
+                }
+              }
+            ]
+          }
+        },
+        "names": [
+          "a"
+        ]
+      }
+    }
+  ]"""
+
+    assert expected in json_plan
+
+    round_trip_substrait_plan = ss.Plan.from_json(json_plan)
+    round_trip_logical_plan = ss.Consumer.from_substrait_plan(
+        ctx, round_trip_substrait_plan
+    )
+
+    assert round_trip_logical_plan == original_logical_plan

python/tests/test_substrait.py

@@ -66,8 +66,8 @@ use crate::expr::unnest::PyUnnest;
 use crate::expr::values::PyValues;
 use crate::expr::window::PyWindowExpr;
 
-#[pyclass(frozen, name = "LogicalPlan", module = "datafusion", subclass)]
-#[derive(Debug, Clone)]
+#[pyclass(frozen, name = "LogicalPlan", module = "datafusion", subclass, eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct PyLogicalPlan {
     pub(crate) plan: Arc<LogicalPlan>,
 }

src/sql/logical.rs

@@ -23,7 +23,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 
 use crate::context::PySessionContext;
-use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err};
+use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err};
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::wait_for_future;
 
@@ -42,6 +42,19 @@ impl PyPlan {
             .map_err(PyDataFusionError::EncodeError)?;
         Ok(PyBytes::new(py, &proto_bytes).into())
     }
+
+    /// Get the JSON representation of the substrait plan
+    fn to_json(&self) -> PyDataFusionResult<String> {
+        let json = serde_json::to_string_pretty(&self.plan).map_err(to_datafusion_err)?;
+        Ok(json)
+    }
+
+    /// Parse a Substrait Plan from its JSON representation
+    #[staticmethod]
+    fn from_json(json: &str) -> PyDataFusionResult<PyPlan> {
+        let plan: Plan = serde_json::from_str(json).map_err(to_datafusion_err)?;
+        Ok(PyPlan { plan })
+    }
 }
 
 impl From<PyPlan> for Plan {