feat: add support for generating JSON formatted substrait plan by Prathamesh9284 · Pull Request #1376 · apache/datafusion-python

Prathamesh9284 · 2026-02-10T09:24:36Z

Which issue does this PR close?

Closes #508

Rationale for this change

Other Substrait producers (Isthmus, DuckDB) support generating JSON-formatted Substrait plans, which makes it much easier to inspect and compare plans across engines. The substrait crate already provides serde support via the pbjson feature, so DataFusion Python can leverage this to expose JSON plan generation and parsing with minimal effort.

What changes are included in this PR?

python/datafusion/substrait.py
- Added to_json() instance method and parse_json() static method on the Plan class for converting plans to JSON strings and parse JSON strings to plans.
src/substrait.rs
- Implemented the Rust backend for the above Python methods, including error mapping.
- Added the following functions:
  - fn to_json(&self) -> PyDataFusionResult<String>
  - #[staticmethod] fn from_json(json: &str) -> PyDataFusionResult<PyPlan>
Cargo.toml
- Added serde_json as a dependency.
pyproject.toml
- Added protobuf and substrait as dev dependencies for test support.
python/tests/test_substrait.py
- Added test for JSON round-trip stability and protobuf compatibility validation.

Are there any user-facing changes?

Yes — two new public API methods:

Plan.to_json() -> str — convert a Substrait plan to a JSON string
Plan.from_json(json: str) -> Plan — create a Substrait plan from a JSON string

timsaucer

I have a few questions/recommendations. My biggest concern is I don't see evidence that this is generating json strings that are actually compatible with the other libraries and described in the original issue. It's entirely possible it's working exactly as needed, but I think that's something we can show in the tests.

timsaucer · 2026-02-11T16:38:23Z

python/datafusion/substrait.py

+    def to_json(self) -> str:
+        """Serialize the plan to a JSON string.
+
+        Returns a JSON string following the Protobuf JSON Mapping.


I find the statement "Returns a JSON string following the Protobuf JSON Mapping" to be hard to understand. What exactly does this mean? Is this LLM generated?

I agree the wording isn’t very clear. I’ll rephrase it to make it clearer that this just returns the Plan in its standard JSON form.

timsaucer · 2026-02-11T16:42:46Z

python/datafusion/substrait.py


+    @staticmethod
+    def serialize_json(sql: str, ctx: SessionContext, path: str | pathlib.Path) -> None:
+        """Serialize a SQL query to a JSON-formatted Substrait plan file.


Is the output a json file? I'm trying to understand the "serialize" in this context.

Yes, in the current implementation it writes a JSON-formatted Substrait plan to a file. I used “serialize” in the same sense as the existing binary serialize, but I agree the naming may not be very clear here.

timsaucer · 2026-02-11T16:44:38Z

python/tests/test_substrait.py

+    # Convert plan to JSON
+    json_str = plan.to_json()
+    assert isinstance(json_str, str)
+    assert "relations" in json_str  # basic sanity check


If these conversions to json are stable as the issue suggests, then I would think we could compare the entire output, right? What do the other libraries produce?

Hi @timsaucer,

I was able to generate a JSON-formatted Substrait plan using my DataFusion implementation. Comparing it with Isthmus:

Both include relations, project expressions, virtualTable values, and correct column mappings.

Minor differences exist: DataFusion uses i64 instead of i32, column names differ, and optional metadata (extensions, typeAliases) are not included — these do not affect compatibility.

DuckDB currently does not provide a working JSON inspection method due to extension availability issues, so direct JSON comparison via DuckDB isn’t possible.

Overall, the DataFusion output is compatible and can be used for reproducibility and validation.

timsaucer · 2026-02-11T16:45:41Z

python/datafusion/substrait.py

+    @staticmethod
+    def serialize_json(sql: str, ctx: SessionContext, path: str | pathlib.Path) -> None:


Is the preferred method here to write to a file or to just produce a string output? Maybe the author of the issue has an opinion on best use case.

I was originally trying to keep it consistent with Serde.serialize, which writes the binary plan to a file. But I’m also fine with just returning a JSON string from Plan.to_json() and letting users handle writing to a file themselves. Do you have a preference for which approach fits better with the project?

timsaucer · 2026-02-11T16:46:23Z

python/datafusion/substrait.py

+    def deserialize_json(path: str | pathlib.Path) -> Plan:
+        """Deserialize a Substrait plan from a JSON file.


Similarly, not sure if this needs to be a path to a file that we read or expect to pass in a string. Since it's python it's fairly easy to handle either/both

Following up on my earlier comment if we keep JSON handling on Plan, would it make sense for from_json() to accept either a JSON string or a file path? That might simplify things instead of having separate helpers on Serde. Curious what you think.

timsaucer · 2026-02-11T16:51:43Z

src/errors.rs

    PythonError(PyErr),
    EncodeError(EncodeError),
+    DecodeError(DecodeError),
+    SerdeJsonError(SerdeJsonError),


Recommend removing SerdeJsonError and mapping to EncodeError or DecodeError. Alternatively putting into DataFusion execution error.

Prathamesh9284 · 2026-02-12T17:58:36Z

Hi @timsaucer

I have added pytest to validate cross-library compatibility. On test_substrait.py ines 94–99 parse the JSON using the official Substrait protobuf library (plan_pb2) via json_format.Parse, which would fail if the JSON didn’t conform to the Substrait protobuf schema. It then re-serializes to binary and asserts byte-for-byte equality with the original DataFusion output. This ensures any tool that supports protobuf-JSON can consume it correctly.

… plan. Switch unit test to focus on json parsing and not byte serialization.

timsaucer · 2026-02-13T12:50:05Z

Thank you for the PR. There was an issue where the internal function name differed from the wrapper, so unit tests were failing. I pushed an update correcting this and making a change to the unit test to focus on the json representation of the logical plan rather than serialization and deserialization.

Prathamesh9284 · 2026-02-13T15:33:39Z

Thank you for fixing the internal function name mismatch and adjusting the unit test. I did notice that one of the CI builds failed. The failure appears to be related to a Clippy lint (redundant_closure). I’ve pushed a small fix replacing the redundant closure with the function directly (map_err(to_datafusion_err)), which should resolve the lint issue across environments.

timsaucer requested changes Feb 11, 2026

View reviewed changes

Prathamesh9284 force-pushed the feat/substrait-json-plan branch 2 times, most recently from 479bbd9 to e9df13e Compare February 12, 2026 14:37

chore: add new dependencies to pyproject.toml

babd63a

Prathamesh9284 force-pushed the feat/substrait-json-plan branch from e9df13e to babd63a Compare February 12, 2026 14:48

Prathamesh9284 added 2 commits February 12, 2026 15:27

chore: rename from_json to parse_json

d2572d5

fix: missin import

63da3ae

Prathamesh9284 requested a review from timsaucer February 12, 2026 17:56

Fix call to internal function. Drive by update to dquality on logical…

a7de15e

… plan. Switch unit test to focus on json parsing and not byte serialization.

fix: resolve clippy redundant closure lint in substrait.rs

6043e97

		@staticmethod
		def serialize_json(sql: str, ctx: SessionContext, path: str \| pathlib.Path) -> None:

		def deserialize_json(path: str \| pathlib.Path) -> Plan:
		"""Deserialize a Substrait plan from a JSON file.

Conversation

Prathamesh9284 commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

timsaucer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Prathamesh9284 commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

timsaucer commented Feb 13, 2026

Uh oh!

Prathamesh9284 commented Feb 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Prathamesh9284 commented Feb 10, 2026 •

edited

Loading

Prathamesh9284 commented Feb 12, 2026 •

edited

Loading