[CALCITE-4366] LatticeSuggester should treat UNION and VALUES sub-expressions as tables

core/src/main/codegen/templates/Parser.jj

@@ -1934,7 +1934,7 @@ SqlNode JoinTable(SqlNode e) :
                 joinType,
                 e2,
                 using,
-                new SqlNodeList(list.getList(), Span.of(using).end(this)));
+                new SqlNodeList(list, Span.of(using).end(this)));
         }
     |
         {

core/src/main/java/org/apache/calcite/materialize/LatticeSuggester.java

@@ -17,7 +17,9 @@
 package org.apache.calcite.materialize;
 
 import org.apache.calcite.jdbc.CalciteSchema;
+import org.apache.calcite.plan.RelOptAbstractTable;
 import org.apache.calcite.plan.RelOptCostImpl;
+import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.plan.hep.HepPlanner;
 import org.apache.calcite.plan.hep.HepProgram;
@@ -30,12 +32,24 @@
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.SetOp;
 import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.TableFunctionScan;
 import org.apache.calcite.rel.core.TableScan;
 import org.apache.calcite.rel.rules.CoreRules;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperatorBinding;
+import org.apache.calcite.sql.SqlTableFunction;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.util.CompositeList;
 import org.apache.calcite.util.ImmutableBitSet;
@@ -64,7 +78,9 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -90,6 +106,15 @@ public class LatticeSuggester {
   /** Whether to try to extend an existing lattice when adding a lattice. */
   private final boolean evolve;
 
+  /** Generates unique names. */
+  private final Supplier<String> nameGenerator = new Supplier<String>() {
+    final AtomicInteger nextInt = new AtomicInteger();
+
+    @Override public String get() {
+      return "$table" + nextInt.incrementAndGet();
+    }
+  };
+
   /** Creates a LatticeSuggester. */
   public LatticeSuggester(FrameworkConfig config) {
     this.evolve = config.isEvolveLattice();
@@ -526,13 +551,39 @@ private ColRef toColRef(RexNode e, String alias) {
     } else if (r instanceof TableScan) {
       final TableScan scan = (TableScan) r;
       final TableRef tableRef = q.tableRef(scan);
-      final int fieldCount = r.getRowType().getFieldCount();
-      return new Frame(fieldCount, ImmutableList.of(),
+      final RelDataType rowType = r.getRowType();
+      return new Frame(rowType.getFieldCount(), ImmutableList.of(),
           ImmutableList.of(), ImmutableSet.of(tableRef)) {
         @Override ColRef column(int offset) {
-          if (offset >= scan.getTable().getRowType().getFieldCount()) {
+          if (offset >= rowType.getFieldCount()) {
             throw new IndexOutOfBoundsException("field " + offset
-                + " out of range in " + scan.getTable().getRowType());
+                + " out of range in " + rowType);
+          }
+          return new BaseColRef(tableRef, offset);
+        }
+      };
+    } else if (r instanceof TableFunctionScan) {
+      final TableFunctionScan scan = (TableFunctionScan) r;
+      final RelNode input = scan.getInputs().get(0);
+      String name = null;
+      if (scan.getCall() instanceof RexCall) {
+        final RexCall call = (RexCall) scan.getCall();
+        if (call.getOperator() == WrapFunction.INSTANCE) {
+          name = ((RexLiteral) call.operands.get(1)).getValueAs(String.class);
+        }
+      }
+      if (name == null) {
+        name = nameGenerator.get();
+      }
+      final RelOptTable dummyTable = new DummyTable(input, name);
+      final TableRef tableRef = q.tableRef(scan, dummyTable);
+      final RelDataType rowType = r.getRowType();
+      return new Frame(rowType.getFieldCount(), ImmutableList.of(),
+          ImmutableList.of(), ImmutableSet.of(tableRef)) {
+        @Override ColRef column(int offset) {
+          if (offset >= rowType.getFieldCount()) {
+            throw new IndexOutOfBoundsException("field " + offset
+                + " out of range in " + rowType);
           }
           return new BaseColRef(tableRef, offset);
         }
@@ -554,11 +605,15 @@ private static class Query {
     }
 
     TableRef tableRef(TableScan scan) {
+      return tableRef(scan, scan.getTable());
+    }
+
+    TableRef tableRef(RelNode scan, RelOptTable table) {
       final TableRef r = tableRefs.get(scan.getId());
       if (r != null) {
         return r;
       }
-      final LatticeTable t = space.register(scan.getTable());
+      final LatticeTable t = space.register(table);
       final TableRef r2 = new TableRef(t, tableRefs.size());
       tableRefs.put(scan.getId(), r2);
       return r2;
@@ -812,4 +867,43 @@ private MutableMeasure(SqlAggFunction aggregate, boolean distinct,
     }
   }
 
+  /** Dummy table that represents a RelNode that is to be treated as a leaf
+   * by the lattice suggester. The name is unique within this run of the
+   * suggester. */
+  private static class DummyTable extends RelOptAbstractTable {
+    final RelNode scan;
+
+    DummyTable(RelNode scan, String name) {
+      super(null, name, scan.getRowType());
+      this.scan = scan;
+    }
+
+    @Override public <T> T unwrap(Class<T> clazz) {
+      if (RelNode.class.isAssignableFrom(clazz)
+          && clazz.isInstance(scan)) {
+        return clazz.cast(scan);
+      }
+      return super.unwrap(clazz);
+    }
+  }
+
+  /** "WRAP" user-defined table function. */
+  public static class WrapFunction extends SqlFunction
+      implements SqlTableFunction {
+    public static final WrapFunction INSTANCE = new WrapFunction();
+
+    WrapFunction() {
+      super("WRAP", SqlKind.OTHER_FUNCTION, ReturnTypes.CURSOR,
+          null, OperandTypes.VARIADIC,
+          SqlFunctionCategory.USER_DEFINED_TABLE_FUNCTION);
+    }
+
+    @Override public SqlReturnTypeInference getRowTypeInference() {
+      return WrapFunction::cursor0;
+    }
+
+    private static RelDataType cursor0(SqlOperatorBinding opBinding) {
+      return opBinding.getCursorOperand(0);
+    }
+  }
 }

core/src/main/java/org/apache/calcite/prepare/PlannerImpl.java

@@ -109,7 +109,9 @@ public PlannerImpl(FrameworkConfig config) {
     this.programs = config.getPrograms();
     this.parserConfig = config.getParserConfig();
     this.sqlValidatorConfig = config.getSqlValidatorConfig();
-    this.sqlToRelConverterConfig = config.getSqlToRelConverterConfig();
+    this.sqlToRelConverterConfig = config.getSqlToRelConverterConfig()
+        .addPostStep(SqlToRelConverter::flattenTypesStep)
+        .addPostStep(RelDecorrelator::decorrelateQueryStep);
     this.state = State.STATE_0_CLOSED;
     this.traitDefs = config.getTraitDefs();
     this.convertletTable = config.getConvertletTable();
@@ -252,7 +254,6 @@ private void ready() {
             createCatalogReader(), cluster, convertletTable, config);
     root =
         sqlToRelConverter.convertQuery(validatedSqlNode, false, true);
-    root = root.withRel(sqlToRelConverter.flattenTypes(root.rel, true));
     final RelBuilder relBuilder =
         config.getRelBuilderFactory().create(cluster, null);
     root = root.withRel(

core/src/main/java/org/apache/calcite/rel/core/RelFactories.java

@@ -511,7 +511,7 @@ private static class TableFunctionScanFactoryImpl
       if (call.getOperator() instanceof SqlTableFunction) {
         final SqlOperatorBinding callBinding =
             new RexCallBinding(cluster.getTypeFactory(), call.getOperator(),
-                call.operands, ImmutableList.of());
+                call.operands, ImmutableList.of(), inputs);
         final SqlTableFunction operator = (SqlTableFunction) call.getOperator();
         final SqlReturnTypeInference rowTypeInference =
             operator.getRowTypeInference();

core/src/main/java/org/apache/calcite/rex/RexBuilder.java

@@ -281,7 +281,7 @@ public RelDataType deriveReturnType(
       SqlOperator op,
       List<? extends RexNode> exprs) {
     return op.inferReturnType(
-        new RexCallBinding(typeFactory, op, exprs,
+        new RexCallBinding(typeFactory, op, exprs, ImmutableList.of(),
             ImmutableList.of()));
   }
 

core/src/main/java/org/apache/calcite/rex/RexCallBinding.java

@@ -18,6 +18,7 @@
 
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.runtime.CalciteException;
@@ -41,19 +42,28 @@ public class RexCallBinding extends SqlOperatorBinding {
   //~ Instance fields --------------------------------------------------------
 
   private final List<RexNode> operands;
-
   private final List<RelCollation> inputCollations;
+  private final List<RelNode> inputs;
 
   //~ Constructors -----------------------------------------------------------
 
+  @Deprecated // to be removed before 2.0
   public RexCallBinding(
       RelDataTypeFactory typeFactory,
       SqlOperator sqlOperator,
       List<? extends RexNode> operands,
       List<RelCollation> inputCollations) {
+    this(typeFactory, sqlOperator, operands, inputCollations,
+        ImmutableList.of());
+  }
+
+  public RexCallBinding(RelDataTypeFactory typeFactory, SqlOperator sqlOperator,
+      List<? extends RexNode> operands, List<RelCollation> inputCollations,
+      List<RelNode> inputs) {
     super(typeFactory, sqlOperator);
     this.operands = ImmutableList.copyOf(operands);
     this.inputCollations = ImmutableList.copyOf(inputCollations);
+    this.inputs = ImmutableList.copyOf(inputs);
   }
 
   /** Creates a binding of the appropriate type. */
@@ -68,7 +78,7 @@ public static RexCallBinding create(RelDataTypeFactory typeFactory,
       break;
     }
     return new RexCallBinding(typeFactory, call.getOperator(),
-        call.getOperands(), inputCollations);
+        call.getOperands(), inputCollations, ImmutableList.of());
   }
 
   //~ Methods ----------------------------------------------------------------
@@ -128,16 +138,18 @@ public List<RexNode> operands() {
     return operands;
   }
 
-  // implement SqlOperatorBinding
   @Override public int getOperandCount() {
     return operands.size();
   }
 
-  // implement SqlOperatorBinding
   @Override public RelDataType getOperandType(int ordinal) {
     return operands.get(ordinal).getType();
   }
 
+  @Override public RelDataType getCursorOperand(int ordinal) {
+    return inputs.get(ordinal).getRowType();
+  }
+
   @Override public CalciteException newError(
       Resources.ExInst<SqlValidatorException> e) {
     return SqlUtil.newContextException(SqlParserPos.ZERO, e);
@@ -153,7 +165,8 @@ private static class RexCastCallBinding extends RexCallBinding {
         SqlOperator sqlOperator, List<? extends RexNode> operands,
         RelDataType type,
         List<RelCollation> inputCollations) {
-      super(typeFactory, sqlOperator, operands, inputCollations);
+      super(typeFactory, sqlOperator, operands, inputCollations,
+          ImmutableList.of());
       this.type = type;
     }
 

core/src/main/java/org/apache/calcite/sql/SqlCallBinding.java

@@ -369,9 +369,8 @@ private <T> T valueAs(SqlNode node, Class<T> clazz) {
     if (!SqlUtil.isCallTo(operand, SqlStdOperatorTable.ROW)) {
       return null;
     }
-    for (SqlNode id : ((SqlCall) operand).getOperandList()) {
-      columnList.add(((SqlIdentifier) id).getSimple());
-    }
+    columnList.addAll(
+        SqlIdentifier.simpleNames(((SqlCall) operand).getOperandList()));
     return validator.getParentCursor(paramName);
   }
 

core/src/main/java/org/apache/calcite/sql/SqlHint.java

@@ -18,14 +18,14 @@
 
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.util.NlsString;
+import org.apache.calcite.util.Util;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
 /**
  * A <code>SqlHint</code> is a node of a parse tree which represents
@@ -111,21 +111,17 @@ public HintOptionFormat getOptionFormat() {
    */
   public List<String> getOptionList() {
     if (optionFormat == HintOptionFormat.ID_LIST) {
-      final List<String> attrs = options.getList().stream()
-          .map(node -> ((SqlIdentifier) node).getSimple())
-          .collect(Collectors.toList());
-      return ImmutableList.copyOf(attrs);
+      return ImmutableList.copyOf(SqlIdentifier.simpleNames(options));
     } else if (optionFormat == HintOptionFormat.LITERAL_LIST) {
-      final List<String> attrs = options.getList().stream()
+      return options.stream()
           .map(node -> {
             SqlLiteral literal = (SqlLiteral) node;
             Comparable<?> comparable = SqlLiteral.value(literal);
             return comparable instanceof NlsString
                 ? ((NlsString) comparable).getValue()
                 : comparable.toString();
           })
-          .collect(Collectors.toList());
-      return ImmutableList.copyOf(attrs);
+          .collect(Util.toImmutableList());
     } else {
       return ImmutableList.of();
     }

core/src/main/java/org/apache/calcite/sql/SqlIdentifier.java

@@ -329,6 +329,18 @@ public String getSimple() {
     return names.get(0);
   }
 
+  /** Returns the simple names in a list of identifiers.
+   * Assumes that the list consists of are not-null, simple identifiers. */
+  public static List<String> simpleNames(List<? extends SqlNode> list) {
+    return Util.transform(list, n -> ((SqlIdentifier) n).getSimple());
+  }
+
+  /** Returns the simple names in a iterable of identifiers.
+   * Assumes that the iterable consists of not-null, simple identifiers. */
+  public static Iterable<String> simpleNames(Iterable<? extends SqlNode> list) {
+    return Util.transform(list, n -> ((SqlIdentifier) n).getSimple());
+  }
+
   /**
    * Returns whether this identifier is a star, such as "*" or "foo.bar.*".
    */

core/src/main/java/org/apache/calcite/sql/SqlNode.java

@@ -363,7 +363,7 @@ public static boolean equalDeep(List<SqlNode> operands0,
    * @return a {@code Collector} that collects all the input elements into a
    * {@link SqlNodeList}, in encounter order
    */
-  public static <T extends SqlNode> Collector<T, ArrayList<T>, SqlNodeList>
+  public static <T extends SqlNode> Collector<T, ArrayList<SqlNode>, SqlNodeList>
       toList() {
     return toList(SqlParserPos.ZERO);
   }
@@ -377,9 +377,9 @@ public static boolean equalDeep(List<SqlNode> operands0,
    * @return a {@code Collector} that collects all the input elements into a
    * {@link SqlNodeList}, in encounter order
    */
-  public static <T extends SqlNode> Collector<T, ArrayList<T>, SqlNodeList>
+  public static <T extends SqlNode> Collector<T, ArrayList<SqlNode>, SqlNodeList>
       toList(SqlParserPos pos) {
     return Collector.of(ArrayList::new, ArrayList::add, Util::combine,
-        list -> new SqlNodeList(list, pos));
+        list -> SqlNodeList.of(pos, list));
   }
 }