build: switch to substrait-antlr

.github/workflows/version-checks.yml

@@ -0,0 +1,24 @@
+name: Substrait Package Version Checks
+
+on:
+  pull_request:
+
+permissions:
+  contents: read
+
+jobs:
+  substrait-package-version-checks:
+    name: Verify substrait-* Package Versions
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+        with:
+          submodules: recursive
+
+      - name: Install uv with python
+        uses: astral-sh/setup-uv@v7
+
+      - name: Verify substrait-* versions
+        run: ./check_substrait_package_versions.sh
+

CONTRIBUTING.md

@@ -45,14 +45,6 @@ Copy the core function extensions into substrait-python
 pixi run copy-extensions
 ```
 
-## Antlr grammar
-
-Substrait uses antlr grammar to derive output types of extension functions. Make sure java is installed and ANTLR_JAR environment variable is set. Take a look at .devcontainer/Dockerfile for example setup.
-
-```
-pixi run antlr
-```
-
 ## Extensions stubs
 
 Substrait uses jsonschema to describe the data model for extension files.

check_substrait_package_versions.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+set -euo pipefail
+
+# Check that all substrait-* packages have the same version
+uv sync
+
+protobuf_version=$(uv pip show substrait-protobuf | grep '^Version:' | awk '{print $2}')
+antlr_version=$(uv pip show substrait-antlr | grep '^Version:' | awk '{print $2}')
+
+if [ "$protobuf_version" != "$antlr_version" ]; then
+    echo "error: substrait-protobuf ($protobuf_version) and substrait-antlr ($antlr_version) versions do not match"
+    exit 1
+fi

pyproject.toml

@@ -15,11 +15,11 @@ dynamic = ["version"]
 write_to = "src/substrait/_version.py"
 
 [project.optional-dependencies]
-extensions = ["antlr4-python3-runtime", "pyyaml"]
+extensions = ["antlr4-python3-runtime", "substrait-antlr==0.79.0", "pyyaml"]
 sql = ["sqloxide; python_version < '3.14'", "deepdiff"]
 
 [dependency-groups]
-dev = ["pytest >= 7.0.0", "antlr4-python3-runtime", "pyyaml", "sqloxide; python_version < '3.14'", "deepdiff", "duckdb<=1.2.2; python_version < '3.14'", "datafusion"]
+dev = ["pytest >= 7.0.0", "antlr4-python3-runtime", "substrait-antlr==0.79.0", "pyyaml", "sqloxide; python_version < '3.14'", "deepdiff", "duckdb<=1.2.2; python_version < '3.14'", "datafusion"]
 
 [tool.pytest.ini_options]
 pythonpath = "src"
@@ -56,19 +56,12 @@ update-substrait = [ { task = "update-submodule" }, { task = "codegen" }]
 
 update-submodule = "./update_submodule.sh"
 
-codegen = [{ task = "antlr" }, { task = "copy-extensions" }, { task = "codegen-extensions" }]
+codegen = [{ task = "copy-extensions" }, { task = "codegen-extensions" }]
 
 check-codegen = "./check_codegen.sh"
 
 copy-extensions = "./copy_extension_yamls.sh"
 
-antlr = """
-	cd third_party/substrait/grammar \
-		&& antlr4 -o ../../../src/substrait/gen/antlr -Dlanguage=Python3 SubstraitType.g4 \
-		&& rm ../../../src/substrait/gen/antlr/*.tokens \
-		&& rm ../../../src/substrait/gen/antlr/*.interp
-"""
-
 codegen-extensions = """
     datamodel-codegen \
 		--input-file-type jsonschema \
@@ -81,7 +74,5 @@ codegen-extensions = """
 """
 
 [tool.pixi.dependencies]
-antlr = ">=4.13.2,<5"
-buf = ">=1.62.1,<2"
 datamodel-code-generator = ">=0.52.0,<0.53"
 ruff = ">=0.14.10,<0.15"

src/substrait/derivation_expression.py

@@ -2,9 +2,8 @@
 
 from antlr4 import CommonTokenStream, InputStream
 from substrait.type_pb2 import NamedStruct, Type
-
-from substrait.gen.antlr.SubstraitTypeLexer import SubstraitTypeLexer
-from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
+from substrait_antlr.substrait_type.SubstraitTypeLexer import SubstraitTypeLexer
+from substrait_antlr.substrait_type.SubstraitTypeParser import SubstraitTypeParser
 
 
 def _evaluate(x, values: dict):

src/substrait/extension_registry/signature_checker_helpers.py

@@ -3,9 +3,9 @@
 from typing import Dict
 
 from substrait.type_pb2 import Type
+from substrait_antlr.substrait_type.SubstraitTypeParser import SubstraitTypeParser
 
 from substrait.derivation_expression import _evaluate
-from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
 
 from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError