Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -20,15 +19,10 @@
public class ExtensionCollector extends AbstractExtensionLookup {
private final BidiMap<Integer, SimpleExtension.FunctionAnchor> funcMap;
private final BidiMap<Integer, SimpleExtension.TypeAnchor> 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);
}
Expand All @@ -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) {
Expand All @@ -67,30 +60,24 @@ public void addExtensionsToPlan(Plan.Builder builder) {
SimpleExtensions simpleExtensions = getExtensions();

builder.addAllExtensionUrns(simpleExtensions.urns.values());
builder.addAllExtensionUris(simpleExtensions.uris.values());
builder.addAllExtensions(simpleExtensions.extensionList);
}

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<String, SimpleExtensionURN> urns = new HashMap<>();
HashMap<String, SimpleExtensionURI> uris = new HashMap<>();

ArrayList<SimpleExtensionDeclaration> extensionList = new ArrayList<>();
for (Map.Entry<Integer, SimpleExtension.FunctionAnchor> e : funcMap.forwardEntrySet()) {
String urn = e.getValue().urn();
String uri = getUriFromUrn(urn);

// Create URN entry
SimpleExtensionURN urnObj =
urns.computeIfAbsent(
urn,
Expand All @@ -100,40 +87,20 @@ 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);
}

for (Map.Entry<Integer, SimpleExtension.TypeAnchor> e : typeMap.forwardEntrySet()) {
String urn = e.getValue().urn();
String uri = getUriFromUrn(urn);

// Create URN entry
SimpleExtensionURN urnObj =
urns.computeIfAbsent(
urn,
Expand All @@ -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<String, SimpleExtensionURN> urns;
final HashMap<String, SimpleExtensionURI> uris;
final ArrayList<SimpleExtensionDeclaration> extensionList;

SimpleExtensions(
HashMap<String, SimpleExtensionURN> urns,
HashMap<String, SimpleExtensionURI> uris,
ArrayList<SimpleExtensionDeclaration> extensionList) {
this.urns = urns;
this.uris = uris;
this.extensionList = extensionList;
}
}
Expand Down
Loading
Loading