track urn<->uri mapping internally

Author: benbellick

core/src/main/java/io/substrait/extension/BidiMap.java

@@ -15,8 +15,8 @@ public BidiMap(Map<T1, T2> forwardMap) {
   }
 
   public BidiMap() {
-      this.forwardMap = new HashMap<>();
-      this.reverseMap = new HashMap<>();
+    this.forwardMap = new HashMap<>();
+    this.reverseMap = new HashMap<>();
   }
 
   public T2 get(T1 t1) {
@@ -28,15 +28,33 @@ public T1 reverseGet(T2 t2) {
   }
 
   public void put(T1 t1, T2 t2) {
+    // Check for conflicting mappings (different values for same key)
+    T2 existingForward = forwardMap.get(t1);
+    T1 existingReverse = reverseMap.get(t2);
+
+    if (existingForward != null && !existingForward.equals(t2)) {
+      throw new IllegalArgumentException("Key already exists in map with different value");
+    }
+    if (existingReverse != null && !existingReverse.equals(t1)) {
+      throw new IllegalArgumentException("Key already exists in map with different value");
+    }
+
+    // Allow identical mappings, only add if not already present
     forwardMap.put(t1, t2);
     reverseMap.put(t2, t1);
   }
 
+  public void merge(BidiMap<T1, T2> other) {
+    for (Map.Entry<T1, T2> entry : other.forwardEntrySet()) {
+      put(entry.getKey(), entry.getValue());
+    }
+  }
+
   public Set<Map.Entry<T1, T2>> forwardEntrySet() {
-      return forwardMap.entrySet();
+    return forwardMap.entrySet();
   }
 
-    public Set<Map.Entry<T2, T1>> reverseEntrySet() {
-        return reverseMap.entrySet();
-    }
+  public Set<Map.Entry<T2, T1>> reverseEntrySet() {
+    return reverseMap.entrySet();
+  }
 }

core/src/main/java/io/substrait/extension/SimpleExtension.java

@@ -52,11 +52,11 @@ private static void validateUrn(String urn) {
       throw new IllegalArgumentException("URN cannot be null or empty");
     }
     if (!urn.matches("^extension:[^:]+:[^:]+$")) {
-      throw new IllegalArgumentException("URN must follow format 'extension:<namespace>:<name>', got: " + urn);
+      throw new IllegalArgumentException(
+          "URN must follow format 'extension:<namespace>:<name>', got: " + urn);
     }
   }
 
-
   private static ObjectMapper objectMapper(String urn) {
     InjectableValues.Std iv = new InjectableValues.Std();
     iv.addValue(URN_LOCATOR_KEY, urn);
@@ -551,6 +551,13 @@ public abstract static class ExtensionSignatures {
     @JsonProperty("urn")
     public abstract String urn();
 
+    // URI is not from YAML, but from the loading context
+    // this only needs to be present temporarily to handle the URI -> URN migration
+    @Value.Default
+    public String uri() {
+      return "";
+    }
+
     @JsonProperty("scalar_functions")
     public abstract List<ScalarFunction> scalars();
 
@@ -580,6 +587,11 @@ public Stream<SimpleExtension.Function> resolve(String urn) {
 
   @Value.Immutable
   public abstract static class ExtensionCollection {
+    @Value.Default
+    BidiMap<String, String> uriUrnMap() {
+      return new BidiMap<>();
+    }
+
     private final Supplier<Set<String>> urnSupplier =
         Util.memoize(
             () -> {
@@ -700,7 +712,55 @@ public WindowFunctionVariant getWindowFunction(FunctionAnchor anchor) {
               anchor.key(), anchor.urn()));
     }
 
+    /**
+     * Gets the URN for a given URI. This is only useful during the URI -> URN migration, and will
+     * be dropped when the migration is complete.
+     *
+     * @param uri The URI to look up
+     * @return The corresponding URN, or null if not found
+     */
+    public String getUrn(String uri) {
+      return uriUrnMap().get(uri);
+    }
+
+    /**
+     * Gets the URI for a given URN. This is only useful during the URI -> URN migration, and will
+     * be dropped when the migration is complete.
+     *
+     * @param urn The URN to look up
+     * @return The corresponding URI, or null if not found
+     */
+    public String getUri(String urn) {
+      return uriUrnMap().reverseGet(urn);
+    }
+
+    /**
+     * Checks if a URI has a corresponding URN mapping. This is only useful during the URI -> URN
+     * migration, and will be dropped when the migration is complete.
+     *
+     * @param uri The URI to check
+     * @return true if the URI has a URN mapping, false otherwise
+     */
+    public boolean hasUrn(String uri) {
+      return uriUrnMap().get(uri) != null;
+    }
+
+    /**
+     * Checks if a URN has a corresponding URI mapping. This is only useful during the URI -> URN
+     * migration, and will be dropped when the migration is complete.
+     *
+     * @param urn The URN to check
+     * @return true if the URN has a URI mapping, false otherwise
+     */
+    public boolean hasUri(String urn) {
+      return uriUrnMap().reverseGet(urn) != null;
+    }
+
     public ExtensionCollection merge(ExtensionCollection extensionCollection) {
+      BidiMap<String, String> mergedUriUrnMap = new BidiMap<>();
+      mergedUriUrnMap.merge(uriUrnMap());
+      mergedUriUrnMap.merge(extensionCollection.uriUrnMap());
+
       return ImmutableSimpleExtension.ExtensionCollection.builder()
           .addAllAggregateFunctions(aggregateFunctions())
           .addAllAggregateFunctions(extensionCollection.aggregateFunctions())
@@ -710,6 +770,7 @@ public ExtensionCollection merge(ExtensionCollection extensionCollection) {
           .addAllWindowFunctions(extensionCollection.windowFunctions())
           .addAllTypes(types())
           .addAllTypes(extensionCollection.types())
+          .uriUrnMap(mergedUriUrnMap)
           .build();
     }
   }
@@ -740,12 +801,12 @@ private static ExtensionCollection load(List<String> resourcePaths) {
       throw new IllegalArgumentException("Require at least one resource path.");
     }
 
-        List<ExtensionCollection> extensions =
+    List<ExtensionCollection> extensions =
         resourcePaths.stream()
             .map(
                 path -> {
                   try (InputStream stream = ExtensionCollection.class.getResourceAsStream(path)) {
-                    return load(stream);
+                    return load(path, stream);
                   } catch (IOException e) {
                     throw new UncheckedIOException(e);
                   }
@@ -758,40 +819,48 @@ private static ExtensionCollection load(List<String> resourcePaths) {
     return complete;
   }
 
-  public static ExtensionCollection load(String content) {
+  public static ExtensionCollection load(String uri, String content) {
     try {
-      // Parse with basic YAML mapper first to extract URN (if present)
+      // Parse with basic YAML mapper first to extract URN
       ObjectMapper basicYamlMapper = new ObjectMapper(new YAMLFactory());
       com.fasterxml.jackson.databind.JsonNode rootNode = basicYamlMapper.readTree(content);
-
-      // URN is required
       com.fasterxml.jackson.databind.JsonNode urnNode = rootNode.get("urn");
       if (urnNode == null) {
         throw new IllegalArgumentException("Extension YAML file must contain a 'urn' field");
       }
       String urn = urnNode.asText();
       validateUrn(urn);
 
-      // Then parse with URN-aware mapper
-      ExtensionSignatures doc = objectMapper(urn).readValue(content, ExtensionSignatures.class);
-      return buildExtensionCollection(urn, doc);
+      ExtensionSignatures docWithoutUri =
+          objectMapper(urn).readValue(content, ExtensionSignatures.class);
+
+      ExtensionSignatures doc =
+          ImmutableSimpleExtension.ExtensionSignatures.builder()
+              .from(docWithoutUri)
+              .uri(uri)
+              .build();
+
+      return buildExtensionCollection(uri, doc);
     } catch (IOException e) {
       throw new IllegalStateException(e);
     }
   }
 
-    public static ExtensionCollection load(InputStream stream) {
-        try (Scanner scanner = new Scanner(stream)) {
-            scanner.useDelimiter("\\A");
-            String content = scanner.next();
-            return load(content);
-        }
+  public static ExtensionCollection load(String uri, InputStream stream) {
+    try (Scanner scanner = new Scanner(stream)) {
+      scanner.useDelimiter("\\A");
+      String content = scanner.next();
+      return load(uri, content);
     }
-
+  }
 
   public static ExtensionCollection buildExtensionCollection(
-      String urn, ExtensionSignatures extensionSignatures) {
+      String uri, ExtensionSignatures extensionSignatures) {
+    String urn = extensionSignatures.urn();
     validateUrn(urn);
+    if (uri == null || uri == "") {
+      throw new IllegalArgumentException("URI cannot be null or empty");
+    }
     List<ScalarFunctionVariant> scalarFunctionVariants =
         extensionSignatures.scalars().stream()
             .flatMap(t -> t.resolve(urn))
@@ -824,23 +893,23 @@ public static ExtensionCollection buildExtensionCollection(
         Stream.concat(windowFunctionVariants, windowAggFunctionVariants)
             .collect(Collectors.toList());
 
+    BidiMap<String, String> uriUrnMap = new BidiMap<>();
+    uriUrnMap.put(uri, urn);
+
     ImmutableSimpleExtension.ExtensionCollection collection =
         ImmutableSimpleExtension.ExtensionCollection.builder()
             .scalarFunctions(scalarFunctionVariants)
             .aggregateFunctions(aggregateFunctionVariants)
             .windowFunctions(allWindowFunctionVariants)
             .addAllTypes(extensionSignatures.types())
+            .uriUrnMap(uriUrnMap)
             .build();
+
     LOGGER.atDebug().log(
         "Loaded {} aggregate functions and {} scalar functions from {}.",
         collection.aggregateFunctions().size(),
-        collection.scalarFunctions().size(), extensionSignatures.urn());
+        collection.scalarFunctions().size(),
+        extensionSignatures.urn());
     return collection;
   }
-
-  public static ExtensionCollection buildExtensionCollection(
-      ExtensionSignatures extensionSignatures) {
-    String urn = extensionSignatures.urn();
-    return buildExtensionCollection(urn, extensionSignatures);
-  }
 }

core/src/test/java/io/substrait/extension/TypeExtensionTest.java

@@ -32,9 +32,9 @@ public class TypeExtensionTest {
   final SimpleExtension.ExtensionCollection extensionCollection;
 
   {
-    InputStream inputStream =
-        this.getClass().getResourceAsStream("/extensions/custom_extensions.yaml");
-    extensionCollection = SimpleExtension.load(inputStream);
+    String path = "/extensions/custom_extensions.yaml";
+    InputStream inputStream = this.getClass().getResourceAsStream(path);
+    extensionCollection = SimpleExtension.load(path, inputStream);
   }
 
   final SubstraitBuilder b = new SubstraitBuilder(extensionCollection);

core/src/test/java/io/substrait/extension/UrnValidationTest.java

@@ -7,38 +7,43 @@ public class UrnValidationTest {
 
   @Test
   public void testMissingUrnThrowsException() {
-    String yamlWithoutUrn = """
-        %YAML 1.2
-        ---
-        scalar_functions:
-          - name: test
-        """;
-    IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load(yamlWithoutUrn));
+    String yamlWithoutUrn = "%YAML 1.2\n" +
+        "---\n" +
+        "scalar_functions:\n" +
+        "  - name: test\n";
+    IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load("test://uri", yamlWithoutUrn));
     assertTrue(exception.getMessage().contains("Extension YAML file must contain a 'urn' field"));
   }
 
   @Test
   public void testInvalidUrnFormatThrowsException() {
-    String yamlWithInvalidUrn = """
-        %YAML 1.2
-        ---
-        urn: invalid:format
-        scalar_functions:
-          - name: test
-        """;
-    IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load(yamlWithInvalidUrn));
+    String yamlWithInvalidUrn = "%YAML 1.2\n" +
+        "---\n" +
+        "urn: invalid:format\n" +
+        "scalar_functions:\n" +
+        "  - name: test\n";
+    IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load("test://uri", yamlWithInvalidUrn));
     assertTrue(exception.getMessage().contains("URN must follow format 'extension:<namespace>:<name>'"));
   }
 
   @Test
   public void testValidUrnWorks() {
-    String yamlWithValidUrn = """
-        %YAML 1.2
-        ---
-        urn: extension:test:valid
-        scalar_functions:
-          - name: test
-        """;
-    assertDoesNotThrow(() -> SimpleExtension.load(yamlWithValidUrn));
+    String yamlWithValidUrn = "%YAML 1.2\n" +
+        "---\n" +
+        "urn: extension:test:valid\n" +
+        "scalar_functions:\n" +
+        "  - name: test\n";
+    assertDoesNotThrow(() -> SimpleExtension.load("test://uri", yamlWithValidUrn));
+  }
+
+  @Test
+  public void testUriUrnMapIsPopulated() {
+    String yamlWithValidUrn = "%YAML 1.2\n" +
+        "---\n" +
+        "urn: extension:test:valid\n" +
+        "scalar_functions:\n" +
+        "  - name: test\n";
+    SimpleExtension.ExtensionCollection collection = SimpleExtension.load("test://uri", yamlWithValidUrn);
+    assertEquals("extension:test:valid", collection.getUrn("test://uri"));
   }
 }

isthmus/src/test/java/io/substrait/isthmus/CustomFunctionTest.java

@@ -56,7 +56,7 @@ public class CustomFunctionTest extends PlanTestBase {
 
   // Load custom extension into an ExtensionCollection
   static final SimpleExtension.ExtensionCollection extensionCollection =
-      SimpleExtension.load(FUNCTIONS_CUSTOM);
+      SimpleExtension.load("custom.yaml", FUNCTIONS_CUSTOM);
 
   final SubstraitBuilder b = new SubstraitBuilder(extensionCollection);
 

isthmus/src/test/java/io/substrait/isthmus/RelCopyOnWriteVisitorTest.java

@@ -31,12 +31,13 @@ public class RelCopyOnWriteVisitorTest extends PlanTestBase {
       SimpleExtension.FunctionAnchor.of(
           "extension:io.substrait:functions_aggregate_approx", "approx_count_distinct:any");
   public static SimpleExtension.FunctionAnchor COUNT =
-      SimpleExtension.FunctionAnchor.of("extension:io.substrait:functions_aggregate_generic", "count:any");
+      SimpleExtension.FunctionAnchor.of(
+          "extension:io.substrait:functions_aggregate_generic", "count:any");
 
   private static final String COUNT_DISTINCT_SUBBQUERY =
       "select\n"
           + "  count(distinct l.l_orderkey),\n"
-          + "  count(distinct l.l_orderkey) + 1,\n"
+          + " count(distinct l.l_orderkey) + 1,\n"
           + "  sum(l.l_extendedprice * (1 - l.l_discount)) as revenue,\n"
           + "  o.o_orderdate,\n"
           + "  count(distinct o.o_shippriority)\n"

spark/src/main/scala/io/substrait/spark/SparkExtension.scala

@@ -29,7 +29,7 @@ object SparkExtension {
   final val file = "/spark.yml"
 
   private val SparkImpls: SimpleExtension.ExtensionCollection =
-    SimpleExtension.load(getClass.getResourceAsStream(file))
+    SimpleExtension.load(file, getClass.getResourceAsStream(file))
 
   private val EXTENSION_COLLECTION: SimpleExtension.ExtensionCollection =
     SimpleExtension.loadDefaults()