From 4746a8f189c1841225cf6d63d902461ab2fb1add Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 16:47:07 -0400 Subject: [PATCH 01/13] feat!: bump substrait to v0.85.0, drop URI support Closes #563 --- .../extension/ExtensionCollector.java | 56 +- .../extension/ImmutableExtensionLookup.java | 193 +----- .../substrait/extension/SimpleExtension.java | 41 -- .../ExtensionCollectionMergeTest.java | 31 +- .../ExtensionCollectionUriUrnTest.java | 60 -- .../ExtensionCollectorUriUrnTest.java | 41 -- .../ImmutableExtensionLookupUriUrnTest.java | 621 ------------------ .../UriUrnMigrationEndToEndTest.java | 110 ---- .../extension/UrnValidationTest.java | 14 - .../io/substrait/plan/PlanRoundtripTest.java | 74 +++ .../complex-expected-plan.json | 9 - .../complex-input-plan.json | 18 +- .../simple-expected-plan.json} | 8 - .../simple-input-plan.json} | 0 .../zero-anchor-expected-plan.json} | 7 - .../zero-anchor-input-plan.json} | 1 - .../mixed-partial-coverage-expected-plan.json | 133 ---- .../mixed-partial-coverage-input-plan.json | 110 ---- .../unresolvable-uri-plan.json | 73 -- .../uri-only-expected-plan.json | 98 --- .../uri-only-input-plan.json | 80 --- substrait | 2 +- 22 files changed, 105 insertions(+), 1675 deletions(-) delete mode 100644 core/src/test/java/io/substrait/extension/ExtensionCollectionUriUrnTest.java delete mode 100644 core/src/test/java/io/substrait/extension/ExtensionCollectorUriUrnTest.java delete mode 100644 core/src/test/java/io/substrait/extension/ImmutableExtensionLookupUriUrnTest.java delete mode 100644 core/src/test/java/io/substrait/extension/UriUrnMigrationEndToEndTest.java create mode 100644 core/src/test/java/io/substrait/plan/PlanRoundtripTest.java rename core/src/test/resources/{uri-urn-migration => plan-roundtrip}/complex-expected-plan.json (94%) rename core/src/test/resources/{uri-urn-migration => plan-roundtrip}/complex-input-plan.json (90%) rename core/src/test/resources/{uri-urn-migration/urn-only-expected-plan.json => plan-roundtrip/simple-expected-plan.json} (94%) rename core/src/test/resources/{uri-urn-migration/urn-only-input-plan.json => plan-roundtrip/simple-input-plan.json} (100%) rename core/src/test/resources/{uri-urn-migration/zero-urn-resolution-expected-plan.json => plan-roundtrip/zero-anchor-expected-plan.json} (94%) rename core/src/test/resources/{uri-urn-migration/zero-urn-resolution-input-plan.json => plan-roundtrip/zero-anchor-input-plan.json} (98%) delete mode 100644 core/src/test/resources/uri-urn-migration/mixed-partial-coverage-expected-plan.json delete mode 100644 core/src/test/resources/uri-urn-migration/mixed-partial-coverage-input-plan.json delete mode 100644 core/src/test/resources/uri-urn-migration/unresolvable-uri-plan.json delete mode 100644 core/src/test/resources/uri-urn-migration/uri-only-expected-plan.json delete mode 100644 core/src/test/resources/uri-urn-migration/uri-only-input-plan.json diff --git a/core/src/main/java/io/substrait/extension/ExtensionCollector.java b/core/src/main/java/io/substrait/extension/ExtensionCollector.java index 7ad07a6b1..8091aceaf 100644 --- a/core/src/main/java/io/substrait/extension/ExtensionCollector.java +++ b/core/src/main/java/io/substrait/extension/ExtensionCollector.java @@ -3,7 +3,6 @@ import io.substrait.proto.ExtendedExpression; import io.substrait.proto.Plan; import io.substrait.proto.SimpleExtensionDeclaration; -import io.substrait.proto.SimpleExtensionURI; import io.substrait.proto.SimpleExtensionURN; import java.util.ArrayList; import java.util.HashMap; @@ -20,15 +19,10 @@ public class ExtensionCollector extends AbstractExtensionLookup { private final BidiMap funcMap; private final BidiMap typeMap; - private final SimpleExtension.ExtensionCollection extensionCollection; // start at 0 to make sure functionAnchors start with 1 according to spec private int counter = 0; - private String getUriFromUrn(String urn) { - return extensionCollection.getUriFromUrn(urn); - } - public ExtensionCollector() { this(DefaultExtensionCatalog.DEFAULT_COLLECTION); } @@ -40,7 +34,6 @@ public ExtensionCollector(SimpleExtension.ExtensionCollection extensionCollectio } funcMap = new BidiMap<>(functionAnchorMap); typeMap = new BidiMap<>(typeAnchorMap); - this.extensionCollection = extensionCollection; } public int getFunctionReference(SimpleExtension.Function declaration) { @@ -67,7 +60,6 @@ public void addExtensionsToPlan(Plan.Builder builder) { SimpleExtensions simpleExtensions = getExtensions(); builder.addAllExtensionUrns(simpleExtensions.urns.values()); - builder.addAllExtensionUris(simpleExtensions.uris.values()); builder.addAllExtensions(simpleExtensions.extensionList); } @@ -75,22 +67,17 @@ public void addExtensionsToExtendedExpression(ExtendedExpression.Builder builder SimpleExtensions simpleExtensions = getExtensions(); builder.addAllExtensionUrns(simpleExtensions.urns.values()); - builder.addAllExtensionUris(simpleExtensions.uris.values()); builder.addAllExtensions(simpleExtensions.extensionList); } private SimpleExtensions getExtensions() { AtomicInteger urnPos = new AtomicInteger(1); - AtomicInteger uriPos = new AtomicInteger(1); HashMap urns = new HashMap<>(); - HashMap uris = new HashMap<>(); ArrayList extensionList = new ArrayList<>(); for (Map.Entry e : funcMap.forwardEntrySet()) { String urn = e.getValue().urn(); - String uri = getUriFromUrn(urn); - // Create URN entry SimpleExtensionURN urnObj = urns.computeIfAbsent( urn, @@ -100,30 +87,12 @@ private SimpleExtensions getExtensions() { .setUrn(k) .build()); - // Create URI entry if mapping exists - SimpleExtensionURI uriObj = null; - if (uri != null) { - uriObj = - uris.computeIfAbsent( - uri, - k -> - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(uriPos.getAndIncrement()) - .setUri(k) - .build()); - } - - // Create function declaration with both URN and URI references SimpleExtensionDeclaration.ExtensionFunction.Builder funcBuilder = SimpleExtensionDeclaration.ExtensionFunction.newBuilder() .setFunctionAnchor(e.getKey()) .setName(e.getValue().key()) .setExtensionUrnReference(urnObj.getExtensionUrnAnchor()); - if (uriObj != null) { - funcBuilder.setExtensionUriReference(uriObj.getExtensionUriAnchor()); - } - SimpleExtensionDeclaration decl = SimpleExtensionDeclaration.newBuilder().setExtensionFunction(funcBuilder).build(); extensionList.add(decl); @@ -131,9 +100,7 @@ private SimpleExtensions getExtensions() { for (Map.Entry e : typeMap.forwardEntrySet()) { String urn = e.getValue().urn(); - String uri = getUriFromUrn(urn); - // Create URN entry SimpleExtensionURN urnObj = urns.computeIfAbsent( urn, @@ -143,48 +110,27 @@ private SimpleExtensions getExtensions() { .setUrn(k) .build()); - // Create URI entry if mapping exists - SimpleExtensionURI uriObj = null; - if (uri != null) { - uriObj = - uris.computeIfAbsent( - uri, - k -> - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(uriPos.getAndIncrement()) - .setUri(k) - .build()); - } - - // Create type declaration with both URN and URI references SimpleExtensionDeclaration.ExtensionType.Builder typeBuilder = SimpleExtensionDeclaration.ExtensionType.newBuilder() .setTypeAnchor(e.getKey()) .setName(e.getValue().key()) .setExtensionUrnReference(urnObj.getExtensionUrnAnchor()); - if (uriObj != null) { - typeBuilder.setExtensionUriReference(uriObj.getExtensionUriAnchor()); - } - SimpleExtensionDeclaration decl = SimpleExtensionDeclaration.newBuilder().setExtensionType(typeBuilder).build(); extensionList.add(decl); } - return new SimpleExtensions(urns, uris, extensionList); + return new SimpleExtensions(urns, extensionList); } private static final class SimpleExtensions { final HashMap urns; - final HashMap uris; final ArrayList extensionList; SimpleExtensions( HashMap urns, - HashMap uris, ArrayList extensionList) { this.urns = urns; - this.uris = uris; this.extensionList = extensionList; } } diff --git a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java index 9b546d9dc..f6bf77386 100644 --- a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java +++ b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java @@ -32,220 +32,71 @@ public static Builder builder(SimpleExtension.ExtensionCollection extensionColle public static class Builder { private final Map functionMap = new HashMap<>(); private final Map typeMap = new HashMap<>(); - private final SimpleExtension.ExtensionCollection extensionCollection; public Builder(SimpleExtension.ExtensionCollection extensionCollection) { if (extensionCollection == null) { throw new IllegalArgumentException("ExtensionCollection is required"); } - this.extensionCollection = extensionCollection; - } - - /** - * Resolves URN from URI using the URI/URN mapping. - * - * @param uri The URI to resolve - * @return The corresponding URN, or null if no mapping exists - */ - private String resolveUrnFromUri(String uri) { - return extensionCollection.getUrnFromUri(uri); } private SimpleExtension.FunctionAnchor resolveFunctionAnchor( - SimpleExtensionDeclaration.ExtensionFunction func, - Map urnMap, - Map uriMap) { - - // 1. Try non-zero URN reference - if (func.getExtensionUrnReference() != 0) { - String urnFromUrnRef = urnMap.get(func.getExtensionUrnReference()); - if (urnFromUrnRef == null) { - throw new IllegalStateException( - String.format( - "Function '%s' references URN anchor %d, but no URN is registered at that anchor", - func.getName(), func.getExtensionUrnReference())); - } - return SimpleExtension.FunctionAnchor.of(urnFromUrnRef, func.getName()); - } - - // 2. Try non-zero URI reference - if (func.getExtensionUriReference() != 0) { - String uriFromUriRef = uriMap.get(func.getExtensionUriReference()); - if (uriFromUriRef == null) { - throw new IllegalStateException( - String.format( - "Function '%s' references URI anchor %d, but no URI is registered at that anchor", - func.getName(), func.getExtensionUriReference())); - } - String urnFromUriRef = resolveUrnFromUri(uriFromUriRef); - if (urnFromUriRef == null) { - throw new IllegalStateException( - String.format( - "Function '%s' references URI anchor %d with URI '%s', but this URI could not be resolved to a URN. " - + "Ensure a URI <-> URN mapping is registered in the ExtensionCollection.", - func.getName(), func.getExtensionUriReference(), uriFromUriRef)); - } - return SimpleExtension.FunctionAnchor.of(urnFromUriRef, func.getName()); - } - - /* At this point, both URI and URN are known be 0. - * With protobufs, we cannot distinguish between 0 as an - * intentional value vs 0 as a default value. - * We perform some additional checks to below to handle this. - */ + SimpleExtensionDeclaration.ExtensionFunction func, Map urnMap) { String urn = urnMap.get(func.getExtensionUrnReference()); - String uri = uriMap.get(func.getExtensionUriReference()); - - // 3. Try both 0 URI and 0 URN if both resolve - if (uri != null && urn != null) { - String resolvedUrn = resolveUrnFromUri(uri); - if (urn.equals(resolvedUrn)) { - return SimpleExtension.FunctionAnchor.of(urn, func.getName()); - } + if (urn == null) { throw new IllegalStateException( String.format( - "Conflicting URI/URN mapping at reference 0: URI '%s' maps to URN '%s', but reference 0 also specifies URN '%s'. " - + "These must be consistent for proper resolution.", - uri, resolvedUrn, urn)); - } - - // 4. Try only 0 URN - if (urn != null) { - return SimpleExtension.FunctionAnchor.of(urn, func.getName()); - } - // 5. Try only 0 URI - if (uri != null && resolveUrnFromUri(uri) != null) { - return SimpleExtension.FunctionAnchor.of(resolveUrnFromUri(uri), func.getName()); + "Function '%s' references URN anchor %d, but no URN is registered at that anchor", + func.getName(), func.getExtensionUrnReference())); } - throw new IllegalStateException( - String.format( - "All resolution strategies failed for URI %s and URN %s (perhaps a URI <-> URN mapping was not registered during the migration) ", - uri, urn)); + return SimpleExtension.FunctionAnchor.of(urn, func.getName()); } private SimpleExtension.TypeAnchor resolveTypeAnchor( - SimpleExtensionDeclaration.ExtensionType type, - Map urnMap, - Map uriMap) { - - // 1. Try non-zero URN reference - if (type.getExtensionUrnReference() != 0) { - String urnFromUrnRef = urnMap.get(type.getExtensionUrnReference()); - if (urnFromUrnRef == null) { - throw new IllegalStateException( - String.format( - "Type '%s' references URN anchor %d, but no URN is registered at that anchor", - type.getName(), type.getExtensionUrnReference())); - } - return SimpleExtension.TypeAnchor.of(urnFromUrnRef, type.getName()); - } - - // 2. Try non-zero URI reference - if (type.getExtensionUriReference() != 0) { - String uriFromUriRef = uriMap.get(type.getExtensionUriReference()); - if (uriFromUriRef == null) { - throw new IllegalStateException( - String.format( - "Type '%s' references URI anchor %d, but no URI is registered at that anchor", - type.getName(), type.getExtensionUriReference())); - } - String urnFromUriRef = resolveUrnFromUri(uriFromUriRef); - if (urnFromUriRef == null) { - throw new IllegalStateException( - String.format( - "Type '%s' references URI anchor %d with URI '%s', but this URI could not be resolved to a URN. " - + "Ensure a URI <-> URN mapping is registered in the ExtensionCollection.", - type.getName(), type.getExtensionUriReference(), uriFromUriRef)); - } - return SimpleExtension.TypeAnchor.of(urnFromUriRef, type.getName()); - } - - /* At this point, both URI and URN are known be 0. - * With protobufs, we cannot distinguish between 0 as an - * intentional value vs 0 as a default value. - * We perform some additional checks to below to handle this. - */ + SimpleExtensionDeclaration.ExtensionType type, Map urnMap) { String urn = urnMap.get(type.getExtensionUrnReference()); - String uri = uriMap.get(type.getExtensionUriReference()); - - // 3. Try both 0 URI and 0 URN if both resolve - if (uri != null && urn != null) { - String resolvedUrn = resolveUrnFromUri(uri); - if (urn.equals(resolvedUrn)) { - return SimpleExtension.TypeAnchor.of(urn, type.getName()); - } + if (urn == null) { throw new IllegalStateException( String.format( - "Conflicting URI/URN mapping at reference 0: URI '%s' maps to URN '%s', but reference 0 also specifies URN '%s'. " - + "These must be consistent for proper resolution.", - uri, resolvedUrn, urn)); + "Type '%s' references URN anchor %d, but no URN is registered at that anchor", + type.getName(), type.getExtensionUrnReference())); } - - // 4. Try only 0 URN - if (urn != null) { - return SimpleExtension.TypeAnchor.of(urn, type.getName()); - } - // 5. Try only 0 URI - if (uri != null && resolveUrnFromUri(uri) != null) { - return SimpleExtension.TypeAnchor.of(resolveUrnFromUri(uri), type.getName()); - } - throw new IllegalStateException( - String.format( - "All resolution strategies failed for URI %s and URN %s (perhaps a URI <-> URN mapping was not registered during the migration) ", - uri, urn)); + return SimpleExtension.TypeAnchor.of(urn, type.getName()); } public Builder from(Plan plan) { - return from( - plan.getExtensionUrnsList(), plan.getExtensionUrisList(), plan.getExtensionsList()); + return from(plan.getExtensionUrnsList(), plan.getExtensionsList()); } public Builder from(ExtendedExpression extendedExpression) { return from( - extendedExpression.getExtensionUrnsList(), - extendedExpression.getExtensionUrisList(), - extendedExpression.getExtensionsList()); + extendedExpression.getExtensionUrnsList(), extendedExpression.getExtensionsList()); } private Builder from( List simpleExtensionURNs, - List simpleExtensionURIs, List simpleExtensionDeclarations) { Map urnMap = new HashMap<>(); - Map uriMap = new HashMap<>(); - // Handle URN format for (SimpleExtensionURN extension : simpleExtensionURNs) { urnMap.put(extension.getExtensionUrnAnchor(), extension.getUrn()); } - // Handle deprecated URI format - for (io.substrait.proto.SimpleExtensionURI extension : simpleExtensionURIs) { - uriMap.put(extension.getExtensionUriAnchor(), extension.getUri()); - } - - // Add all functions used in plan to the functionMap for (SimpleExtensionDeclaration extension : simpleExtensionDeclarations) { - if (!extension.hasExtensionFunction()) { - continue; + if (extension.hasExtensionFunction()) { + SimpleExtensionDeclaration.ExtensionFunction func = extension.getExtensionFunction(); + int reference = func.getFunctionAnchor(); + SimpleExtension.FunctionAnchor anchor = resolveFunctionAnchor(func, urnMap); + functionMap.put(reference, anchor); } - SimpleExtensionDeclaration.ExtensionFunction func = extension.getExtensionFunction(); - int reference = func.getFunctionAnchor(); - SimpleExtension.FunctionAnchor anchor = resolveFunctionAnchor(func, urnMap, uriMap); - functionMap.put(reference, anchor); - } - // Add all types used in plan to the typeMap - for (SimpleExtensionDeclaration extension : simpleExtensionDeclarations) { - if (!extension.hasExtensionType()) { - continue; + if (extension.hasExtensionType()) { + SimpleExtensionDeclaration.ExtensionType type = extension.getExtensionType(); + int reference = type.getTypeAnchor(); + SimpleExtension.TypeAnchor anchor = resolveTypeAnchor(type, urnMap); + typeMap.put(reference, anchor); } - SimpleExtensionDeclaration.ExtensionType type = extension.getExtensionType(); - int reference = type.getTypeAnchor(); - SimpleExtension.TypeAnchor anchor = resolveTypeAnchor(type, urnMap, uriMap); - typeMap.put(reference, anchor); } return this; diff --git a/core/src/main/java/io/substrait/extension/SimpleExtension.java b/core/src/main/java/io/substrait/extension/SimpleExtension.java index ab9058e77..556548f27 100644 --- a/core/src/main/java/io/substrait/extension/SimpleExtension.java +++ b/core/src/main/java/io/substrait/extension/SimpleExtension.java @@ -638,11 +638,6 @@ public abstract static class ExtensionCollection { Function::getAnchor, java.util.function.Function.identity())); }); - @Value.Default - BidiMap uriUrnMap() { - return new BidiMap<>(); - } - public abstract List types(); public abstract List scalarFunctions(); @@ -719,31 +714,7 @@ public WindowFunctionVariant getWindowFunction(FunctionAnchor anchor) { anchor.key(), anchor.urn())); } - /** - * Gets the URI for a given URN. This is for internal framework use during URI/URN migration. - * - * @param urn The URN to look up - * @return The corresponding URI, or null if not found - */ - String getUriFromUrn(String urn) { - return uriUrnMap().reverseGet(urn); - } - - /** - * Gets the URN for a given URI. This is for internal framework use during URI/URN migration. - * - * @param uri The URI to look up - * @return The corresponding URN, or null if not found - */ - String getUrnFromUri(String uri) { - return uriUrnMap().get(uri); - } - public ExtensionCollection merge(ExtensionCollection extensionCollection) { - BidiMap mergedUriUrnMap = new BidiMap<>(); - mergedUriUrnMap.merge(uriUrnMap()); - mergedUriUrnMap.merge(extensionCollection.uriUrnMap()); - return ImmutableSimpleExtension.ExtensionCollection.builder() .addAllAggregateFunctions(aggregateFunctions()) .addAllAggregateFunctions(extensionCollection.aggregateFunctions()) @@ -753,7 +724,6 @@ public ExtensionCollection merge(ExtensionCollection extensionCollection) { .addAllWindowFunctions(extensionCollection.windowFunctions()) .addAllTypes(types()) .addAllTypes(extensionCollection.types()) - .uriUrnMap(mergedUriUrnMap) .build(); } } @@ -783,10 +753,6 @@ public static ExtensionCollection load(List resourcePaths) { public static ExtensionCollection load(String uri, String content) { try { - if (uri == null || uri.isEmpty()) { - throw new IllegalArgumentException("URI cannot be null or empty"); - } - // Parse with basic YAML mapper first to extract URN ObjectMapper basicYamlMapper = new ObjectMapper(new YAMLFactory()); com.fasterxml.jackson.databind.JsonNode rootNode = basicYamlMapper.readTree(content); @@ -821,9 +787,6 @@ public static ExtensionCollection buildExtensionCollection( String uri, ExtensionSignatures extensionSignatures) { String urn = extensionSignatures.urn(); validateUrn(urn); - if (uri == null || uri == "") { - throw new IllegalArgumentException("URI cannot be null or empty"); - } List scalarFunctionVariants = extensionSignatures.scalars().stream() .flatMap(t -> t.resolve(urn)) @@ -856,16 +819,12 @@ public static ExtensionCollection buildExtensionCollection( Stream.concat(windowFunctionVariants, windowAggFunctionVariants) .collect(Collectors.toList()); - BidiMap 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( diff --git a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java b/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java index c2c1a58b2..20aa84f51 100644 --- a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java +++ b/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java @@ -1,8 +1,6 @@ package io.substrait.extension; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import org.junit.jupiter.api.Test; @@ -10,7 +8,7 @@ class ExtensionCollectionMergeTest { @Test - void testMergeCollectionsWithDifferentUriUrnMappings() { + void testMergeCollectionsWithDifferentUrns() { String yaml1 = "%YAML 1.2\n" + "---\n" @@ -38,16 +36,11 @@ void testMergeCollectionsWithDifferentUriUrnMappings() { SimpleExtension.ExtensionCollection merged = collection1.merge(collection2); - assertEquals("extension:ns1:collection1", merged.getUrnFromUri("uri1://extensions")); - assertEquals("extension:ns2:collection2", merged.getUrnFromUri("uri2://extensions")); - assertEquals("uri1://extensions", merged.getUriFromUrn("extension:ns1:collection1")); - assertEquals("uri2://extensions", merged.getUriFromUrn("extension:ns2:collection2")); - assertTrue(merged.scalarFunctions().size() >= 2); } @Test - void testMergeCollectionsWithIdenticalMappings() { + void testMergeCollectionsWithIdenticalUrns() { String yaml = "%YAML 1.2\n" + "---\n" @@ -64,24 +57,6 @@ void testMergeCollectionsWithIdenticalMappings() { SimpleExtension.ExtensionCollection merged = assertDoesNotThrow(() -> collection1.merge(collection2)); - assertEquals("extension:shared:extension", merged.getUrnFromUri("shared://uri")); - assertEquals("shared://uri", merged.getUriFromUrn("extension:shared:extension")); - } - - @Test - void testMergeCollectionsWithConflictingMappings() { - String yaml1 = - "%YAML 1.2\n" + "---\n" + "urn: extension:conflict:urn1\n" + "scalar_functions: []\n"; - - String yaml2 = - "%YAML 1.2\n" + "---\n" + "urn: extension:conflict:urn2\n" + "scalar_functions: []\n"; - - SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load("conflict://uri", yaml1); - SimpleExtension.ExtensionCollection collection2 = - SimpleExtension.load("conflict://uri", yaml2); // Same URI, different URN - - IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> collection1.merge(collection2)); - assertTrue(exception.getMessage().contains("Key already exists in map with different value")); + assertTrue(merged.scalarFunctions().size() >= 1); } } diff --git a/core/src/test/java/io/substrait/extension/ExtensionCollectionUriUrnTest.java b/core/src/test/java/io/substrait/extension/ExtensionCollectionUriUrnTest.java deleted file mode 100644 index 8be2d6138..000000000 --- a/core/src/test/java/io/substrait/extension/ExtensionCollectionUriUrnTest.java +++ /dev/null @@ -1,60 +0,0 @@ -package io.substrait.extension; - -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import org.junit.jupiter.api.Test; - -class ExtensionCollectionUriUrnTest { - - @Test - void testHasUrnAndHasUri() { - String yamlContent = - "%YAML 1.2\n" - + "---\n" - + "urn: extension:test:exists\n" - + "scalar_functions:\n" - + " - name: test_function\n"; - - SimpleExtension.ExtensionCollection collection = - SimpleExtension.load("file:///tmp/test.yaml", yamlContent); - - assertTrue(collection.getUrnFromUri("file:///tmp/test.yaml") != null); - assertTrue(collection.getUriFromUrn("extension:test:exists") != null); - assertFalse(collection.getUrnFromUri("nonexistent://uri") != null); - assertFalse(collection.getUriFromUrn("extension:nonexistent:urn") != null); - } - - @Test - void testGetNonexistentMappings() { - String yamlContent = - "%YAML 1.2\n" + "---\n" + "urn: extension:test:minimal\n" + "scalar_functions: []\n"; - - SimpleExtension.ExtensionCollection collection = - SimpleExtension.load("minimal://extension", yamlContent); - - assertNull(collection.getUrnFromUri("nonexistent://uri")); - assertNull(collection.getUriFromUrn("extension:nonexistent:urn")); - } - - @Test - void testEmptyUriThrowsException() { - String yamlContent = - "%YAML 1.2\n" + "---\n" + "urn: extension:test:empty\n" + "scalar_functions: []\n"; - - IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load("", yamlContent)); - assertTrue(exception.getMessage().contains("URI cannot be null or empty")); - } - - @Test - void testNullUriThrowsException() { - String yamlContent = - "%YAML 1.2\n" + "---\n" + "urn: extension:test:null\n" + "scalar_functions: []\n"; - - // The system throws NPE when null is passed, which is expected behavior - assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load(null, yamlContent)); - } -} diff --git a/core/src/test/java/io/substrait/extension/ExtensionCollectorUriUrnTest.java b/core/src/test/java/io/substrait/extension/ExtensionCollectorUriUrnTest.java deleted file mode 100644 index abfec1310..000000000 --- a/core/src/test/java/io/substrait/extension/ExtensionCollectorUriUrnTest.java +++ /dev/null @@ -1,41 +0,0 @@ -package io.substrait.extension; - -import static org.junit.jupiter.api.Assertions.assertEquals; - -import io.substrait.proto.Plan; -import org.junit.jupiter.api.Test; - -class ExtensionCollectorUriUrnTest { - - @Test - void testExtensionCollectorScalarFuncWithoutURI() { - String uri = "test://uri"; - BidiMap uriUrnMap = new BidiMap(); - uriUrnMap.put(uri, "extension:test:basic"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - ExtensionCollector collector = new ExtensionCollector(extensionCollection); - - SimpleExtension.ScalarFunctionVariant func = - ImmutableSimpleExtension.ScalarFunctionVariant.builder() - .urn("extension:test:basic") - .name("test_func") - .returnType(io.substrait.function.TypeExpressionCreator.REQUIRED.BOOLEAN) - .build(); - - int functionRef = collector.getFunctionReference(func); - assertEquals(1, functionRef); - - Plan.Builder planBuilder = Plan.newBuilder(); - collector.addExtensionsToPlan(planBuilder); - - Plan plan = planBuilder.build(); - assertEquals(1, plan.getExtensionUrnsCount()); - assertEquals("extension:test:basic", plan.getExtensionUrns(0).getUrn()); - - assertEquals(1, plan.getExtensionUrisCount()); - assertEquals("test://uri", plan.getExtensionUris(0).getUri()); - } -} diff --git a/core/src/test/java/io/substrait/extension/ImmutableExtensionLookupUriUrnTest.java b/core/src/test/java/io/substrait/extension/ImmutableExtensionLookupUriUrnTest.java deleted file mode 100644 index 71cc9a337..000000000 --- a/core/src/test/java/io/substrait/extension/ImmutableExtensionLookupUriUrnTest.java +++ /dev/null @@ -1,621 +0,0 @@ -package io.substrait.extension; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import io.substrait.proto.Plan; -import io.substrait.proto.SimpleExtensionDeclaration; -import io.substrait.proto.SimpleExtensionURI; -import io.substrait.proto.SimpleExtensionURN; -import org.junit.jupiter.api.Test; - -class ImmutableExtensionLookupUriUrnTest { - - @Test - void testUrnResolutionWorks() { - // Create URN-only plan (normal case) - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(1) - .setUrn("extension:test:urn") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("test_func") - .setExtensionUrnReference(1) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUrns(urnProto).addExtensions(decl).build(); - - // Test with no ExtensionCollection (no URI/URN mapping available) - ImmutableExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); - - assertEquals("extension:test:urn", lookup.functionAnchorMap.get(1).urn()); - assertEquals("test_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testUriToUrnFallbackWorks() { - // Create an ExtensionCollection with URI/URN mapping - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/extensions/test", "extension:test:mapped"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - // Create URI-only plan (legacy case) - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(1) - .setUri("http://example.com/extensions/test") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("legacy_func") - .setExtensionUriReference(1) // References the URI anchor (deprecated field) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - // Test with URI/URN mapping - should resolve URI to URN - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:mapped", lookup.functionAnchorMap.get(1).urn()); - assertEquals("legacy_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testUriWithoutMappingThrowsError() { - // Create URI-only plan without mapping - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(1) - .setUri("http://example.com/unmapped") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("unmapped_func") - .setExtensionUriReference(1) // References the URI anchor - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - // Should throw error - URI present but no mapping available - IllegalStateException exception = - assertThrows( - IllegalStateException.class, - () -> { - ImmutableExtensionLookup.builder().from(plan).build(); - }); - - assertTrue(exception.getMessage().contains("could not be resolved to a URN")); - assertTrue(exception.getMessage().contains("http://example.com/unmapped")); - assertTrue(exception.getMessage().contains("URI <-> URN mapping")); - } - - @Test - void testMissingUrnAndUriThrowsError() { - // Create plan with missing URN/URI reference - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("missing_func") - .setExtensionUrnReference(999) // Non-existent reference - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensions(decl).build(); - - // Should throw error - neither URN nor URI found - IllegalStateException exception = - assertThrows( - IllegalStateException.class, - () -> { - ImmutableExtensionLookup.builder().from(plan).build(); - }); - - assertTrue(exception.getMessage().contains("no URN is registered at that anchor")); - assertTrue(exception.getMessage().contains("999")); // The missing anchor reference - } - - // ========================================================================== - // Simple tests for all 5 resolution cases - Functions - // ========================================================================== - - @Test - void testFunctionCase1_NonZeroUrnReference() { - // Case 1: Non-zero URN reference resolves - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(1) - .setUrn("extension:test:case1") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("case1_func") - .setExtensionUrnReference(1) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUrns(urnProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); - - assertEquals("extension:test:case1", lookup.functionAnchorMap.get(1).urn()); - assertEquals("case1_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testFunctionCase2_NonZeroUriReference() { - // Case 2: Non-zero URI reference resolves via mapping - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case2", "extension:test:case2"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(1) - .setUri("http://example.com/case2") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("case2_func") - .setExtensionUriReference(1) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case2", lookup.functionAnchorMap.get(1).urn()); - assertEquals("case2_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testFunctionCase3_ZeroBothResolveConsistent() { - // Case 3: Both 0 references resolve to consistent URN - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case3", "extension:test:case3"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:case3") - .build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/case3") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("case3_func") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = - Plan.newBuilder() - .addExtensionUrns(urnProto) - .addExtensionUris(uriProto) - .addExtensions(decl) - .build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case3", lookup.functionAnchorMap.get(1).urn()); - assertEquals("case3_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testFunctionCase3_ZeroBothResolveConflict() { - // Case 3: Both 0 references resolve but to different URNs - should throw - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/conflict", "extension:test:different"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:original") - .build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/conflict") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("conflict_func") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = - Plan.newBuilder() - .addExtensionUrns(urnProto) - .addExtensionUris(uriProto) - .addExtensions(decl) - .build(); - - IllegalStateException exception = - assertThrows( - IllegalStateException.class, - () -> { - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - }); - - assertTrue(exception.getMessage().contains("Conflicting URI/URN mapping")); - assertTrue(exception.getMessage().contains("These must be consistent")); - } - - @Test - void testFunctionCase4_ZeroUrnOnly() { - // Case 4: Only 0 URN reference resolves - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:case4") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("case4_func") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUrns(urnProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); - - assertEquals("extension:test:case4", lookup.functionAnchorMap.get(1).urn()); - assertEquals("case4_func", lookup.functionAnchorMap.get(1).key()); - } - - @Test - void testFunctionCase5_ZeroUriOnly() { - // Case 5: Only 0 URI reference resolves - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case5", "extension:test:case5"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/case5") - .build(); - - SimpleExtensionDeclaration.ExtensionFunction func = - SimpleExtensionDeclaration.ExtensionFunction.newBuilder() - .setFunctionAnchor(1) - .setName("case5_func") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionFunction(func).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case5", lookup.functionAnchorMap.get(1).urn()); - assertEquals("case5_func", lookup.functionAnchorMap.get(1).key()); - } - - // ========================================================================== - // Simple tests for all 5 resolution cases - Types - // ========================================================================== - - @Test - void testTypeCase1_NonZeroUrnReference() { - // Case 1: Non-zero URN reference resolves - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(1) - .setUrn("extension:test:case1") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("case1_type") - .setExtensionUrnReference(1) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = Plan.newBuilder().addExtensionUrns(urnProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); - - assertEquals("extension:test:case1", lookup.typeAnchorMap.get(1).urn()); - assertEquals("case1_type", lookup.typeAnchorMap.get(1).key()); - } - - @Test - void testTypeCase2_NonZeroUriReference() { - // Case 2: Non-zero URI reference resolves via mapping - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case2", "extension:test:case2"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(1) - .setUri("http://example.com/case2") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("case2_type") - .setExtensionUriReference(1) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case2", lookup.typeAnchorMap.get(1).urn()); - assertEquals("case2_type", lookup.typeAnchorMap.get(1).key()); - } - - @Test - void testTypeCase3_ZeroBothResolveConsistent() { - // Case 3: Both 0 references resolve to consistent URN - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case3", "extension:test:case3"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:case3") - .build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/case3") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("case3_type") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = - Plan.newBuilder() - .addExtensionUrns(urnProto) - .addExtensionUris(uriProto) - .addExtensions(decl) - .build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case3", lookup.typeAnchorMap.get(1).urn()); - assertEquals("case3_type", lookup.typeAnchorMap.get(1).key()); - } - - @Test - void testTypeCase3_ZeroBothResolveConflict() { - // Case 3: Both 0 references resolve but to different URNs - should throw - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/conflict", "extension:test:different"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:original") - .build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/conflict") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("conflict_type") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = - Plan.newBuilder() - .addExtensionUrns(urnProto) - .addExtensionUris(uriProto) - .addExtensions(decl) - .build(); - - IllegalStateException exception = - assertThrows( - IllegalStateException.class, - () -> { - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - }); - - assertTrue(exception.getMessage().contains("Conflicting URI/URN mapping")); - assertTrue(exception.getMessage().contains("These must be consistent")); - } - - @Test - void testTypeCase4_ZeroUrnOnly() { - // Case 4: Only 0 URN reference resolves - SimpleExtensionURN urnProto = - SimpleExtensionURN.newBuilder() - .setExtensionUrnAnchor(0) - .setUrn("extension:test:case4") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("case4_type") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = Plan.newBuilder().addExtensionUrns(urnProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); - - assertEquals("extension:test:case4", lookup.typeAnchorMap.get(1).urn()); - assertEquals("case4_type", lookup.typeAnchorMap.get(1).key()); - } - - @Test - void testTypeCase5_ZeroUriOnly() { - // Case 5: Only 0 URI reference resolves - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/case5", "extension:test:case5"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(0) - .setUri("http://example.com/case5") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("case5_type") - .setExtensionUrnReference(0) - .setExtensionUriReference(0) - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:test:case5", lookup.typeAnchorMap.get(1).urn()); - assertEquals("case5_type", lookup.typeAnchorMap.get(1).key()); - } - - @Test - void testTypeUriToUrnFallbackWorks() { - // Test the same logic but for types instead of functions - BidiMap uriUrnMap = new BidiMap<>(); - uriUrnMap.put("http://example.com/types/test", "extension:types:mapped"); - - SimpleExtension.ExtensionCollection extensionCollection = - SimpleExtension.ExtensionCollection.builder().uriUrnMap(uriUrnMap).build(); - - SimpleExtensionURI uriProto = - SimpleExtensionURI.newBuilder() - .setExtensionUriAnchor(1) - .setUri("http://example.com/types/test") - .build(); - - SimpleExtensionDeclaration.ExtensionType type = - SimpleExtensionDeclaration.ExtensionType.newBuilder() - .setTypeAnchor(1) - .setName("legacy_type") - .setExtensionUriReference(1) // References the URI anchor - .build(); - - SimpleExtensionDeclaration decl = - SimpleExtensionDeclaration.newBuilder().setExtensionType(type).build(); - - Plan plan = Plan.newBuilder().addExtensionUris(uriProto).addExtensions(decl).build(); - - ImmutableExtensionLookup lookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); - - assertEquals("extension:types:mapped", lookup.typeAnchorMap.get(1).urn()); - assertEquals("legacy_type", lookup.typeAnchorMap.get(1).key()); - } -} diff --git a/core/src/test/java/io/substrait/extension/UriUrnMigrationEndToEndTest.java b/core/src/test/java/io/substrait/extension/UriUrnMigrationEndToEndTest.java deleted file mode 100644 index e96b8133f..000000000 --- a/core/src/test/java/io/substrait/extension/UriUrnMigrationEndToEndTest.java +++ /dev/null @@ -1,110 +0,0 @@ -package io.substrait.extension; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import com.google.protobuf.util.JsonFormat; -import io.substrait.plan.PlanProtoConverter; -import io.substrait.plan.ProtoPlanConverter; -import io.substrait.proto.Plan; -import java.io.BufferedReader; -import java.io.IOException; -import java.io.InputStream; -import java.io.InputStreamReader; -import java.nio.charset.StandardCharsets; -import java.util.Arrays; -import java.util.List; -import java.util.stream.Collectors; -import org.junit.jupiter.api.Test; - -/** - * End-to-end tests demonstrating the full URI/URN migration workflow: 1. Consume plans with mixed - * URI/URN references 2. Convert proto -> POJO using ImmutableExtensionLookup with URI/URN mapping - * 3. Convert POJO -> proto using PlanProtoConverter 4. Verify output contains proper - * extensioninformation - */ -class UriUrnMigrationEndToEndTest { - - /** Load a proto Plan from a JSON resource file using JsonFormat */ - private Plan loadPlanFromJson(String resourcePath) throws IOException { - try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourcePath)) { - if (inputStream == null) { - throw new IOException("Resource not found: " + resourcePath); - } - - String jsonContent = - new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)) - .lines() - .collect(Collectors.joining("\n")); - - Plan.Builder planBuilder = Plan.newBuilder(); - JsonFormat.parser().merge(jsonContent, planBuilder); - return planBuilder.build(); - } - } - - @Test - void testUriUrnMigrationEndToEnd() throws IOException { - - // List of (inputPath, expectedPath, extensionCollection) tuples - List testCases = - Arrays.asList( - new String[] { - "uri-urn-migration/uri-only-input-plan.json", - "uri-urn-migration/uri-only-expected-plan.json" - }, - new String[] { - "uri-urn-migration/complex-input-plan.json", - "uri-urn-migration/complex-expected-plan.json" - }, - new String[] { - "uri-urn-migration/urn-only-input-plan.json", - "uri-urn-migration/urn-only-expected-plan.json" - }, - new String[] { - "uri-urn-migration/mixed-partial-coverage-input-plan.json", - "uri-urn-migration/mixed-partial-coverage-expected-plan.json" - }, - new String[] { - "uri-urn-migration/zero-urn-resolution-input-plan.json", - "uri-urn-migration/zero-urn-resolution-expected-plan.json" - }); - - for (String[] testCase : testCases) { - String inputPath = testCase[0]; - String expectedPath = testCase[1]; - - Plan inputPlan = loadPlanFromJson(inputPath); - Plan expectedPlan = loadPlanFromJson(expectedPath); - - ProtoPlanConverter protoToPojo = - new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); - io.substrait.plan.Plan pojoPlan = protoToPojo.from(inputPlan); - - PlanProtoConverter pojoToProto = - new PlanProtoConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); - Plan actualPlan = pojoToProto.toProto(pojoPlan); - - assertEquals(expectedPlan, actualPlan); - } - } - - @Test - void testUnresolvableUriThrowsException() throws IOException { - Plan inputPlan = loadPlanFromJson("uri-urn-migration/unresolvable-uri-plan.json"); - - ProtoPlanConverter protoToPojo = - new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); - - IllegalStateException exception = - assertThrows( - IllegalStateException.class, - () -> { - protoToPojo.from(inputPlan); - }); - - assertTrue(exception.getMessage().contains("could not be resolved to a URN")); - assertTrue(exception.getMessage().contains("/functions_nonexistent.yaml")); - } -} diff --git a/core/src/test/java/io/substrait/extension/UrnValidationTest.java b/core/src/test/java/io/substrait/extension/UrnValidationTest.java index 55541c2d5..44bac5093 100644 --- a/core/src/test/java/io/substrait/extension/UrnValidationTest.java +++ b/core/src/test/java/io/substrait/extension/UrnValidationTest.java @@ -1,7 +1,6 @@ package io.substrait.extension; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -44,17 +43,4 @@ void testValidUrnWorks() { + " - name: test\n"; assertDoesNotThrow(() -> SimpleExtension.load("some/uri", yamlWithValidUrn)); } - - @Test - 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.getUrnFromUri("test://uri")); - } } diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java new file mode 100644 index 000000000..cf8a03ec9 --- /dev/null +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -0,0 +1,74 @@ +package io.substrait.plan; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.google.protobuf.util.JsonFormat; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.proto.Plan; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Roundtrip tests: parse a JSON proto plan, convert to POJO, convert back to proto, and compare + * with the expected output. + */ +class PlanRoundtripTest { + + private Plan loadPlanFromJson(String resourcePath) throws IOException { + try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourcePath)) { + if (inputStream == null) { + throw new IOException("Resource not found: " + resourcePath); + } + + String jsonContent = + new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)) + .lines() + .collect(Collectors.joining("\n")); + + Plan.Builder planBuilder = Plan.newBuilder(); + JsonFormat.parser().merge(jsonContent, planBuilder); + return planBuilder.build(); + } + } + + static Stream roundtripCases() { + return Stream.of( + Arguments.of( + "simple", + "plan-roundtrip/simple-input-plan.json", + "plan-roundtrip/simple-expected-plan.json"), + Arguments.of( + "complex", + "plan-roundtrip/complex-input-plan.json", + "plan-roundtrip/complex-expected-plan.json"), + Arguments.of( + "zero-anchor", + "plan-roundtrip/zero-anchor-input-plan.json", + "plan-roundtrip/zero-anchor-expected-plan.json")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("roundtripCases") + void testPlanRoundtrip(String name, String inputPath, String expectedPath) throws IOException { + Plan inputPlan = loadPlanFromJson(inputPath); + Plan expectedPlan = loadPlanFromJson(expectedPath); + + ProtoPlanConverter protoToPojo = + new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); + io.substrait.plan.Plan pojoPlan = protoToPojo.from(inputPlan); + + PlanProtoConverter pojoToProto = + new PlanProtoConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); + Plan actualPlan = pojoToProto.toProto(pojoPlan); + + assertEquals(expectedPlan, actualPlan); + } +} diff --git a/core/src/test/resources/uri-urn-migration/complex-expected-plan.json b/core/src/test/resources/plan-roundtrip/complex-expected-plan.json similarity index 94% rename from core/src/test/resources/uri-urn-migration/complex-expected-plan.json rename to core/src/test/resources/plan-roundtrip/complex-expected-plan.json index 09833ace9..9f4c47a18 100644 --- a/core/src/test/resources/uri-urn-migration/complex-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/complex-expected-plan.json @@ -1,10 +1,4 @@ { - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], "extensionUrns": [ { "extensionUrnAnchor": 1, @@ -16,7 +10,6 @@ "extensionFunction": { "functionAnchor": 1, "name": "add:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } }, @@ -24,7 +17,6 @@ "extensionFunction": { "functionAnchor": 2, "name": "subtract:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } }, @@ -32,7 +24,6 @@ "extensionFunction": { "functionAnchor": 3, "name": "multiply:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } } diff --git a/core/src/test/resources/uri-urn-migration/complex-input-plan.json b/core/src/test/resources/plan-roundtrip/complex-input-plan.json similarity index 90% rename from core/src/test/resources/uri-urn-migration/complex-input-plan.json rename to core/src/test/resources/plan-roundtrip/complex-input-plan.json index da28cbb33..17a1984bd 100644 --- a/core/src/test/resources/uri-urn-migration/complex-input-plan.json +++ b/core/src/test/resources/plan-roundtrip/complex-input-plan.json @@ -1,17 +1,7 @@ { - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - }, - { - "extensionUriAnchor": 0, - "uri": "/functions_string.yaml" - } - ], "extensionUrns": [ { - "extensionUrnAnchor": 2, + "extensionUrnAnchor": 1, "urn": "extension:io.substrait:functions_arithmetic" } ], @@ -20,21 +10,21 @@ "extensionFunction": { "functionAnchor": 1, "name": "add:i32_i32", - "extensionUriReference": 1 + "extensionUrnReference": 1 } }, { "extensionFunction": { "functionAnchor": 2, "name": "subtract:i32_i32", - "extensionUrnReference": 2 + "extensionUrnReference": 1 } }, { "extensionFunction": { "functionAnchor": 3, "name": "multiply:i32_i32", - "extensionUriReference": 1 + "extensionUrnReference": 1 } } ], diff --git a/core/src/test/resources/uri-urn-migration/urn-only-expected-plan.json b/core/src/test/resources/plan-roundtrip/simple-expected-plan.json similarity index 94% rename from core/src/test/resources/uri-urn-migration/urn-only-expected-plan.json rename to core/src/test/resources/plan-roundtrip/simple-expected-plan.json index 9e47d1e74..5ae8daa3f 100644 --- a/core/src/test/resources/uri-urn-migration/urn-only-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/simple-expected-plan.json @@ -1,10 +1,4 @@ { - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], "extensionUrns": [ { "extensionUrnAnchor": 1, @@ -16,7 +10,6 @@ "extensionFunction": { "functionAnchor": 1, "name": "add:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } }, @@ -24,7 +17,6 @@ "extensionFunction": { "functionAnchor": 2, "name": "subtract:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } } diff --git a/core/src/test/resources/uri-urn-migration/urn-only-input-plan.json b/core/src/test/resources/plan-roundtrip/simple-input-plan.json similarity index 100% rename from core/src/test/resources/uri-urn-migration/urn-only-input-plan.json rename to core/src/test/resources/plan-roundtrip/simple-input-plan.json diff --git a/core/src/test/resources/uri-urn-migration/zero-urn-resolution-expected-plan.json b/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json similarity index 94% rename from core/src/test/resources/uri-urn-migration/zero-urn-resolution-expected-plan.json rename to core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json index b8788c850..ff6b8ab0e 100644 --- a/core/src/test/resources/uri-urn-migration/zero-urn-resolution-expected-plan.json +++ b/core/src/test/resources/plan-roundtrip/zero-anchor-expected-plan.json @@ -1,10 +1,4 @@ { - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], "extensionUrns": [ { "extensionUrnAnchor": 1, @@ -16,7 +10,6 @@ "extensionFunction": { "functionAnchor": 1, "name": "add:i32_i32", - "extensionUriReference": 1, "extensionUrnReference": 1 } } diff --git a/core/src/test/resources/uri-urn-migration/zero-urn-resolution-input-plan.json b/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json similarity index 98% rename from core/src/test/resources/uri-urn-migration/zero-urn-resolution-input-plan.json rename to core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json index bc4983846..202814110 100644 --- a/core/src/test/resources/uri-urn-migration/zero-urn-resolution-input-plan.json +++ b/core/src/test/resources/plan-roundtrip/zero-anchor-input-plan.json @@ -10,7 +10,6 @@ "extensionFunction": { "functionAnchor": 1, "name": "add:i32_i32", - "extensionUriReference": 0, "extensionUrnReference": 0 } } diff --git a/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-expected-plan.json b/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-expected-plan.json deleted file mode 100644 index 2898243b3..000000000 --- a/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-expected-plan.json +++ /dev/null @@ -1,133 +0,0 @@ -{ - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], - "extensionUrns": [ - { - "extensionUrnAnchor": 1, - "urn": "extension:io.substrait:functions_arithmetic" - } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "add:i32_i32", - "extensionUriReference": 1, - "extensionUrnReference": 1 - } - }, - { - "extensionFunction": { - "functionAnchor": 2, - "name": "multiply:i32_i32", - "extensionUriReference": 1, - "extensionUrnReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "project": { - "common": { - "direct": {} - }, - "input": { - "read": { - "common": { - "direct": {} - }, - "baseSchema": { - "names": [ - "x", - "y" - ], - "struct": { - "types": [ - { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - }, - { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - } - ], - "nullability": "NULLABILITY_REQUIRED" - } - }, - "namedTable": { - "names": [ - "test_table" - ] - } - } - }, - "expressions": [ - { - "scalarFunction": { - "functionReference": 1, - "outputType": { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 10 - } - } - }, - { - "value": { - "scalarFunction": { - "functionReference": 2, - "outputType": { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 2 - } - } - }, - { - "value": { - "literal": { - "i32": 3 - } - } - } - ] - } - } - } - ] - } - } - ] - } - }, - "names": [ - "calculation" - ] - } - } - ], - "version": { - "minorNumber": 75 - } -} diff --git a/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-input-plan.json b/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-input-plan.json deleted file mode 100644 index e02019925..000000000 --- a/core/src/test/resources/uri-urn-migration/mixed-partial-coverage-input-plan.json +++ /dev/null @@ -1,110 +0,0 @@ -{ - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "add:i32_i32", - "extensionUriReference": 1 - } - }, - { - "extensionFunction": { - "functionAnchor": 2, - "name": "multiply:i32_i32", - "extensionUriReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "project": { - "input": { - "read": { - "baseSchema": { - "names": [ - "x", - "y" - ], - "struct": { - "types": [ - { - "i32": {} - }, - { - "i32": {} - } - ] - } - }, - "namedTable": { - "names": [ - "test_table" - ] - } - } - }, - "expressions": [ - { - "scalarFunction": { - "functionReference": 1, - "outputType": { - "i32": {} - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 10 - } - } - }, - { - "value": { - "scalarFunction": { - "functionReference": 2, - "outputType": { - "i32": {} - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 2 - } - } - }, - { - "value": { - "literal": { - "i32": 3 - } - } - } - ] - } - } - } - ] - } - } - ] - } - }, - "names": [ - "calculation" - ] - } - } - ], - "version": { - "minorNumber": 75 - } -} diff --git a/core/src/test/resources/uri-urn-migration/unresolvable-uri-plan.json b/core/src/test/resources/uri-urn-migration/unresolvable-uri-plan.json deleted file mode 100644 index 769d359e7..000000000 --- a/core/src/test/resources/uri-urn-migration/unresolvable-uri-plan.json +++ /dev/null @@ -1,73 +0,0 @@ -{ - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_nonexistent.yaml" - } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "nonexistent_function:i32_i32", - "extensionUriReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "project": { - "input": { - "read": { - "baseSchema": { - "names": [ - "dummy" - ], - "struct": { - "types": [ - { - "i32": {} - } - ] - } - }, - "namedTable": { - "names": [ - "test_table" - ] - } - } - }, - "expressions": [ - { - "scalarFunction": { - "functionReference": 1, - "outputType": { - "i32": {} - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 7 - } - } - } - ] - } - } - ] - } - }, - "names": [ - "result" - ] - } - } - ], - "version": { - "minorNumber": 75 - } -} diff --git a/core/src/test/resources/uri-urn-migration/uri-only-expected-plan.json b/core/src/test/resources/uri-urn-migration/uri-only-expected-plan.json deleted file mode 100644 index 1ad0150e3..000000000 --- a/core/src/test/resources/uri-urn-migration/uri-only-expected-plan.json +++ /dev/null @@ -1,98 +0,0 @@ -{ - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], - "extensionUrns": [ - { - "extensionUrnAnchor": 1, - "urn": "extension:io.substrait:functions_arithmetic" - } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "add:i32_i32", - "extensionUriReference": 1, - "extensionUrnReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "project": { - "common": { - "direct": {} - }, - "input": { - "read": { - "common": { - "direct": {} - }, - "baseSchema": { - "names": [ - "dummy" - ], - "struct": { - "types": [ - { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - } - ], - "nullability": "NULLABILITY_REQUIRED" - } - }, - "namedTable": { - "names": [ - "test_table" - ] - } - } - }, - "expressions": [ - { - "scalarFunction": { - "functionReference": 1, - "outputType": { - "i32": { - "nullability": "NULLABILITY_REQUIRED" - } - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 7 - } - } - }, - { - "value": { - "literal": { - "i32": 3 - } - } - } - ] - } - } - ] - } - }, - "names": [ - "result" - ] - } - } - ], - "version": { - "minorNumber": 75 - } -} diff --git a/core/src/test/resources/uri-urn-migration/uri-only-input-plan.json b/core/src/test/resources/uri-urn-migration/uri-only-input-plan.json deleted file mode 100644 index 5630a00ca..000000000 --- a/core/src/test/resources/uri-urn-migration/uri-only-input-plan.json +++ /dev/null @@ -1,80 +0,0 @@ -{ - "extensionUris": [ - { - "extensionUriAnchor": 1, - "uri": "/functions_arithmetic.yaml" - } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "add:i32_i32", - "extensionUriReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "project": { - "input": { - "read": { - "baseSchema": { - "names": [ - "dummy" - ], - "struct": { - "types": [ - { - "i32": {} - } - ] - } - }, - "namedTable": { - "names": [ - "test_table" - ] - } - } - }, - "expressions": [ - { - "scalarFunction": { - "functionReference": 1, - "outputType": { - "i32": {} - }, - "arguments": [ - { - "value": { - "literal": { - "i32": 7 - } - } - }, - { - "value": { - "literal": { - "i32": 3 - } - } - } - ] - } - } - ] - } - }, - "names": [ - "result" - ] - } - } - ], - "version": { - "minorNumber": 75 - } -} diff --git a/substrait b/substrait index e4ce3f871..2aaae7c0d 160000 --- a/substrait +++ b/substrait @@ -1 +1 @@ -Subproject commit e4ce3f8710d91cb95cce010c6b59eb1d618bb2a6 +Subproject commit 2aaae7c0d369a23b5b469badce0644097f1769a1 From 080b97a1ba7621148b60f74fd5764e9886aedd63 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:00:01 -0400 Subject: [PATCH 02/13] refactor: remove uri parameter from SimpleExtension.load() --- .../ExtendedExpressionProtoConverter.java | 3 +-- .../substrait/extension/SimpleExtension.java | 18 +++++++----------- .../ExtensionCollectionMergeTest.java | 10 ++++------ .../substrait/extension/TypeExtensionTest.java | 2 +- .../substrait/extension/UrnValidationTest.java | 8 +++----- .../extension/VariadicBehaviorTest.java | 6 ++---- .../io/substrait/plan/PlanConverterTest.java | 2 +- .../type/proto/LiteralRoundtripTest.java | 2 +- .../substrait/isthmus/CustomFunctionTest.java | 2 +- .../isthmus/DuplicateFunctionUrnTest.java | 5 ++--- .../SimpleExtensionToSqlOperatorTest.java | 1 - .../UserDefinedLiteralRoundtripTest.java | 2 +- .../io/substrait/spark/SparkExtension.scala | 2 +- 13 files changed, 25 insertions(+), 38 deletions(-) diff --git a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java index e0cb9d1c5..d38ac04b1 100644 --- a/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java +++ b/core/src/main/java/io/substrait/extendedexpression/ExtendedExpressionProtoConverter.java @@ -58,8 +58,7 @@ public ExtendedExpression toProto( builder.setBaseSchema( extendedExpression.getBaseSchema().toProto(new TypeProtoConverter(functionCollector))); - // the process of adding simple extensions, such as extensionURIs and extensions, is handled on - // the fly + // the process of adding simple extensions (URNs and declarations) is handled on the fly functionCollector.addExtensionsToExtendedExpression(builder); if (extendedExpression.getAdvancedExtension().isPresent()) { builder.setAdvancedExtensions(extendedExpression.getAdvancedExtension().get()); diff --git a/core/src/main/java/io/substrait/extension/SimpleExtension.java b/core/src/main/java/io/substrait/extension/SimpleExtension.java index 556548f27..1af82526d 100644 --- a/core/src/main/java/io/substrait/extension/SimpleExtension.java +++ b/core/src/main/java/io/substrait/extension/SimpleExtension.java @@ -738,7 +738,7 @@ public static ExtensionCollection load(List resourcePaths) { .map( path -> { try (InputStream stream = ExtensionCollection.class.getResourceAsStream(path)) { - return load(path, stream); + return load(stream); } catch (IOException e) { throw new UncheckedIOException(e); } @@ -751,7 +751,7 @@ public static ExtensionCollection load(List resourcePaths) { return complete; } - public static ExtensionCollection load(String uri, String content) { + public static ExtensionCollection load(String content) { try { // Parse with basic YAML mapper first to extract URN ObjectMapper basicYamlMapper = new ObjectMapper(new YAMLFactory()); @@ -763,28 +763,24 @@ public static ExtensionCollection load(String uri, String content) { String urn = urnNode.asText(); validateUrn(urn); - ExtensionSignatures docWithoutUri = - objectMapper(urn).readValue(content, ExtensionSignatures.class); + ExtensionSignatures doc = objectMapper(urn).readValue(content, ExtensionSignatures.class); - ExtensionSignatures doc = - ImmutableSimpleExtension.ExtensionSignatures.builder().from(docWithoutUri).build(); - - return buildExtensionCollection(uri, doc); + return buildExtensionCollection(doc); } catch (IOException e) { throw new IllegalStateException(e); } } - public static ExtensionCollection load(String uri, InputStream stream) { + public static ExtensionCollection load(InputStream stream) { try (Scanner scanner = new Scanner(stream)) { scanner.useDelimiter(READ_WHOLE_FILE); String content = scanner.next(); - return load(uri, content); + return load(content); } } public static ExtensionCollection buildExtensionCollection( - String uri, ExtensionSignatures extensionSignatures) { + ExtensionSignatures extensionSignatures) { String urn = extensionSignatures.urn(); validateUrn(urn); List scalarFunctionVariants = diff --git a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java b/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java index 20aa84f51..dfbbd571c 100644 --- a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java +++ b/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java @@ -29,10 +29,8 @@ void testMergeCollectionsWithDifferentUrns() { + " - args: []\n" + " return: i32\n"; - SimpleExtension.ExtensionCollection collection1 = - SimpleExtension.load("uri1://extensions", yaml1); - SimpleExtension.ExtensionCollection collection2 = - SimpleExtension.load("uri2://extensions", yaml2); + SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load(yaml1); + SimpleExtension.ExtensionCollection collection2 = SimpleExtension.load(yaml2); SimpleExtension.ExtensionCollection merged = collection1.merge(collection2); @@ -51,8 +49,8 @@ void testMergeCollectionsWithIdenticalUrns() { + " - args: []\n" + " return: boolean\n"; - SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load("shared://uri", yaml); - SimpleExtension.ExtensionCollection collection2 = SimpleExtension.load("shared://uri", yaml); + SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load(yaml); + SimpleExtension.ExtensionCollection collection2 = SimpleExtension.load(yaml); SimpleExtension.ExtensionCollection merged = assertDoesNotThrow(() -> collection1.merge(collection2)); diff --git a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java index 4b3b94a80..3a3c2ec63 100644 --- a/core/src/test/java/io/substrait/extension/TypeExtensionTest.java +++ b/core/src/test/java/io/substrait/extension/TypeExtensionTest.java @@ -32,7 +32,7 @@ class TypeExtensionTest extends TestBase { static { try { String customExtensionStr = asString("extensions/custom_extensions.yaml"); - CUSTOM_EXTENSION = SimpleExtension.load(URN, customExtensionStr); + CUSTOM_EXTENSION = SimpleExtension.load(customExtensionStr); } catch (IOException e) { throw new UncheckedIOException(e); } diff --git a/core/src/test/java/io/substrait/extension/UrnValidationTest.java b/core/src/test/java/io/substrait/extension/UrnValidationTest.java index 44bac5093..a8abc58be 100644 --- a/core/src/test/java/io/substrait/extension/UrnValidationTest.java +++ b/core/src/test/java/io/substrait/extension/UrnValidationTest.java @@ -12,8 +12,7 @@ class UrnValidationTest { void testMissingUrnThrowsException() { String yamlWithoutUrn = "%YAML 1.2\n" + "---\n" + "scalar_functions:\n" + " - name: test\n"; IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, () -> SimpleExtension.load("some/uri", yamlWithoutUrn)); + assertThrows(IllegalArgumentException.class, () -> SimpleExtension.load(yamlWithoutUrn)); assertTrue(exception.getMessage().contains("Extension YAML file must contain a 'urn' field")); } @@ -27,8 +26,7 @@ void testInvalidUrnFormatThrowsException() { + " - name: test\n"; IllegalArgumentException exception = assertThrows( - IllegalArgumentException.class, - () -> SimpleExtension.load("some/uri", yamlWithInvalidUrn)); + IllegalArgumentException.class, () -> SimpleExtension.load(yamlWithInvalidUrn)); assertTrue( exception.getMessage().contains("URN must follow format 'extension::'")); } @@ -41,6 +39,6 @@ void testValidUrnWorks() { + "urn: extension:test:valid\n" + "scalar_functions:\n" + " - name: test\n"; - assertDoesNotThrow(() -> SimpleExtension.load("some/uri", yamlWithValidUrn)); + assertDoesNotThrow(() -> SimpleExtension.load(yamlWithValidUrn)); } } diff --git a/core/src/test/java/io/substrait/extension/VariadicBehaviorTest.java b/core/src/test/java/io/substrait/extension/VariadicBehaviorTest.java index cc044f350..ad48c9a68 100644 --- a/core/src/test/java/io/substrait/extension/VariadicBehaviorTest.java +++ b/core/src/test/java/io/substrait/extension/VariadicBehaviorTest.java @@ -22,8 +22,7 @@ void testParameterConsistencyLoadingConsistent() { + " parameterConsistency: CONSISTENT\n" + " return: string\n"; - SimpleExtension.ExtensionCollection collection = - SimpleExtension.load("test://example", yamlContent); + SimpleExtension.ExtensionCollection collection = SimpleExtension.load(yamlContent); assertEquals( SimpleExtension.VariadicBehavior.ParameterConsistency.CONSISTENT, @@ -45,8 +44,7 @@ void testParameterConsistencyLoadingInconsistent() { + " parameterConsistency: INCONSISTENT\n" + " return: string\n"; - SimpleExtension.ExtensionCollection collection = - SimpleExtension.load("test://example", yamlContent); + SimpleExtension.ExtensionCollection collection = SimpleExtension.load(yamlContent); assertEquals( SimpleExtension.VariadicBehavior.ParameterConsistency.INCONSISTENT, diff --git a/core/src/test/java/io/substrait/plan/PlanConverterTest.java b/core/src/test/java/io/substrait/plan/PlanConverterTest.java index 9f3e3546e..ab0a321b4 100644 --- a/core/src/test/java/io/substrait/plan/PlanConverterTest.java +++ b/core/src/test/java/io/substrait/plan/PlanConverterTest.java @@ -226,7 +226,7 @@ void nestedUserDefinedTypesShareExtensionCollector() { + " y: T\n" + " z: T\n"; - SimpleExtension.ExtensionCollection extensions = SimpleExtension.load("test.yaml", yaml); + SimpleExtension.ExtensionCollection extensions = SimpleExtension.load(yaml); // Create type objects Type pointType = Type.UserDefined.builder().nullable(false).urn(urn).name("point").build(); diff --git a/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java index 2ea3b7fd6..58bdd1a70 100644 --- a/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java +++ b/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java @@ -64,7 +64,7 @@ class LiteralRoundtripTest extends TestBase { + " value: T\n"; private static final SimpleExtension.ExtensionCollection NESTED_TYPES_EXTENSIONS = - SimpleExtension.load("nested_types.yaml", NESTED_TYPES_YAML); + SimpleExtension.load(NESTED_TYPES_YAML); private static final ExtensionCollector NESTED_TYPES_FUNCTION_COLLECTOR = new ExtensionCollector(); diff --git a/isthmus/src/test/java/io/substrait/isthmus/CustomFunctionTest.java b/isthmus/src/test/java/io/substrait/isthmus/CustomFunctionTest.java index 3297e3c3d..b5a7317b9 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/CustomFunctionTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/CustomFunctionTest.java @@ -53,7 +53,7 @@ class CustomFunctionTest extends PlanTestBase { // Load custom extension into an ExtensionCollection static final SimpleExtension.ExtensionCollection CUSTOM_EXTENSIONS = - SimpleExtension.load(URN, FUNCTIONS_CUSTOM); + SimpleExtension.load(FUNCTIONS_CUSTOM); // Create user-defined types static final String aTypeName = "a_type"; diff --git a/isthmus/src/test/java/io/substrait/isthmus/DuplicateFunctionUrnTest.java b/isthmus/src/test/java/io/substrait/isthmus/DuplicateFunctionUrnTest.java index 1d8cca4cc..8f8cc26cd 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/DuplicateFunctionUrnTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/DuplicateFunctionUrnTest.java @@ -30,9 +30,8 @@ class DuplicateFunctionUrnTest extends PlanTestBase { try { String extensions1 = asString("extensions/functions_duplicate_urn1.yaml"); String extensions2 = asString("extensions/functions_duplicate_urn2.yaml"); - collection1 = - SimpleExtension.load("urn:extension:io.substrait:functions_string", extensions1); - collection2 = SimpleExtension.load("urn:extension:com.domain:string", extensions2); + collection1 = SimpleExtension.load(extensions1); + collection2 = SimpleExtension.load(extensions2); collection = collection1.merge(collection2); // Verify that the merged collection contains duplicate concat functions with different URNs diff --git a/isthmus/src/test/java/io/substrait/isthmus/SimpleExtensionToSqlOperatorTest.java b/isthmus/src/test/java/io/substrait/isthmus/SimpleExtensionToSqlOperatorTest.java index 030e501d8..f97a381cd 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SimpleExtensionToSqlOperatorTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SimpleExtensionToSqlOperatorTest.java @@ -39,7 +39,6 @@ class SimpleExtensionToSqlOperatorTest { static { final SimpleExtension.ExtensionCollection extensions = SimpleExtension.load( - CUSTOM_FUNCTION_PATH, SimpleExtensionToSqlOperatorTest.class.getResourceAsStream(CUSTOM_FUNCTION_PATH)); FUNCTION_DEFS = diff --git a/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java b/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java index 24c085e29..f97616ab6 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java @@ -73,7 +73,7 @@ class UserDefinedLiteralRoundtripTest extends PlanTestBase { + " attributes: map\n"; private static final SimpleExtension.ExtensionCollection NESTED_TYPES_EXTENSIONS = - SimpleExtension.load("nested_types.yaml", NESTED_TYPES_YAML); + SimpleExtension.load(NESTED_TYPES_YAML); private static final Map USER_TYPE_FACTORIES = Map.of( diff --git a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala index 3857874f4..0f224e735 100644 --- a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala +++ b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala @@ -28,7 +28,7 @@ object SparkExtension { final val file = "/spark.yml" private val SparkImpls: SimpleExtension.ExtensionCollection = - SimpleExtension.load(file, getClass.getResourceAsStream(file)) + SimpleExtension.load(getClass.getResourceAsStream(file)) private val EXTENSION_COLLECTION: SimpleExtension.ExtensionCollection = DefaultExtensionCatalog.DEFAULT_COLLECTION; From 62f40dd1e553d079246c94eac9db08265afb3587 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:04:35 -0400 Subject: [PATCH 03/13] refactor: simplify PlanRoundtripTest --- .../io/substrait/plan/PlanRoundtripTest.java | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java index cf8a03ec9..342a06869 100644 --- a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -10,11 +10,9 @@ import java.io.InputStream; import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.stream.Collectors; -import java.util.stream.Stream; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.api.Test; /** * Roundtrip tests: parse a JSON proto plan, convert to POJO, convert back to proto, and compare @@ -22,6 +20,19 @@ */ class PlanRoundtripTest { + private static final List TEST_CASES = + List.of( + new String[] { + "plan-roundtrip/simple-input-plan.json", "plan-roundtrip/simple-expected-plan.json" + }, + new String[] { + "plan-roundtrip/complex-input-plan.json", "plan-roundtrip/complex-expected-plan.json" + }, + new String[] { + "plan-roundtrip/zero-anchor-input-plan.json", + "plan-roundtrip/zero-anchor-expected-plan.json" + }); + private Plan loadPlanFromJson(String resourcePath) throws IOException { try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourcePath)) { if (inputStream == null) { @@ -39,28 +50,7 @@ private Plan loadPlanFromJson(String resourcePath) throws IOException { } } - static Stream roundtripCases() { - return Stream.of( - Arguments.of( - "simple", - "plan-roundtrip/simple-input-plan.json", - "plan-roundtrip/simple-expected-plan.json"), - Arguments.of( - "complex", - "plan-roundtrip/complex-input-plan.json", - "plan-roundtrip/complex-expected-plan.json"), - Arguments.of( - "zero-anchor", - "plan-roundtrip/zero-anchor-input-plan.json", - "plan-roundtrip/zero-anchor-expected-plan.json")); - } - - @ParameterizedTest(name = "{0}") - @MethodSource("roundtripCases") - void testPlanRoundtrip(String name, String inputPath, String expectedPath) throws IOException { - Plan inputPlan = loadPlanFromJson(inputPath); - Plan expectedPlan = loadPlanFromJson(expectedPath); - + private void testPlanRoundtrip(Plan inputPlan, Plan expectedPlan) { ProtoPlanConverter protoToPojo = new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); io.substrait.plan.Plan pojoPlan = protoToPojo.from(inputPlan); @@ -71,4 +61,13 @@ void testPlanRoundtrip(String name, String inputPath, String expectedPath) throw assertEquals(expectedPlan, actualPlan); } + + @Test + void testAllPlanRoundtrips() throws IOException { + for (String[] testCase : TEST_CASES) { + Plan inputPlan = loadPlanFromJson(testCase[0]); + Plan expectedPlan = loadPlanFromJson(testCase[1]); + testPlanRoundtrip(inputPlan, expectedPlan); + } + } } From 29e9ff8b5b624a06469251c2cb42d22c7ce1ae95 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:16:09 -0400 Subject: [PATCH 04/13] refactor: drop unused ExtensionCollection from ImmutableExtensionLookup Add test showing plans with unknown URNs are accepted when those functions aren't referenced by relations. --- .../ProtoExtendedExpressionConverter.java | 2 +- .../extension/ImmutableExtensionLookup.java | 12 +--- .../io/substrait/plan/ProtoPlanConverter.java | 3 +- .../io/substrait/plan/UnknownUrnPlanTest.java | 61 +++++++++++++++++++ 4 files changed, 64 insertions(+), 14 deletions(-) create mode 100644 core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java diff --git a/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java b/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java index 175b2d705..a41a15c8d 100644 --- a/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java +++ b/core/src/main/java/io/substrait/extendedexpression/ProtoExtendedExpressionConverter.java @@ -38,7 +38,7 @@ public ExtendedExpression from(io.substrait.proto.ExtendedExpression extendedExp // fill in simple extension information through a discovery in the current proto-extended // expression ExtensionLookup functionLookup = - ImmutableExtensionLookup.builder(extensionCollection).from(extendedExpression).build(); + ImmutableExtensionLookup.builder().from(extendedExpression).build(); NamedStruct baseSchemaProto = extendedExpression.getBaseSchema(); diff --git a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java index f6bf77386..eb6ffdb24 100644 --- a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java +++ b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java @@ -22,23 +22,13 @@ private ImmutableExtensionLookup( } public static Builder builder() { - return builder(DefaultExtensionCatalog.DEFAULT_COLLECTION); - } - - public static Builder builder(SimpleExtension.ExtensionCollection extensionCollection) { - return new Builder(extensionCollection); + return new Builder(); } public static class Builder { private final Map functionMap = new HashMap<>(); private final Map typeMap = new HashMap<>(); - public Builder(SimpleExtension.ExtensionCollection extensionCollection) { - if (extensionCollection == null) { - throw new IllegalArgumentException("ExtensionCollection is required"); - } - } - private SimpleExtension.FunctionAnchor resolveFunctionAnchor( SimpleExtensionDeclaration.ExtensionFunction func, Map urnMap) { diff --git a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java index 12501c1ac..2e5334b3e 100644 --- a/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java +++ b/core/src/main/java/io/substrait/plan/ProtoPlanConverter.java @@ -67,8 +67,7 @@ protected ProtoRelConverter getProtoRelConverter(final ExtensionLookup functionL } public Plan from(io.substrait.proto.Plan plan) { - ExtensionLookup functionLookup = - ImmutableExtensionLookup.builder(extensionCollection).from(plan).build(); + ExtensionLookup functionLookup = ImmutableExtensionLookup.builder().from(plan).build(); ProtoRelConverter relConverter = getProtoRelConverter(functionLookup); List roots = new ArrayList<>(); for (PlanRel planRel : plan.getRelationsList()) { diff --git a/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java b/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java new file mode 100644 index 000000000..c359e82dc --- /dev/null +++ b/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java @@ -0,0 +1,61 @@ +package io.substrait.plan; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; + +import com.google.protobuf.util.JsonFormat; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.proto.Plan; +import org.junit.jupiter.api.Test; + +/** + * Verifies that plans declaring extensions with unknown URNs can be deserialized, as long as the + * relations don't reference those unknown functions. + */ +class UnknownUrnPlanTest { + + @Test + void testPlanWithUnusedUnknownUrn() throws Exception { + String json = + """ + { + "extensionUrns": [ + { "extensionUrnAnchor": 1, "urn": "extension:com.unknown:functions_custom" } + ], + "extensions": [ + { + "extensionFunction": { + "functionAnchor": 1, + "name": "custom_func:i32_i32", + "extensionUrnReference": 1 + } + } + ], + "relations": [ + { + "root": { + "input": { + "read": { + "baseSchema": { + "names": ["a"], + "struct": { "types": [{ "i32": {} }] } + }, + "namedTable": { "names": ["t"] } + } + }, + "names": ["a"] + } + } + ], + "version": { "minorNumber": 75 } + } + """; + + Plan.Builder builder = Plan.newBuilder(); + JsonFormat.parser().merge(json, builder); + Plan plan = builder.build(); + + ProtoPlanConverter converter = + new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); + assertDoesNotThrow(() -> converter.from(plan)); + } +} From 62d4d17f9edfef95efd0d9e733f24f9ae52f8b53 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:18:39 -0400 Subject: [PATCH 05/13] chore: remove UnknownUrnPlanTest --- .../io/substrait/plan/UnknownUrnPlanTest.java | 61 ------------------- 1 file changed, 61 deletions(-) delete mode 100644 core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java diff --git a/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java b/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java deleted file mode 100644 index c359e82dc..000000000 --- a/core/src/test/java/io/substrait/plan/UnknownUrnPlanTest.java +++ /dev/null @@ -1,61 +0,0 @@ -package io.substrait.plan; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; - -import com.google.protobuf.util.JsonFormat; -import io.substrait.extension.DefaultExtensionCatalog; -import io.substrait.proto.Plan; -import org.junit.jupiter.api.Test; - -/** - * Verifies that plans declaring extensions with unknown URNs can be deserialized, as long as the - * relations don't reference those unknown functions. - */ -class UnknownUrnPlanTest { - - @Test - void testPlanWithUnusedUnknownUrn() throws Exception { - String json = - """ - { - "extensionUrns": [ - { "extensionUrnAnchor": 1, "urn": "extension:com.unknown:functions_custom" } - ], - "extensions": [ - { - "extensionFunction": { - "functionAnchor": 1, - "name": "custom_func:i32_i32", - "extensionUrnReference": 1 - } - } - ], - "relations": [ - { - "root": { - "input": { - "read": { - "baseSchema": { - "names": ["a"], - "struct": { "types": [{ "i32": {} }] } - }, - "namedTable": { "names": ["t"] } - } - }, - "names": ["a"] - } - } - ], - "version": { "minorNumber": 75 } - } - """; - - Plan.Builder builder = Plan.newBuilder(); - JsonFormat.parser().merge(json, builder); - Plan plan = builder.build(); - - ProtoPlanConverter converter = - new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); - assertDoesNotThrow(() -> converter.from(plan)); - } -} From 91397363ab8474465ef905d65a5a275a29b52cc4 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:23:14 -0400 Subject: [PATCH 06/13] refactor: use record for PlanRoundtripTest cases --- .../io/substrait/plan/PlanRoundtripTest.java | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java index 342a06869..aa7833686 100644 --- a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -20,18 +20,19 @@ */ class PlanRoundtripTest { - private static final List TEST_CASES = + record TestCase(String input, String expected) {} + + private static final List TEST_CASES = List.of( - new String[] { - "plan-roundtrip/simple-input-plan.json", "plan-roundtrip/simple-expected-plan.json" - }, - new String[] { - "plan-roundtrip/complex-input-plan.json", "plan-roundtrip/complex-expected-plan.json" - }, - new String[] { - "plan-roundtrip/zero-anchor-input-plan.json", - "plan-roundtrip/zero-anchor-expected-plan.json" - }); + new TestCase( + "plan-roundtrip/simple-input-plan.json", + "plan-roundtrip/simple-expected-plan.json"), + new TestCase( + "plan-roundtrip/complex-input-plan.json", + "plan-roundtrip/complex-expected-plan.json"), + new TestCase( + "plan-roundtrip/zero-anchor-input-plan.json", + "plan-roundtrip/zero-anchor-expected-plan.json")); private Plan loadPlanFromJson(String resourcePath) throws IOException { try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourcePath)) { @@ -64,9 +65,9 @@ private void testPlanRoundtrip(Plan inputPlan, Plan expectedPlan) { @Test void testAllPlanRoundtrips() throws IOException { - for (String[] testCase : TEST_CASES) { - Plan inputPlan = loadPlanFromJson(testCase[0]); - Plan expectedPlan = loadPlanFromJson(testCase[1]); + for (TestCase tc : TEST_CASES) { + Plan inputPlan = loadPlanFromJson(tc.input()); + Plan expectedPlan = loadPlanFromJson(tc.expected()); testPlanRoundtrip(inputPlan, expectedPlan); } } From b25f533a4a2d4194e1c5b659b12c253b298253be Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Mon, 9 Mar 2026 17:31:02 -0400 Subject: [PATCH 07/13] fix: move record after fields to satisfy PMD --- core/src/test/java/io/substrait/plan/PlanRoundtripTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java index aa7833686..cdb287a70 100644 --- a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -20,8 +20,6 @@ */ class PlanRoundtripTest { - record TestCase(String input, String expected) {} - private static final List TEST_CASES = List.of( new TestCase( @@ -63,6 +61,8 @@ private void testPlanRoundtrip(Plan inputPlan, Plan expectedPlan) { assertEquals(expectedPlan, actualPlan); } + record TestCase(String input, String expected) {} + @Test void testAllPlanRoundtrips() throws IOException { for (TestCase tc : TEST_CASES) { From 3feecaf44baec9a958813f5bd8784b6a84d55040 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Tue, 10 Mar 2026 09:23:11 -0400 Subject: [PATCH 08/13] fix: spotless formatting --- core/src/test/java/io/substrait/plan/PlanRoundtripTest.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java index cdb287a70..1d4afda66 100644 --- a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -23,8 +23,7 @@ class PlanRoundtripTest { private static final List TEST_CASES = List.of( new TestCase( - "plan-roundtrip/simple-input-plan.json", - "plan-roundtrip/simple-expected-plan.json"), + "plan-roundtrip/simple-input-plan.json", "plan-roundtrip/simple-expected-plan.json"), new TestCase( "plan-roundtrip/complex-input-plan.json", "plan-roundtrip/complex-expected-plan.json"), From 215022f0b8cdf216c5f9072941ec5ff6b806c74c Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Tue, 10 Mar 2026 09:35:18 -0400 Subject: [PATCH 09/13] docs: add javadoc to ExtensionCollection --- .../main/java/io/substrait/extension/SimpleExtension.java | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/core/src/main/java/io/substrait/extension/SimpleExtension.java b/core/src/main/java/io/substrait/extension/SimpleExtension.java index 1af82526d..e85e0e060 100644 --- a/core/src/main/java/io/substrait/extension/SimpleExtension.java +++ b/core/src/main/java/io/substrait/extension/SimpleExtension.java @@ -592,6 +592,13 @@ public Stream resolve(String urn) { } } + /** + * The catalog of function and type definitions loaded from YAML extension files. Maps URN+name + * pairs to full definitions (argument types, return types, etc.). + * + *

Used by {@link AbstractExtensionLookup#getScalarFunction} and similar methods to resolve a + * {@link FunctionAnchor} into a complete {@link Function} with signature metadata. + */ @Value.Immutable public abstract static class ExtensionCollection { private final Supplier> urnSupplier = From 8ea6ff5d80e72aef4bf09945d3cf28e5e114e6c0 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Tue, 10 Mar 2026 10:21:20 -0400 Subject: [PATCH 10/13] test: add ExtensionLookupTest for missing URN anchors and unregistered references --- .../extension/ImmutableExtensionLookup.java | 1 + .../extension/ExtensionLookupTest.java | 93 +++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 core/src/test/java/io/substrait/extension/ExtensionLookupTest.java diff --git a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java index eb6ffdb24..92acf580f 100644 --- a/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java +++ b/core/src/main/java/io/substrait/extension/ImmutableExtensionLookup.java @@ -69,6 +69,7 @@ private Builder from( List simpleExtensionDeclarations) { Map urnMap = new HashMap<>(); + // TODO: consider validating URN format (e.g. must start with "extension:") for (SimpleExtensionURN extension : simpleExtensionURNs) { urnMap.put(extension.getExtensionUrnAnchor(), extension.getUrn()); } diff --git a/core/src/test/java/io/substrait/extension/ExtensionLookupTest.java b/core/src/test/java/io/substrait/extension/ExtensionLookupTest.java new file mode 100644 index 000000000..36e69ba46 --- /dev/null +++ b/core/src/test/java/io/substrait/extension/ExtensionLookupTest.java @@ -0,0 +1,93 @@ +package io.substrait.extension; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.proto.Plan; +import io.substrait.proto.SimpleExtensionDeclaration; +import io.substrait.proto.SimpleExtensionURN; +import org.junit.jupiter.api.Test; + +class ExtensionLookupTest { + + @Test + void functionReferencingMissingUrnAnchorThrows() { + // function declaration references URN anchor 99, but no URN is registered at that anchor + Plan plan = + Plan.newBuilder() + .addExtensions( + SimpleExtensionDeclaration.newBuilder() + .setExtensionFunction( + SimpleExtensionDeclaration.ExtensionFunction.newBuilder() + .setFunctionAnchor(1) + .setName("add:i32_i32") + .setExtensionUrnReference(99))) + .build(); + + IllegalStateException e = + assertThrows( + IllegalStateException.class, + () -> ImmutableExtensionLookup.builder().from(plan).build()); + assertTrue(e.getMessage().contains("no URN is registered at that anchor")); + } + + @Test + void typeReferencingMissingUrnAnchorThrows() { + // type declaration references URN anchor 99, but no URN is registered at that anchor + Plan plan = + Plan.newBuilder() + .addExtensions( + SimpleExtensionDeclaration.newBuilder() + .setExtensionType( + SimpleExtensionDeclaration.ExtensionType.newBuilder() + .setTypeAnchor(1) + .setName("point") + .setExtensionUrnReference(99))) + .build(); + + IllegalStateException e = + assertThrows( + IllegalStateException.class, + () -> ImmutableExtensionLookup.builder().from(plan).build()); + assertTrue(e.getMessage().contains("no URN is registered at that anchor")); + } + + @Test + void lookupOfUnregisteredFunctionAnchorThrows() { + // build a lookup with one function at anchor 1, then ask for anchor 999 + Plan plan = + Plan.newBuilder() + .addExtensionUrns( + SimpleExtensionURN.newBuilder().setExtensionUrnAnchor(1).setUrn("extension:x:y")) + .addExtensions( + SimpleExtensionDeclaration.newBuilder() + .setExtensionFunction( + SimpleExtensionDeclaration.ExtensionFunction.newBuilder() + .setFunctionAnchor(1) + .setName("add:i32_i32") + .setExtensionUrnReference(1))) + .build(); + + ExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); + SimpleExtension.ExtensionCollection extensions = + SimpleExtension.ExtensionCollection.builder().build(); + + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, () -> lookup.getScalarFunction(999, extensions)); + assertTrue(e.getMessage().contains("Unknown function id")); + } + + @Test + void lookupOfUnregisteredTypeAnchorThrows() { + Plan plan = Plan.newBuilder().build(); + + ExtensionLookup lookup = ImmutableExtensionLookup.builder().from(plan).build(); + SimpleExtension.ExtensionCollection extensions = + SimpleExtension.ExtensionCollection.builder().build(); + + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> lookup.getType(999, extensions)); + assertTrue(e.getMessage().contains("Unknown type id")); + } +} From 1da70fbcaeba6ab4d230dae4c32835ac40a9269e Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Tue, 10 Mar 2026 10:29:26 -0400 Subject: [PATCH 11/13] test: remove trivial ExtensionCollectionMergeTest --- .../ExtensionCollectionMergeTest.java | 60 ------------------- 1 file changed, 60 deletions(-) delete mode 100644 core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java diff --git a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java b/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java deleted file mode 100644 index dfbbd571c..000000000 --- a/core/src/test/java/io/substrait/extension/ExtensionCollectionMergeTest.java +++ /dev/null @@ -1,60 +0,0 @@ -package io.substrait.extension; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import org.junit.jupiter.api.Test; - -class ExtensionCollectionMergeTest { - - @Test - void testMergeCollectionsWithDifferentUrns() { - String yaml1 = - "%YAML 1.2\n" - + "---\n" - + "urn: extension:ns1:collection1\n" - + "scalar_functions:\n" - + " - name: func1\n" - + " impls:\n" - + " - args: []\n" - + " return: boolean\n"; - - String yaml2 = - "%YAML 1.2\n" - + "---\n" - + "urn: extension:ns2:collection2\n" - + "scalar_functions:\n" - + " - name: func2\n" - + " impls:\n" - + " - args: []\n" - + " return: i32\n"; - - SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load(yaml1); - SimpleExtension.ExtensionCollection collection2 = SimpleExtension.load(yaml2); - - SimpleExtension.ExtensionCollection merged = collection1.merge(collection2); - - assertTrue(merged.scalarFunctions().size() >= 2); - } - - @Test - void testMergeCollectionsWithIdenticalUrns() { - String yaml = - "%YAML 1.2\n" - + "---\n" - + "urn: extension:shared:extension\n" - + "scalar_functions:\n" - + " - name: shared_func\n" - + " impls:\n" - + " - args: []\n" - + " return: boolean\n"; - - SimpleExtension.ExtensionCollection collection1 = SimpleExtension.load(yaml); - SimpleExtension.ExtensionCollection collection2 = SimpleExtension.load(yaml); - - SimpleExtension.ExtensionCollection merged = - assertDoesNotThrow(() -> collection1.merge(collection2)); - - assertTrue(merged.scalarFunctions().size() >= 1); - } -} From 51e50953c0f0780884ef65c5f06e254adbf6f3c9 Mon Sep 17 00:00:00 2001 From: Ben Bellick Date: Tue, 10 Mar 2026 10:36:21 -0400 Subject: [PATCH 12/13] refactor: use parameterized test in PlanRoundtripTest --- .../io/substrait/plan/PlanRoundtripTest.java | 51 ++++++++++--------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java index 1d4afda66..cf8a03ec9 100644 --- a/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java +++ b/core/src/test/java/io/substrait/plan/PlanRoundtripTest.java @@ -10,9 +10,11 @@ import java.io.InputStream; import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; -import java.util.List; import java.util.stream.Collectors; -import org.junit.jupiter.api.Test; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; /** * Roundtrip tests: parse a JSON proto plan, convert to POJO, convert back to proto, and compare @@ -20,17 +22,6 @@ */ class PlanRoundtripTest { - private static final List TEST_CASES = - List.of( - new TestCase( - "plan-roundtrip/simple-input-plan.json", "plan-roundtrip/simple-expected-plan.json"), - new TestCase( - "plan-roundtrip/complex-input-plan.json", - "plan-roundtrip/complex-expected-plan.json"), - new TestCase( - "plan-roundtrip/zero-anchor-input-plan.json", - "plan-roundtrip/zero-anchor-expected-plan.json")); - private Plan loadPlanFromJson(String resourcePath) throws IOException { try (InputStream inputStream = getClass().getClassLoader().getResourceAsStream(resourcePath)) { if (inputStream == null) { @@ -48,7 +39,28 @@ private Plan loadPlanFromJson(String resourcePath) throws IOException { } } - private void testPlanRoundtrip(Plan inputPlan, Plan expectedPlan) { + static Stream roundtripCases() { + return Stream.of( + Arguments.of( + "simple", + "plan-roundtrip/simple-input-plan.json", + "plan-roundtrip/simple-expected-plan.json"), + Arguments.of( + "complex", + "plan-roundtrip/complex-input-plan.json", + "plan-roundtrip/complex-expected-plan.json"), + Arguments.of( + "zero-anchor", + "plan-roundtrip/zero-anchor-input-plan.json", + "plan-roundtrip/zero-anchor-expected-plan.json")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("roundtripCases") + void testPlanRoundtrip(String name, String inputPath, String expectedPath) throws IOException { + Plan inputPlan = loadPlanFromJson(inputPath); + Plan expectedPlan = loadPlanFromJson(expectedPath); + ProtoPlanConverter protoToPojo = new ProtoPlanConverter(DefaultExtensionCatalog.DEFAULT_COLLECTION); io.substrait.plan.Plan pojoPlan = protoToPojo.from(inputPlan); @@ -59,15 +71,4 @@ private void testPlanRoundtrip(Plan inputPlan, Plan expectedPlan) { assertEquals(expectedPlan, actualPlan); } - - record TestCase(String input, String expected) {} - - @Test - void testAllPlanRoundtrips() throws IOException { - for (TestCase tc : TEST_CASES) { - Plan inputPlan = loadPlanFromJson(tc.input()); - Plan expectedPlan = loadPlanFromJson(tc.expected()); - testPlanRoundtrip(inputPlan, expectedPlan); - } - } } From 19003101b3e5bf932f025ddafb973934a76b0716 Mon Sep 17 00:00:00 2001 From: Ben Bellick <36523439+benbellick@users.noreply.github.com> Date: Fri, 13 Mar 2026 07:54:17 -0400 Subject: [PATCH 13/13] Update SimpleExtension.java Co-authored-by: Niels Pardon --- core/src/main/java/io/substrait/extension/SimpleExtension.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/java/io/substrait/extension/SimpleExtension.java b/core/src/main/java/io/substrait/extension/SimpleExtension.java index 70182d36e..d98c632f6 100644 --- a/core/src/main/java/io/substrait/extension/SimpleExtension.java +++ b/core/src/main/java/io/substrait/extension/SimpleExtension.java @@ -614,7 +614,7 @@ public Stream resolve(String urn) { } /** - * The catalog of function and type definitions loaded from YAML extension files. Maps URN+name + * The catalog of function and type definitions loaded from YAML extension files. Maps URN + name * pairs to full definitions (argument types, return types, etc.). * *

Used by {@link AbstractExtensionLookup#getScalarFunction} and similar methods to resolve a