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.

Author: timsaucer

Cargo.toml

@@ -90,4 +90,4 @@ crate-type = ["cdylib", "rlib"]
 
 [profile.release]
 lto = true
-codegen-units = 1
+codegen-units = 1

pyproject.toml

@@ -185,8 +185,6 @@ dev = [
   "toml>=0.10.2",
   "pygithub==2.5.0",
   "codespell==2.4.1",
-  "protobuf>=6.33.5",
-  "substrait>=0.27.0",
 ]
 docs = [
   "sphinx>=7.1.2",
@@ -198,4 +196,4 @@ docs = [
   "pickleshare>=0.7.5",
   "sphinx-autoapi>=3.4.0",
   "setuptools>=75.3.0",
-]
+]

python/datafusion/plan.py

@@ -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/substrait.py

@@ -76,17 +76,17 @@ def to_json(self) -> str:
         return self.plan_internal.to_json()
 
     @staticmethod
-    def parse_json(json: str) -> Plan:
-        """
-        Parse a plan from a JSON string.
+    def from_json(json: str) -> Plan:
+        """Parse a plan from a JSON string representation.
 
         Args:
             json: JSON representation of a Substrait plan.
 
         Returns:
-            PyPlan object representing the Substrait plan.
+            Plan object representing the Substrait plan.
         """
-        return Plan(substrait_internal.Plan.parse_json(json))
+        return Plan(substrait_internal.Plan.from_json(json))
+
 
 @deprecated("Use `Plan` instead.")
 class plan(Plan):  # noqa: N801
@@ -210,4 +210,4 @@ def from_substrait_plan(ctx: SessionContext, plan: Plan) -> LogicalPlan:
 
 @deprecated("Use `Consumer` instead.")
 class consumer(Consumer):  # noqa: N801
-    """Use `Consumer` instead."""
+    """Use `Consumer` instead."""

python/tests/test_substrait.py

@@ -19,8 +19,7 @@
 import pytest
 from datafusion import SessionContext
 from datafusion import substrait as ss
-from substrait import plan_pb2
-from google.protobuf import json_format
+
 
 @pytest.fixture
 def ctx():
@@ -77,23 +76,74 @@ def test_substrait_file_serialization(ctx, tmp_path, path_to_str):
     assert str(expected_logical_plan) == str(expected_actual_plan)
 
 
-def test_plan_json_stability_and_validator_compatibility(ctx):
-    sql = "SELECT * FROM (VALUES (1, 4), (2, 5), (3, 6)) AS t(a, b)"
-
-    expected_binary_plan = ss.Serde.serialize_bytes(sql, ctx)
-    actual_plan = ss.Serde.serialize_to_plan(sql, ctx)
-
-    json_str = actual_plan.to_json()
-    reconstructed_plan = ss.Plan.parse_json(json_str)  
-
-    expected_logical_plan = ss.Consumer.from_substrait_plan(ctx, actual_plan)
-    actual_logical_plan = ss.Consumer.from_substrait_plan(ctx, reconstructed_plan)
-
-    assert str(expected_logical_plan) == str(actual_logical_plan)
-
-    # Verify that the JSON can be parsed by the Substrait protobuf library
-    proto_plan = plan_pb2.Plan()
-    json_format.Parse(json_str, proto_plan)
-    actual_binary_plan = proto_plan.SerializeToString()
+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 expected_binary_plan == actual_binary_plan
+    assert round_trip_logical_plan == original_logical_plan

src/sql/logical.rs

@@ -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/substrait.rs

@@ -45,16 +45,14 @@ impl PyPlan {
 
     /// 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(|e| to_datafusion_err(e))?;
+        let json = serde_json::to_string_pretty(&self.plan).map_err(|e| to_datafusion_err(e))?;
         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(|e| to_datafusion_err(e))?;
+        let plan: Plan = serde_json::from_str(json).map_err(|e| to_datafusion_err(e))?;
         Ok(PyPlan { plan })
     }
 }