cleanup

Author: benbellick

.gitignore

@@ -142,4 +142,3 @@ src/substrait/_version.py
 .directory
 .gdb_history
 .DS_Store
-/.envrc

src/substrait/bimap.py

@@ -2,7 +2,7 @@
 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.
+between URIs and URNs.
 
 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
@@ -34,8 +34,8 @@ def put(self, uri: str, urn: str) -> None:
             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 self.contains_uri(uri):
+            existing_urn = self.get_urn(uri)
             if existing_urn != urn:
                 raise ValueError(
                     f"URI '{uri}' is already mapped to URN '{existing_urn}', "
@@ -44,8 +44,8 @@ def put(self, uri: str, urn: str) -> None:
             # Already have this exact mapping, nothing to do
             return
 
-        if urn in self._urn_to_uri:
-            existing_uri = self._urn_to_uri[urn]
+        if self.contains_urn(urn):
+            existing_uri = self.get_uri(urn)
             if existing_uri != uri:
                 raise ValueError(
                     f"URN '{urn}' is already mapped to URI '{existing_uri}', "
@@ -54,7 +54,6 @@ def put(self, uri: str, urn: str) -> None:
             # 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
 
@@ -101,11 +100,3 @@ def contains_urn(self, urn: str) -> bool:
             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

@@ -230,15 +230,14 @@ 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)
+        # Create URI extension for backwards compatibility during URI -> URN migration
+        uri = registry._uri_urn_bimap.get_uri(urn)
         func_extension_uris = []
         if uri:
             uri_anchor = registry.lookup_uri_anchor(uri)
@@ -247,7 +246,6 @@ def resolve(
                     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(
@@ -259,7 +257,6 @@ def resolve(
             )
         ]
 
-        # Merge extensions from all expressions
         extension_urns = merge_extension_urns(
             func_extension_urns, *[b.extension_urns for b in bound_expressions]
         )
@@ -329,15 +326,14 @@ 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)
+        # Create URI extension for backwards compatibility during URI -> URN migration
+        uri = registry._uri_urn_bimap.get_uri(urn)
         func_extension_uris = []
         if uri:
             uri_anchor = registry.lookup_uri_anchor(uri)
@@ -346,7 +342,6 @@ def resolve(
                     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(
@@ -358,7 +353,6 @@ def resolve(
             )
         ]
 
-        # Merge extensions from all expressions
         extension_urns = merge_extension_urns(
             func_extension_urns, *[b.extension_urns for b in bound_expressions]
         )
@@ -430,15 +424,14 @@ 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)
+        # Create URI extension for backwards compatibility during URI -> URN migration
+        uri = registry._uri_urn_bimap.get_uri(urn)
         func_extension_uris = []
         if uri:
             uri_anchor = registry.lookup_uri_anchor(uri)
@@ -447,7 +440,6 @@ def resolve(
                     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(
@@ -459,7 +451,7 @@ def resolve(
             )
         ]
 
-        # Merge extensions from all expressions
+
         extension_urns = merge_extension_urns(
             func_extension_urns,
             *[b.extension_urns for b in bound_expressions],

src/substrait/builders/plan.py

@@ -28,8 +28,8 @@
 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.
+    During the URI -> URN migration period, we maintain both URI and URN references
+    for backwards compatibility. 
     """
     return {
         "extension_uris": merge_extension_uris(*[b.extension_uris for b in objs if b]),

src/substrait/extension_registry.py

@@ -251,6 +251,7 @@ def satisfies_signature(self, signature: tuple) -> Optional[str]:
 class ExtensionRegistry:
     def __init__(self, load_default_extensions=True) -> None:
         self._urn_mapping: dict = defaultdict(dict)  # URN -> anchor ID
+        # NOTE: during the URI -> URN migration, we only need an id generator for URN. We can use the same anchor for plan construction for URIs.
         self._urn_id_generator = itertools.count(1)
 
         self._function_mapping: dict = defaultdict(dict)
@@ -278,7 +279,7 @@ def register_extension_yaml(
 
         Args:
             fname: Path to the YAML file
-            uri: URI for the extension (for URI/URN bimap)
+            uri: URI for the extension (this is required during the URI -> URN migration)
         """
         fname = Path(fname)
         with open(fname) as f:  # type: ignore
@@ -297,13 +298,10 @@ def register_extension_dict(self, definitions: dict, uri: str) -> None:
         if not urn:
             raise ValueError("Extension definitions must contain a 'urn' field")
 
-        # Validate URN format
         self._validate_urn_format(urn)
 
-        # Assign anchor to URN (URI will use the same anchor during output)
         self._urn_mapping[urn] = next(self._urn_id_generator)
 
-        # Store URI <-> URN mapping for output generation
         self._uri_urn_bimap.put(uri, urn)
 
         simple_extensions = build_simple_extensions(definitions)
@@ -369,28 +367,6 @@ def lookup_uri_anchor(self, uri: str) -> Optional[int]:
             return self._urn_mapping.get(urn)
         return None
 
-    def uri_to_urn(self, uri: str) -> Optional[str]:
-        """Convert a URI to its corresponding URN using the bimap.
-
-        Args:
-            uri: The extension URI
-
-        Returns:
-            The corresponding URN, or None if not in the bimap
-        """
-        return self._uri_urn_bimap.get_urn(uri)
-
-    def urn_to_uri(self, urn: str) -> Optional[str]:
-        """Convert a URN to its corresponding URI using the bimap.
-
-        Args:
-            urn: The extension URN
-
-        Returns:
-            The corresponding URI, or None if not in the bimap
-        """
-        return self._uri_urn_bimap.get_uri(urn)
-
     def _validate_urn_format(self, urn: str) -> None:
         """Validate that a URN follows the expected format.
 

tests/test_bimap.py

@@ -78,24 +78,3 @@ def test_contains(self):
 
         assert bimap.contains_uri(uri)
         assert bimap.contains_urn(urn)
-
-    def test_len(self):
-        """Test the length of the bimap."""
-        bimap = UriUrnBiDiMap()
-
-        assert len(bimap) == 0
-
-        bimap.put("uri1", "urn1")
-        assert len(bimap) == 1
-
-        bimap.put("uri2", "urn2")
-        assert len(bimap) == 2
-
-    def test_repr(self):
-        """Test the string representation of the bimap."""
-        bimap = UriUrnBiDiMap()
-
-        assert repr(bimap) == "UriUrnBiDiMap(0 mappings)"
-
-        bimap.put("uri1", "urn1")
-        assert repr(bimap) == "UriUrnBiDiMap(1 mappings)"

tests/test_extension_registry.py

@@ -363,7 +363,7 @@ def test_registry_uri_to_urn_conversion():
     registry.register_extension_dict(yaml.safe_load(content_with_urn), uri=uri)
 
     # Test URI to URN conversion
-    assert registry.uri_to_urn(uri) == "extension:test:bimap"
+    assert registry._uri_urn_bimap.get_urn(uri) == "extension:test:bimap"
 
 
 def test_registry_urn_to_uri_conversion():
@@ -384,7 +384,7 @@ def test_registry_urn_to_uri_conversion():
     registry.register_extension_dict(yaml.safe_load(content_with_urn), uri=uri)
 
     # Test URN to URI conversion
-    assert registry.urn_to_uri("extension:test:bimap2") == uri
+    assert registry._uri_urn_bimap.get_uri("extension:test:bimap2") == uri
 
 
 def test_registry_uri_anchor_lookup():
@@ -408,8 +408,8 @@ def test_registry_nonexistent_uri_urn_returns_none():
     """Test that looking up non-existent URI/URN returns None."""
     registry = ExtensionRegistry(load_default_extensions=False)
 
-    assert registry.uri_to_urn("https://nonexistent.com/test.yaml") is None
-    assert registry.urn_to_uri("extension:nonexistent:test") is None
+    assert registry._uri_urn_bimap.get_urn("https://nonexistent.com/test.yaml") is None
+    assert registry._uri_urn_bimap.get_uri("extension:nonexistent:test") is None
     assert registry.lookup_uri_anchor("https://nonexistent.com/test.yaml") is None
 
 
@@ -419,14 +419,14 @@ def test_registry_default_extensions_have_uri_mappings():
 
     # Check that at least one default extension has a URI mapping
     urn = "extension:io.substrait:functions_comparison"
-    uri = registry.urn_to_uri(urn)
+    uri = registry._uri_urn_bimap.get_uri(urn)
 
     assert uri is not None
     assert "https://github.com/substrait-io/substrait/blob/main/extensions" in uri
     assert "functions_comparison.yaml" in uri
 
     # Verify reverse mapping works
-    assert registry.uri_to_urn(uri) == urn
+    assert registry._uri_urn_bimap.get_urn(uri) == urn
 
 
 # ============================================================================

tests/test_uri_urn_migration.py

@@ -53,8 +53,8 @@ def test_register_with_uri(self):
         assert registry.lookup_uri_anchor(uri) is not None
 
         # Test bimap conversions
-        assert registry.uri_to_urn(uri) == "extension:example:test"
-        assert registry.urn_to_uri("extension:example:test") == uri
+        assert registry._uri_urn_bimap.get_urn(uri) == "extension:example:test"
+        assert registry._uri_urn_bimap.get_uri("extension:example:test") == uri
 
     def test_register_requires_uri(self):
         """Test that registering an extension requires a URI during migration."""
@@ -75,7 +75,7 @@ def test_default_extensions_have_uris(self):
 
         # Check one of the default extensions
         urn = "extension:io.substrait:functions_comparison"
-        uri_from_bimap = registry.urn_to_uri(urn)
+        uri_from_bimap = registry._uri_urn_bimap.get_uri(urn)
 
         # Should have a URI derived from DEFAULT_URN_PREFIX
         assert uri_from_bimap is not None
@@ -165,8 +165,8 @@ def test_uri_and_urn_always_paired(self):
         assert registry.lookup_uri_anchor(uri) is not None
 
         # Verify bimap has both directions
-        assert registry.uri_to_urn(uri) == "extension:test:functions_paired"
-        assert registry.urn_to_uri("extension:test:functions_paired") == uri
+        assert registry._uri_urn_bimap.get_urn(uri) == "extension:test:functions_paired"
+        assert registry._uri_urn_bimap.get_uri("extension:test:functions_paired") == uri
 
 
 # ============================================================================