feat: graceful URI -> URN migration (#114)

Introduces a graceful migration from the previous usage of URIs to URNs

- requires extensions to present a URI explicitly and a URN inside of
the yaml file #95
- removes the usage of uri/urn aliases in the `ExtensionRegistry`
- introduce action to ensure examples run #121
- introduce action to ensure generated code up to date #120

Closes #95, #120, #121

Author: benbellick

.devcontainer/devcontainer.json

@@ -14,7 +14,7 @@
 	// "forwardPorts": [],
 
 	// Use 'postCreateCommand' to run commands after the container is created.
-	"postCreateCommand": "uv venv --clear && uv sync --extra test",
+	"postCreateCommand": "uv venv --clear && uv sync --extra test --extra gen_proto",
 
 	// Configure tool-specific properties.
 	"customizations": {

.github/workflows/codegen-check.yml

@@ -0,0 +1,44 @@
+name: Code Generation Check
+
+on:
+  pull_request:
+
+permissions:
+  contents: read
+
+jobs:
+  codegen-check:
+    name: Verify Code Generation
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v5
+        with:
+          submodules: recursive
+
+      - name: Run code generation in devcontainer
+        uses: devcontainers/[email protected]
+        with:
+          runCmd: |
+            # Ensure dependencies are installed
+            uv sync --extra test --extra gen_proto
+            # Run all code generation steps
+            make antlr
+            ./gen_proto.sh
+            make codegen-extensions
+
+      - name: Check for uncommitted changes
+        run: |
+          # Check for diffs, ignoring timestamp lines
+          if ! git diff --quiet --exit-code src/substrait/gen/; then
+            echo "Code generation produced changes. Generated code is out of sync!"
+            echo ""
+            git diff src/substrait/gen/
+            echo ""
+            echo "To fix this, run:"
+            echo "  make antlr"
+            echo "  ./gen_proto.sh"
+            echo "  make codegen-extensions"
+            echo "Then commit the changes."
+            exit 1
+          fi

.github/workflows/example.yml

@@ -0,0 +1,36 @@
+name: Run examples
+
+on:
+  pull_request:
+  push:
+    branches: [ main ]
+
+permissions:
+  contents: read
+
+jobs:
+  example:
+    name: Run ${{ matrix.example }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        example:
+          - builder_example.py
+          - duckdb_example.py
+          - adbc_example.py
+          - pyarrow_example.py
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v5
+        with:
+          submodules: recursive
+      - name: Install uv with python
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: "3.10"
+      - name: Install package dependencies
+        run: |
+          uv sync --frozen --extra extensions
+      - name: Run ${{ matrix.example }}
+        run: |
+          uv run examples/${{ matrix.example }}

Makefile

@@ -9,7 +9,8 @@ codegen-extensions:
 		--input-file-type jsonschema \
 		--input third_party/substrait/text/simple_extensions_schema.yaml \
 		--output src/substrait/gen/json/simple_extensions.py \
-		--output-model-type dataclasses.dataclass
+		--output-model-type dataclasses.dataclass \
+		--disable-timestamp
 
 lint:
 	uvx [email protected] check

examples/adbc_example.py

@@ -45,7 +45,7 @@ def read_adbc_named_table(name: str, conn):
         table = filter(
             table,
             expression=scalar_function(
-                "functions_comparison.yaml",
+                "extension:io.substrait:functions_comparison",
                 "gte",
                 expressions=[column("ints"), literal(3, i64())],
             ),

examples/builder_example.py

@@ -21,17 +21,17 @@
 def basic_example():
     ns = named_struct(
         names=["id", "is_applicable"],
-        struct=struct(types=[i64(nullable=False), boolean()]),
+        struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
     )
 
     table = read_named_table("example_table", ns)
     table = filter(table, expression=column("is_applicable"))
     table = filter(
         table,
         expression=scalar_function(
-            "functions_comparison.yaml",
+            "extension:io.substrait:functions_comparison",
             "lt",
-            expressions=[column("id"), literal(100, i64())],
+            expressions=[column("id"), literal(100, i64(nullable=False))],
         ),
     )
     table = project(table, expressions=[column("id")])
@@ -41,14 +41,15 @@ def basic_example():
 
     """
   extension_uris {
-    extension_uri_anchor: 13
-    uri: "functions_comparison.yaml"
+    extension_uri_anchor: 2
+    uri: "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml"
   }
   extensions {
     extension_function {
-      extension_uri_reference: 13
-      function_anchor: 495
-      name: "lt"
+      extension_uri_reference: 2
+      function_anchor: 124
+      name: "lt:any_any"
+      extension_urn_reference: 2
     }
   }
   relations {
@@ -84,7 +85,7 @@ def basic_example():
                               nullability: NULLABILITY_NULLABLE
                             }
                           }
-                          nullability: NULLABILITY_NULLABLE
+                          nullability: NULLABILITY_REQUIRED
                         }
                       }
                       named_table {
@@ -107,10 +108,10 @@ def basic_example():
               }
               condition {
                 scalar_function {
-                  function_reference: 495
+                  function_reference: 124
                   output_type {
                     bool {
-                      nullability: NULLABILITY_NULLABLE
+                      nullability: NULLABILITY_REQUIRED
                     }
                   }
                   arguments {
@@ -129,7 +130,6 @@ def basic_example():
                     value {
                       literal {
                         i64: 100
-                        nullable: true
                       }
                     }
                   }
@@ -152,25 +152,29 @@ def basic_example():
       names: "id"
     }
   }
-  """
+  extension_urns {
+    extension_urn_anchor: 2
+    urn: "extension:io.substrait:functions_comparison"
+  }
+    """
 
 
 def advanced_example():
     print("=== Simple Example ===")
     # Simple example (original)
     ns = named_struct(
         names=["id", "is_applicable"],
-        struct=struct(types=[i64(nullable=False), boolean()]),
+        struct=struct(types=[i64(nullable=False), boolean()], nullable=False),
     )
 
     table = read_named_table("example_table", ns)
     table = filter(table, expression=column("is_applicable"))
     table = filter(
         table,
         expression=scalar_function(
-            "functions_comparison.yaml",
+            "extension:io.substrait:functions_comparison",
             "lt",
-            expressions=[column("id"), literal(100, i64())],
+            expressions=[column("id"), literal(100, i64(nullable=False))],
         ),
     )
     table = project(table, expressions=[column("id")])
@@ -190,7 +194,8 @@ def advanced_example():
                 string(nullable=False),  # name
                 i64(nullable=False),  # age
                 fp64(nullable=False),  # salary
-            ]
+            ],
+            nullable=False,
         ),
     )
 
@@ -200,7 +205,7 @@ def advanced_example():
     adult_users = filter(
         users,
         expression=scalar_function(
-            "functions_comparison.yaml",
+            "extension:io.substrait:functions_comparison",
             "gt",
             expressions=[column("age"), literal(25, i64())],
         ),
@@ -216,7 +221,7 @@ def advanced_example():
             column("salary"),
             # Add a calculated field (this would show function options if available)
             scalar_function(
-                "functions_arithmetic.yaml",
+                "extension:io.substrait:functions_arithmetic",
                 "multiply",
                 expressions=[column("salary"), literal(1.1, fp64())],
                 alias="salary_with_bonus",
@@ -238,7 +243,8 @@ def advanced_example():
                 i64(nullable=False),  # order_id
                 fp64(nullable=False),  # amount
                 string(nullable=False),  # status
-            ]
+            ],
+            nullable=False,
         ),
     )
 
@@ -248,7 +254,7 @@ def advanced_example():
     high_value_orders = filter(
         orders,
         expression=scalar_function(
-            "functions_comparison.yaml",
+            "extension:io.substrait:functions_comparison",
             "gt",
             expressions=[column("amount"), literal(50.0, fp64())],
         ),
@@ -280,16 +286,16 @@ def expression_only_example():
     print("=== Expression-Only Example ===")
     # Show complex expression structure
     complex_expr = scalar_function(
-        "functions_arithmetic.yaml",
+        "extension:io.substrait:functions_arithmetic",
         "multiply",
         expressions=[
             scalar_function(
-                "functions_arithmetic.yaml",
+                "extension:io.substrait:functions_arithmetic",
                 "add",
                 expressions=[
                     column("base_salary"),
                     scalar_function(
-                        "functions_arithmetic.yaml",
+                        "extension:io.substrait:functions_arithmetic",
                         "multiply",
                         expressions=[
                             column("base_salary"),
@@ -299,7 +305,7 @@ def expression_only_example():
                 ],
             ),
             scalar_function(
-                "functions_arithmetic.yaml",
+                "extension:io.substrait:functions_arithmetic",
                 "subtract",
                 expressions=[
                     literal(1.0, fp64()),
@@ -312,7 +318,8 @@ def expression_only_example():
     print("Complex salary calculation expression:")
     # Create a simple plan to wrap the expression
     dummy_schema = named_struct(
-        names=["base_salary"], struct=struct(types=[fp64(nullable=False)])
+        names=["base_salary"],
+        struct=struct(types=[fp64(nullable=False)], nullable=False),
     )
     dummy_table = read_named_table("dummy", dummy_schema)
     dummy_plan = project(dummy_table, expressions=[complex_expr])

examples/duckdb_example.py

@@ -13,7 +13,6 @@
 from substrait.builders.extended_expression import column, scalar_function, literal
 from substrait.builders.type import i32
 from substrait.extension_registry import ExtensionRegistry
-from substrait.json import dump_json
 import pyarrow.substrait as pa_substrait
 
 try:
@@ -42,14 +41,13 @@ def read_duckdb_named_table(name: str, conn):
 table = filter(
     table,
     expression=scalar_function(
-        "functions_comparison.yaml",
+        "extension:io.substrait:functions_comparison",
         "equal",
         expressions=[column("c_nationkey"), literal(3, i32())],
     ),
 )
 table = project(
     table, expressions=[column("c_name"), column("c_address"), column("c_nationkey")]
 )
-
-sql = f"CALL from_substrait_json('{dump_json(table(registry))}')"
-print(duckdb.sql(sql))
+sql = "CALL from_substrait(?)"
+print(duckdb.sql(sql, params=[table(registry).SerializeToString()]))

gen_proto.sh

@@ -18,7 +18,7 @@ rm -rf "$dest_dir/proto"
 
 # Generate the new python protobuf files
 buf generate
-protol --in-place --create-package --python-out "$dest_dir" buf
+uv run protol --in-place --create-package --python-out "$dest_dir" buf
 
 # Remove the old extension files
 rm -rf "$extension_dir"

src/substrait/__init__.py

@@ -3,6 +3,6 @@
 except ImportError:
     pass
 
-__substrait_version__ = "0.74.0"
-__substrait_hash__ = "793c64b"
+__substrait_version__ = "0.77.0"
+__substrait_hash__ = "3c25b1b"
 __minimum_substrait_version__ = "0.30.0"