feat: graceful migration from URI -> URN

This PR introduces handling of migrating from the usage of URI to URN
for extension references. As an intermediate step, both URI and URN
and emitted from produced plans.

Closes #95

Author: benbellick

src/substrait/bimap.py

@@ -0,0 +1,111 @@
+"""
+Bidirectional map for URI <-> URN conversion during the migration period.
+
+This module provides a UriUrnBiDiMap class that maintains a bidirectional mapping
+between URIs and URNs, ensuring consistency and detecting conflicts.
+
+NOTE: This file is temporary and can be removed once the URI -> URN migration
+is complete across all Substrait implementations. At that point, only URN-based
+extension references will be used.
+"""
+
+from typing import Optional
+
+
+class UriUrnBiDiMap:
+    """Bidirectional map for URI <-> URN mappings.
+
+    Maintains two internal dictionaries to enable O(1) lookups in both directions.
+    Enforces that each URI maps to exactly one URN and vice versa.
+    """
+
+    def __init__(self):
+        self._uri_to_urn: dict[str, str] = {}
+        self._urn_to_uri: dict[str, str] = {}
+
+    def put(self, uri: str, urn: str) -> None:
+        """Add a bidirectional URI <-> URN mapping.
+
+        Args:
+            uri: The extension URI (e.g., "https://github.com/.../functions_arithmetic.yaml")
+            urn: The extension URN (e.g., "extension:io.substrait:functions_arithmetic")
+
+        Raises:
+            ValueError: If the URI or URN already exists with a different mapping
+        """
+        # Check for conflicts
+        if uri in self._uri_to_urn:
+            existing_urn = self._uri_to_urn[uri]
+            if existing_urn != urn:
+                raise ValueError(
+                    f"URI '{uri}' is already mapped to URN '{existing_urn}', "
+                    f"cannot remap to '{urn}'"
+                )
+            # Already have this exact mapping, nothing to do
+            return
+
+        if urn in self._urn_to_uri:
+            existing_uri = self._urn_to_uri[urn]
+            if existing_uri != uri:
+                raise ValueError(
+                    f"URN '{urn}' is already mapped to URI '{existing_uri}', "
+                    f"cannot remap to '{uri}'"
+                )
+            # Already have this exact mapping, nothing to do
+            return
+
+        # Add the bidirectional mapping
+        self._uri_to_urn[uri] = urn
+        self._urn_to_uri[urn] = uri
+
+    def get_urn(self, uri: str) -> Optional[str]:
+        """Convert a URI to its corresponding URN.
+
+        Args:
+            uri: The extension URI to look up
+
+        Returns:
+            The corresponding URN, or None if the URI is not in the map
+        """
+        return self._uri_to_urn.get(uri)
+
+    def get_uri(self, urn: str) -> Optional[str]:
+        """Convert a URN to its corresponding URI.
+
+        Args:
+            urn: The extension URN to look up
+
+        Returns:
+            The corresponding URI, or None if the URN is not in the map
+        """
+        return self._urn_to_uri.get(urn)
+
+    def contains_uri(self, uri: str) -> bool:
+        """Check if a URI exists in the map.
+
+        Args:
+            uri: The URI to check
+
+        Returns:
+            True if the URI is in the map, False otherwise
+        """
+        return uri in self._uri_to_urn
+
+    def contains_urn(self, urn: str) -> bool:
+        """Check if a URN exists in the map.
+
+        Args:
+            urn: The URN to check
+
+        Returns:
+            True if the URN is in the map, False otherwise
+        """
+        return urn in self._urn_to_uri
+
+    def __len__(self) -> int:
+        """Return the number of mappings in the bimap."""
+        return len(self._uri_to_urn)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the bimap."""
+        return f"UriUrnBiDiMap({len(self)} mappings)"

src/substrait/builders/extended_expression.py

@@ -8,6 +8,7 @@
 from substrait.utils import (
     type_num_names,
     merge_extension_urns,
+    merge_extension_uris,
     merge_extension_declarations,
 )
 from substrait.type_inference import infer_extended_expression_schema
@@ -229,26 +230,44 @@ def resolve(
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        # Create URN extension
         func_extension_urns = [
             ste.SimpleExtensionURN(
                 extension_urn_anchor=registry.lookup_urn(urn), urn=urn
             )
         ]
 
+        # Create URI extension (convert URN to URI via bimap)
+        uri = registry.urn_to_uri(urn)
+        func_extension_uris = []
+        if uri:
+            uri_anchor = registry.lookup_uri_anchor(uri)
+            if uri_anchor:
+                func_extension_uris = [
+                    ste.SimpleExtensionURI(extension_uri_anchor=uri_anchor, uri=uri)
+                ]
+
+        # Create extension function declaration with both URI and URN references
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
                     extension_urn_reference=registry.lookup_urn(urn),
+                    extension_uri_reference=registry.lookup_uri_anchor(uri) if uri else 0,
                     function_anchor=func[0].anchor,
                     name=str(func[0]),
                 )
             )
         ]
 
+        # Merge extensions from all expressions
         extension_urns = merge_extension_urns(
             func_extension_urns, *[b.extension_urns for b in bound_expressions]
         )
 
+        extension_uris = merge_extension_uris(
+            func_extension_uris, *[b.extension_uris for b in bound_expressions]
+        )
+
         extensions = merge_extension_declarations(
             func_extensions, *[b.extensions for b in bound_expressions]
         )
@@ -277,6 +296,7 @@ def resolve(
             ],
             base_schema=base_schema,
             extension_urns=extension_urns,
+            extension_uris=extension_uris,
             extensions=extensions,
         )
 
@@ -309,26 +329,44 @@ def resolve(
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        # Create URN extension
         func_extension_urns = [
             ste.SimpleExtensionURN(
                 extension_urn_anchor=registry.lookup_urn(urn), urn=urn
             )
         ]
 
+        # Create URI extension (convert URN to URI via bimap)
+        uri = registry.urn_to_uri(urn)
+        func_extension_uris = []
+        if uri:
+            uri_anchor = registry.lookup_uri_anchor(uri)
+            if uri_anchor:
+                func_extension_uris = [
+                    ste.SimpleExtensionURI(extension_uri_anchor=uri_anchor, uri=uri)
+                ]
+
+        # Create extension function declaration with both URI and URN references
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
                     extension_urn_reference=registry.lookup_urn(urn),
+                    extension_uri_reference=registry.lookup_uri_anchor(uri) if uri else 0,
                     function_anchor=func[0].anchor,
                     name=str(func[0]),
                 )
             )
         ]
 
+        # Merge extensions from all expressions
         extension_urns = merge_extension_urns(
             func_extension_urns, *[b.extension_urns for b in bound_expressions]
         )
 
+        extension_uris = merge_extension_uris(
+            func_extension_uris, *[b.extension_uris for b in bound_expressions]
+        )
+
         extensions = merge_extension_declarations(
             func_extensions, *[b.extensions for b in bound_expressions]
         )
@@ -353,6 +391,7 @@ def resolve(
             ],
             base_schema=base_schema,
             extension_urns=extension_urns,
+            extension_uris=extension_uris,
             extensions=extensions,
         )
 
@@ -391,28 +430,48 @@ def resolve(
         if not func:
             raise Exception(f"Unknown function {function} for {signature}")
 
+        # Create URN extension
         func_extension_urns = [
             ste.SimpleExtensionURN(
                 extension_urn_anchor=registry.lookup_urn(urn), urn=urn
             )
         ]
 
+        # Create URI extension (convert URN to URI via bimap)
+        uri = registry.urn_to_uri(urn)
+        func_extension_uris = []
+        if uri:
+            uri_anchor = registry.lookup_uri_anchor(uri)
+            if uri_anchor:
+                func_extension_uris = [
+                    ste.SimpleExtensionURI(extension_uri_anchor=uri_anchor, uri=uri)
+                ]
+
+        # Create extension function declaration with both URI and URN references
         func_extensions = [
             ste.SimpleExtensionDeclaration(
                 extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction(
                     extension_urn_reference=registry.lookup_urn(urn),
+                    extension_uri_reference=registry.lookup_uri_anchor(uri) if uri else 0,
                     function_anchor=func[0].anchor,
                     name=str(func[0]),
                 )
             )
         ]
 
+        # Merge extensions from all expressions
         extension_urns = merge_extension_urns(
             func_extension_urns,
             *[b.extension_urns for b in bound_expressions],
             *[b.extension_urns for b in bound_partitions],
         )
 
+        extension_uris = merge_extension_uris(
+            func_extension_uris,
+            *[b.extension_uris for b in bound_expressions],
+            *[b.extension_uris for b in bound_partitions],
+        )
+
         extensions = merge_extension_declarations(
             func_extensions,
             *[b.extensions for b in bound_expressions],
@@ -446,6 +505,7 @@ def resolve(
             ],
             base_schema=base_schema,
             extension_urns=extension_urns,
+            extension_uris=extension_uris,
             extensions=extensions,
         )
 

src/substrait/builders/plan.py

@@ -18,15 +18,21 @@
     resolve_expression,
 )
 from substrait.type_inference import infer_plan_schema
-from substrait.utils import merge_extension_declarations, merge_extension_urns
+from substrait.utils import merge_extension_declarations, merge_extension_urns, merge_extension_uris
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 
 PlanOrUnbound = Union[stp.Plan, UnboundPlan]
 
 
 def _merge_extensions(*objs):
+    """Merge extension URIs, URNs, and declarations from multiple plan/expression objects.
+
+    During the URI/URN migration period, we maintain both URI and URN references
+    for maximum compatibility.
+    """
     return {
+        "extension_uris": merge_extension_uris(*[b.extension_uris for b in objs if b]),
         "extension_urns": merge_extension_urns(*[b.extension_urns for b in objs if b]),
         "extensions": merge_extension_declarations(*[b.extensions for b in objs if b]),
     }