diff --git a/gradle.properties b/gradle.properties index e2062da03..95a0cc381 100644 --- a/gradle.properties +++ b/gradle.properties @@ -15,7 +15,7 @@ com.github.vlsi.vlsi-release-plugins.version=1.74 # library version antlr.version=4.13.1 -calcite.version=1.38.0 +calcite.version=1.39.0 guava.version=32.1.3-jre immutables.version=2.10.1 jackson.version=2.16.1 diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitCalciteSchema.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitCalciteSchema.java new file mode 100644 index 000000000..d0dca9f56 --- /dev/null +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitCalciteSchema.java @@ -0,0 +1,150 @@ +package io.substrait.isthmus; + +import io.substrait.relation.NamedScan; +import io.substrait.relation.Rel; +import io.substrait.relation.RelCopyOnWriteVisitor; +import io.substrait.type.NamedStruct; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Collectors; +import org.apache.calcite.jdbc.CalciteSchema; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.schema.Table; +import org.apache.calcite.schema.impl.AbstractSchema; +import org.apache.calcite.schema.lookup.Lookup; + +/** + * A subclass of the Calcite Schema for creation from a Substrait relation + * + *

Implementation note: + * + *

The external Isthmus API can take a function that will return the table schema when needed, + * rather than it being available up front. + * + *

This was implemented by a special subclass of the Calcite simple schema. Since this was + * changed in Calcite 1.39.0; it failed to work; the protected methods it extended from changed. + * + *

The new feature of the Calcite schema is the 'lazy' or delayed lookup of tables. This feature + * has not be exploited here + */ +public class SubstraitCalciteSchema extends AbstractSchema { + + private Map tables; + + protected SubstraitCalciteSchema(Map tables) { + this.tables = tables; + } + + @Override + protected Map getTableMap() { + return tables; + } + + @Override + public Lookup tables() { + return super.tables(); + } + + /** + * Turn this into a root Calciteschema Choice of settings is based on current isthmus behaviour + */ + public CalciteSchema getRootSchema() { + return CalciteSchema.createRootSchema(false, false, "", this); + } + + public static Builder builder() { + return new Builder(); + } + + /** + * Builder class to assist with creating the CalciteSchema + * + *

Can be created from a Rel or a Lookup function + */ + public static class Builder { + + private Rel rel; + private RelDataTypeFactory typeFactory; + private TypeConverter typeConverter; + + public Builder withTypeFactory(RelDataTypeFactory typeFactory) { + this.typeFactory = typeFactory; + return this; + } + + public Builder withTypeConverter(TypeConverter typeConverter) { + this.typeConverter = typeConverter; + return this; + } + + public Builder withSubstraitRel(Rel rel) { + this.rel = rel; + return this; + } + + public SubstraitCalciteSchema build() { + if (typeConverter == null) { + throw new IllegalArgumentException("'TypeConverter' must be specified"); + } + + if (typeFactory == null) { + throw new IllegalArgumentException("'TypeFactory' must be specified"); + } + + if (rel == null) { + throw new IllegalArgumentException("'rel' must be specified"); + } + + // If there are any named structs within the relation, gather these and convert + // them to a map of tables + // index by name; note that the name of the table is 'un-namespaced' here. + // This was the existing logic so it has not been altered. + Map, NamedStruct> tableMap = NamedStructGatherer.gatherTables(rel); + + Map tables = + tableMap.entrySet().stream() + .map( + entry -> { + var id = entry.getKey(); + var name = id.get(id.size() - 1); + var table = entry.getValue(); + var value = + new SqlConverterBase.DefinedTable( + name, + typeFactory, + typeConverter.toCalcite(typeFactory, table.struct(), table.names())); + return Map.entry(name, value); + }) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + return new SubstraitCalciteSchema(tables); + } + } + + private static final class NamedStructGatherer extends RelCopyOnWriteVisitor { + Map, NamedStruct> tableMap; + + private NamedStructGatherer() { + super(); + this.tableMap = new HashMap<>(); + } + + public static Map, NamedStruct> gatherTables(Rel rel) { + var visitor = new NamedStructGatherer(); + rel.accept(visitor); + return visitor.tableMap; + } + + @Override + public Optional visit(NamedScan namedScan) { + Optional result = super.visit(namedScan); + + List tableName = namedScan.getNames(); + tableMap.put(tableName, namedScan.getInitialSchema()); + + return result; + } + } +} diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java index a96185a22..fd8066168 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java @@ -2,24 +2,15 @@ import io.substrait.extension.SimpleExtension; import io.substrait.plan.Plan; -import io.substrait.relation.NamedScan; import io.substrait.relation.Rel; -import io.substrait.relation.RelCopyOnWriteVisitor; -import io.substrait.type.NamedStruct; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.function.Function; import org.apache.calcite.jdbc.CalciteSchema; -import org.apache.calcite.jdbc.LookupCalciteSchema; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rel.type.RelDataTypeField; -import org.apache.calcite.schema.Table; import org.apache.calcite.sql.SqlKind; import org.apache.calcite.tools.Frameworks; import org.apache.calcite.tools.RelBuilder; @@ -57,19 +48,12 @@ public SubstraitToCalcite( *

Override this method to customize schema extraction. */ protected CalciteSchema toSchema(Rel rel) { - Map, NamedStruct> tableMap = NamedStructGatherer.gatherTables(rel); - Function, Table> lookup = - id -> { - NamedStruct table = tableMap.get(id); - if (table == null) { - return null; - } - return new SqlConverterBase.DefinedTable( - id.get(id.size() - 1), - typeFactory, - typeConverter.toCalcite(typeFactory, table.struct(), table.names())); - }; - return LookupCalciteSchema.createRootSchema(lookup); + return SubstraitCalciteSchema.builder() + .withSubstraitRel(rel) + .withTypeFactory(typeFactory) + .withTypeConverter(typeConverter) + .build() + .getRootSchema(); } /** @@ -179,29 +163,4 @@ private Pair renameFields( return Pair.of(currentIndex, type); } } - - private static class NamedStructGatherer extends RelCopyOnWriteVisitor { - Map, NamedStruct> tableMap; - - private NamedStructGatherer() { - super(); - this.tableMap = new HashMap<>(); - } - - public static Map, NamedStruct> gatherTables(Rel rel) { - var visitor = new NamedStructGatherer(); - rel.accept(visitor); - return visitor.tableMap; - } - - @Override - public Optional visit(NamedScan namedScan) { - Optional result = super.visit(namedScan); - - List tableName = namedScan.getNames(); - tableMap.put(tableName, namedScan.getInitialSchema()); - - return result; - } - } } diff --git a/isthmus/src/main/java/org/apache/calcite/jdbc/LookupCalciteSchema.java b/isthmus/src/main/java/org/apache/calcite/jdbc/LookupCalciteSchema.java deleted file mode 100644 index 7222dbc6f..000000000 --- a/isthmus/src/main/java/org/apache/calcite/jdbc/LookupCalciteSchema.java +++ /dev/null @@ -1,55 +0,0 @@ -package org.apache.calcite.jdbc; - -import com.google.common.collect.Maps; -import java.util.List; -import java.util.Map; -import java.util.function.Function; -import org.apache.calcite.schema.Schema; -import org.apache.calcite.schema.Table; -import org.apache.calcite.schema.impl.AbstractSchema; -import org.checkerframework.checker.nullness.qual.Nullable; - -public class LookupCalciteSchema extends SimpleCalciteSchema { - private final Function, Table> lookup; - private final Map, Table> cache = Maps.newHashMap(); - - LookupCalciteSchema( - @Nullable CalciteSchema parent, - Schema schema, - String name, - Function, Table> lookup) { - super(parent, schema, name); - this.lookup = lookup; - } - - @Override - public CalciteSchema add(String name, Schema schema) { - final CalciteSchema calciteSchema = new LookupCalciteSchema(this, schema, name, lookup); - subSchemaMap.put(name, calciteSchema); - return calciteSchema; - } - - @Override - protected @Nullable CalciteSchema getImplicitSubSchema(String schemaName, boolean caseSensitive) { - if (cache.computeIfAbsent(path(schemaName), lookup) != null) { - return null; - } - plus().add(schemaName, AbstractSchema.Factory.INSTANCE.create(null, null, null)); - return super.getSubSchema(schemaName, caseSensitive); - } - - @Override - protected @Nullable TableEntry getImplicitTable(String tableName, boolean caseSensitive) { - Table table = cache.computeIfAbsent(path(tableName), lookup); - if (table == null) { - return null; - } - add(tableName, table); - return getTable(tableName, caseSensitive); - } - - public static CalciteSchema createRootSchema(Function, Table> lookup) { - Schema rootSchema = new CalciteConnectionImpl.RootSchema(); - return new LookupCalciteSchema(null, rootSchema, "", lookup); - } -}