feat: add a module to read all Substrait plan formats

.github/workflows/release.yml

@@ -25,6 +25,14 @@ jobs:
       - name: Install build dependencies
         run: |
           python -m pip install build --user
+      - name: Build Substrait planloader library
+        run: |
+          cd ${{ github.workspace }}/third_party/substrait-cpp
+          make release
+      - name: Install Substrait planloader library
+        run: |
+          cd ${{ github.workspace }}/third_party/substrait-cpp/build-Release/export/planloader
+          make install
       - name: Build package
         run: |
           python -m build

.github/workflows/test.yml

@@ -26,6 +26,14 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python }}
+      - name: Build Substrait planloader library
+        run: |
+          cd ${{ github.workspace }}/third_party/substrait-cpp
+          make release
+      - name: Install Substrait planloader library
+        run: |
+          cd ${{ github.workspace }}/third_party/substrait-cpp/build-Release/export/planloader
+          make install
       - name: Install package and test dependencies
         run: |
           python -m pip install --upgrade pip

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "third_party/substrait"]
 	path = third_party/substrait
 	url = https://github.com/substrait-io/substrait
+[submodule "third_party/substrait-cpp"]
+	path = third_party/substrait-cpp
+	url = git@github.com:substrait-io/substrait-cpp.git

CONTRIBUTING.md

@@ -50,6 +50,16 @@ Generate the protobuf files manually. Requires protobuf `v3.20.1`.
 
 
 # Build
+
+## Build and install the textplan loader dynamic library
+```commandline
+pushd third_party/substrait-cpp
+make release
+cd build-Release/export/planloader
+make install
+popd
+```
+
 ## Python package
 Editable installation.
 ```

README.md

@@ -27,7 +27,7 @@ This project is not an execution engine for Substrait Plans.
 This is an experimental package that is still under development.
 
 # Example
-At the moment, this project contains only generated Python classes for the Substrait protobuf messages. Let's use an existing Substrait producer, [Ibis](https://ibis-project.org), to provide an example using Python Substrait as the consumer.
+At the moment, this project contains generated Python classes for the Substrait protobuf messages and a library for loading and saving them in various formats. Let's use an existing Substrait producer, [Ibis](https://ibis-project.org), to provide an example using Python Substrait as the consumer.
 ## Produce a Substrait Plan with Ibis
 ```
 In [1]: import ibis

pyproject.toml

@@ -17,6 +17,7 @@ test = ["pytest >= 7.0.0"]
 
 [tool.pytest.ini_options]
 pythonpath = "src"
+addopts = "--ignore=third_party"
 
 [build-system]
 requires = ["setuptools>=61.0.0", "setuptools_scm[toml]>=6.2.0"]

src/substrait/planloader/planloader.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Routines for loading and saving Substrait plans."""
+import ctypes
+import ctypes.util as ctutil
+import enum
+import substrait.gen.proto.plan_pb2 as plan_pb2
+import sys
+
+
+class PlanFileFormat(enum.Enum):
+    BINARY = ctypes.c_int32(0)
+    JSON = ctypes.c_int32(1)
+    PROTOTEXT = ctypes.c_int32(2)
+    TEXT = ctypes.c_int32(3)
+
+
+class PlanFileException(Exception):
+    pass
+
+
+class SerializedPlan(ctypes.Structure):
+    pass
+
+
+SerializedPlan._fields_ = [
+    ("buffer", ctypes.POINTER(ctypes.c_byte)),
+    ("size", ctypes.c_int32),
+    ("errorMessage", ctypes.c_char_p),
+]
+
+
+# Load the C++ library
+planloader_path = ctutil.find_library("planloader")
+planloader_lib = ctypes.CDLL(planloader_path)
+if planloader_lib is None:
+    print('Failed to find planloader library')
+    sys.exit(1)
+
+# Declare the function signatures for the external functions.
+external_load_substrait_plan = planloader_lib.load_substrait_plan
+external_load_substrait_plan.argtypes = [ctypes.c_char_p]
+external_load_substrait_plan.restype = ctypes.POINTER(SerializedPlan)
+
+external_free_substrait_plan = planloader_lib.free_substrait_plan
+external_free_substrait_plan.argtypes = [ctypes.POINTER(SerializedPlan)]
+external_free_substrait_plan.restype = None
+
+external_save_substrait_plan = planloader_lib.save_substrait_plan
+external_save_substrait_plan.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_char_p, ctypes.c_int32]
+external_save_substrait_plan.restype = ctypes.c_char_p
+
+
+def load_substrait_plan(filename: str) -> plan_pb2.Plan:
+    """
+    Loads a Substrait plan (in any format) from disk.
+
+    Returns:
+        A Plan protobuf object if successful.
+    Raises:
+        PlanFileException if an except occurs while converting or reading from disk.
+    """
+    result = external_load_substrait_plan(filename.encode('UTF-8'))
+    try:
+      if result.contents.errorMessage:
+          raise PlanFileException(result.contents.errorMessage)
+      data = ctypes.string_at(result.contents.buffer, result.contents.size)
+      plan = plan_pb2.Plan()
+      plan.ParseFromString(data)
+    finally:
+      external_free_substrait_plan(result)
+    return plan
+
+
+def save_substrait_plan(plan: plan_pb2.Plan, filename: str, file_format: PlanFileFormat):
+    """
+    Saves the given plan to disk in the specified file format.
+
+    Raises:
+        PlanFileException if an except occurs while converting or writing to disk.
+    """
+    data = plan.SerializeToString()
+    err = external_save_substrait_plan(data, len(data), filename.encode('UTF-8'), file_format.value)
+    if err:
+        raise PlanFileException(err)

tests/test_planloader.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+from substrait.planloader import planloader
+
+
+def test_main():
+    testplan = planloader.load_substrait_plan('tests/tpch-plan01.json')
+    planloader.save_substrait_plan(testplan, 'myoutfile.splan', planloader.PlanFileFormat.TEXT.value)