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)