Java Code Examples for org.apache.calcite.tools.RelBuilder#build()

@Override
public RelNode visit(final LogicalFilter filter) {
  final RelBuilder relBuilder = newCalciteRelBuilderWithoutContext(filter.getCluster());
  RelNode input = filter.getInput().accept(this);
  relBuilder.push(input);

  RexNode newCondition = filter.getCondition().accept(new RexShuttle() {
    @Override
    public RexNode visitInputRef(RexInputRef inputRef) {
      return relBuilder.field(filter.getRowType().getFieldNames().get(inputRef.getIndex()));
    }
  });

  relBuilder.filter(newCondition);
  return relBuilder.build();
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final Join join = call.rel(0);
  final RelNode left = join.getLeft();
  final RelNode right = join.getRight();
  final RelNode convertedLeft = convertIfNecessary(left);
  final RelNode convertedRight = convertIfNecessary(right);
  final RelBuilder builder = factory.create(join.getCluster(), null);

  builder.pushAll(ImmutableList.of(convertedLeft, convertedRight));
  builder.join(join.getJoinType(), join.getCondition());

  final RelNode newJoin = builder.build();
  if(newJoin != null) {
    call.transformTo(newJoin);
  }
}
 

@Test void testTableFunctionScanZeroInputs() {
  // Equivalent SQL:
  //   SELECT *
  //   FROM TABLE(RAMP(3))
  final RelBuilder builder = RelBuilder.create(config().build());
  final SqlOperator rampFunction = new MockSqlOperatorTable.RampFunction();
  RelNode root = builder.functionScan(rampFunction, 0, builder.literal(3))
      .build();
  final String expected = "LogicalTableFunctionScan(invocation=[RAMP(3)], "
      + "rowType=[RecordType(INTEGER I)])\n";
  assertThat(root, hasTree(expected));

  // Make sure that the builder's stack is empty.
  try {
    RelNode node = builder.build();
    fail("expected error, got " + node);
  } catch (NoSuchElementException e) {
    assertNull(e.getMessage());
  }
}
 

@Override
protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder, RelNode topProject,
        RelNode unionInputQuery, RelNode unionInputView) {
    relBuilder.push(unionInputQuery);
    relBuilder.push(unionInputView);
    relBuilder.union(true);
    List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
        // We can take unionInputQuery as it is query based.
        RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
        exprList.add(
                rexBuilder.ensureType(field.getType(), rexBuilder.makeInputRef(relBuilder.peek(), i), true));
        nameList.add(field.getName());
    }
    relBuilder.project(exprList, nameList);
    return relBuilder.build();
}
 

@Override protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder,
    RelNode topProject, RelNode unionInputQuery, RelNode unionInputView) {
  relBuilder.push(unionInputQuery);
  relBuilder.push(unionInputView);
  relBuilder.union(true);
  List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
  for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
    // We can take unionInputQuery as it is query based.
    RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
    exprList.add(
        rexBuilder.ensureType(
            field.getType(),
            rexBuilder.makeInputRef(relBuilder.peek(), i),
            true));
    nameList.add(field.getName());
  }
  relBuilder.project(exprList, nameList);
  return relBuilder.build();
}
 

public void runAllExamples() {
  // Create a builder. The config contains a schema mapped
  // to the SCOTT database, with tables EMP and DEPT.
  final FrameworkConfig config = RelBuilderTest.config().build();
  final RelBuilder builder = RelBuilder.create(config);
  for (int i = 0; i < 4; i++) {
    doExample(builder, i);
    final RelNode node = builder.build();
    if (verbose) {
      System.out.println(RelOptUtil.toString(node));
    }
  }
}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final Project proj = call.rel(0);
  final ScanCrel scan = call.rel(1);

  ProjectPushInfo columnInfo = PrelUtil.getColumns(scan.getRowType(), proj.getProjects());

  // get TableBase, either wrapped in RelOptTable, or TranslatableTable. TableBase table = scan.getTable().unwrap(TableBase.class);
  if (columnInfo == null || columnInfo.isStarQuery()) {
    return;
  }

  ScanCrel newScan = scan.cloneWithProject(columnInfo.columns);

  List<RexNode> newProjects = Lists.newArrayList();
  for (RexNode n : proj.getChildExps()) {
    newProjects.add(n.accept(columnInfo.getInputRewriter()));
  }

  final RelBuilder relBuilder = relBuilderFactory.create(proj.getCluster(), null);
  relBuilder.push(newScan);
  relBuilder.project(newProjects, proj.getRowType().getFieldNames());
  final RelNode newProj = relBuilder.build();

  if (newProj instanceof Project
      && ProjectRemoveRule.isTrivial((Project) newProj)
      && newScan.getRowType().getFullTypeString().equals(newProj.getRowType().getFullTypeString())) {
      call.transformTo(newScan);
  } else {
    if(newScan.getProjectedColumns().equals(scan.getProjectedColumns())) {
      // no point in doing a pushdown that doesn't change anything.
      return;
    }

    call.transformTo(newProj);
  }
}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

/**
 * Attempt to create a new join with a canonicalized join expression, and a possible filter
 * on top
 * @param builder
 * @param join
 * @return a new join tree (or same as original argument if no change)
 */
private RelNode getNewJoinCondition(RelBuilder builder, Join join) {
  final List<Integer> leftKeys = Lists.newArrayList();
  final List<Integer> rightKeys = Lists.newArrayList();
  final List<Boolean> filterNulls = Lists.newArrayList();

  final RexNode remaining = RelOptUtil.splitJoinCondition(join.getLeft(), join.getRight(), join.getCondition(), leftKeys, rightKeys, filterNulls);
  final boolean hasEquiJoins = leftKeys.size() == rightKeys.size() && leftKeys.size() > 0 ;
  final JoinRelType joinType = join.getJoinType();

  // If join has no equi-join condition, do not transform
  if (!hasEquiJoins) {
    return join;
  }

  // Create a new partial condition for the equi-join
  final RexNode partialCondition = JoinFilterCanonicalizationRule.buildJoinCondition(
      builder.getRexBuilder(),
      join.getLeft().getRowType(),
      join.getRight().getRowType(),
      leftKeys,
      rightKeys,
      filterNulls);

  // We do not know how to add filter for non-INNER joins (see DRILL-1337)
  if (joinType != JoinRelType.INNER) {
    final RexNode newJoinCondition = RexUtil.composeConjunction(builder.getRexBuilder(), ImmutableList.of(partialCondition, remaining), false);
    if (RexUtil.eq(join.getCondition(), newJoinCondition)) {
      // Condition is the same, do not create a new rel
      return join;
    }
    builder.pushAll(ImmutableList.of(join.getLeft(), join.getRight()));
    builder.join(joinType, newJoinCondition);

    return builder.build();
  }

  // Check if join condition has changed if pure equi-join
  if (remaining.isAlwaysTrue() && RexUtil.eq(join.getCondition(), partialCondition)) {
    return join;
  }

  // Return the new join with a filter on top
  builder.pushAll(ImmutableList.of(join.getLeft(), join.getRight()));
  builder.join(joinType, partialCondition);
  builder.filter(remaining);
  return builder.build();
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Join join = call.rel(0);
  final int leftFieldCount = join.getLeft().getRowType().getFieldCount();
  final RelOptCluster cluster = join.getCluster();
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  final RelBuilder relBuilder = call.builder();

  final Set<CorrelationId> correlationIds = new HashSet<>();
  final ArrayList<RexNode> corrVar = new ArrayList<>();

  for (int i = 0; i < batchSize; i++) {
    CorrelationId correlationId = cluster.createCorrel();
    correlationIds.add(correlationId);
    corrVar.add(
        rexBuilder.makeCorrel(join.getLeft().getRowType(),
            correlationId));
  }

  final ImmutableBitSet.Builder requiredColumns = ImmutableBitSet.builder();

  // Generate first condition
  final RexNode condition = join.getCondition().accept(new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef input) {
      int field = input.getIndex();
      if (field >= leftFieldCount) {
        return rexBuilder.makeInputRef(input.getType(),
            input.getIndex() - leftFieldCount);
      }
      requiredColumns.set(field);
      return rexBuilder.makeFieldAccess(corrVar.get(0), field);
    }
  });

  List<RexNode> conditionList = new ArrayList<>();
  conditionList.add(condition);

  // Add batchSize-1 other conditions
  for (int i = 1; i < batchSize; i++) {
    final int corrIndex = i;
    final RexNode condition2 = condition.accept(new RexShuttle() {
      @Override public RexNode visitCorrelVariable(RexCorrelVariable variable) {
        return corrVar.get(corrIndex);
      }
    });
    conditionList.add(condition2);
  }

  // Push a filter with batchSize disjunctions
  relBuilder.push(join.getRight()).filter(relBuilder.or(conditionList));
  RelNode right = relBuilder.build();

  JoinRelType joinType = join.getJoinType();
  call.transformTo(
      EnumerableBatchNestedLoopJoin.create(
          convert(join.getLeft(), join.getLeft().getTraitSet()
              .replace(EnumerableConvention.INSTANCE)),
          convert(right, right.getTraitSet()
              .replace(EnumerableConvention.INSTANCE)),
          join.getCondition(),
          requiredColumns.build(),
          correlationIds,
          joinType));
}
 

public void onMatch(RelOptRuleCall call) {
  final Project origProject;
  final Filter filter;
  if (call.rels.length >= 2) {
    origProject = call.rel(0);
    filter = call.rel(1);
  } else {
    origProject = null;
    filter = call.rel(0);
  }
  final RelNode input = filter.getInput();
  final RexNode origFilter = filter.getCondition();

  if ((origProject != null)
      && RexOver.containsOver(origProject.getProjects(), null)) {
    // Cannot push project through filter if project contains a windowed
    // aggregate -- it will affect row counts. Abort this rule
    // invocation; pushdown will be considered after the windowed
    // aggregate has been implemented. It's OK if the filter contains a
    // windowed aggregate.
    return;
  }

  if ((origProject != null)
      && origProject.getRowType().isStruct()
      && origProject.getRowType().getFieldList().stream()
        .anyMatch(RelDataTypeField::isDynamicStar)) {
    // The PushProjector would change the plan:
    //
    //    prj(**=[$0])
    //    : - filter
    //        : - scan
    //
    // to form like:
    //
    //    prj(**=[$0])                    (1)
    //    : - filter                      (2)
    //        : - prj(**=[$0], ITEM= ...) (3)
    //            :  - scan
    // This new plan has more cost that the old one, because of the new
    // redundant project (3), if we also have FilterProjectTransposeRule in
    // the rule set, it will also trigger infinite match of the ProjectMergeRule
    // for project (1) and (3).
    return;
  }

  final RelBuilder builder = call.builder();
  final RelNode topProject;
  if (origProject != null && (wholeProject || wholeFilter)) {
    builder.push(input);

    final Set<RexNode> set = new LinkedHashSet<>();
    final RelOptUtil.InputFinder refCollector = new RelOptUtil.InputFinder();

    if (wholeFilter) {
      set.add(filter.getCondition());
    } else {
      filter.getCondition().accept(refCollector);
    }
    if (wholeProject) {
      set.addAll(origProject.getProjects());
    } else {
      refCollector.visitEach(origProject.getProjects());
    }

    // Build a list with inputRefs, in order, first, then other expressions.
    final List<RexNode> list = new ArrayList<>();
    final ImmutableBitSet refs = refCollector.build();
    for (RexNode field : builder.fields()) {
      if (refs.get(((RexInputRef) field).getIndex()) || set.contains(field)) {
        list.add(field);
      }
    }
    set.removeAll(list);
    list.addAll(set);
    builder.project(list);
    final Replacer replacer = new Replacer(list, builder);
    builder.filter(replacer.visit(filter.getCondition()));
    builder.project(replacer.visitList(origProject.getProjects()),
        origProject.getRowType().getFieldNames());
    topProject = builder.build();
  } else {
    // The traditional mode of operation of this rule: push down field
    // references. The effect is similar to RelFieldTrimmer.
    final PushProjector pushProjector =
        new PushProjector(origProject, origFilter, input,
            preserveExprCondition, builder);
    topProject = pushProjector.convertProject(null);
  }

  if (topProject != null) {
    call.transformTo(topProject);
  }
}
 

protected void perform(RelOptRuleCall call, Project project,
    Join join, RelNode left, Aggregate aggregate) {
  final RelOptCluster cluster = join.getCluster();
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  if (project != null) {
    final ImmutableBitSet bits =
        RelOptUtil.InputFinder.bits(project.getProjects(), null);
    final ImmutableBitSet rightBits =
        ImmutableBitSet.range(left.getRowType().getFieldCount(),
            join.getRowType().getFieldCount());
    if (bits.intersects(rightBits)) {
      return;
    }
  } else {
    if (join.getJoinType().projectsRight()
        && !IS_EMPTY_AGGREGATE.test(aggregate)) {
      return;
    }
  }
  final JoinInfo joinInfo = join.analyzeCondition();
  if (!joinInfo.rightSet().equals(
      ImmutableBitSet.range(aggregate.getGroupCount()))) {
    // Rule requires that aggregate key to be the same as the join key.
    // By the way, neither a super-set nor a sub-set would work.
    return;
  }
  if (!joinInfo.isEqui()) {
    return;
  }
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(left);
  switch (join.getJoinType()) {
  case SEMI:
  case INNER:
    final List<Integer> newRightKeyBuilder = new ArrayList<>();
    final List<Integer> aggregateKeys = aggregate.getGroupSet().asList();
    for (int key : joinInfo.rightKeys) {
      newRightKeyBuilder.add(aggregateKeys.get(key));
    }
    final ImmutableIntList newRightKeys = ImmutableIntList.copyOf(newRightKeyBuilder);
    relBuilder.push(aggregate.getInput());
    final RexNode newCondition =
        RelOptUtil.createEquiJoinCondition(relBuilder.peek(2, 0),
            joinInfo.leftKeys, relBuilder.peek(2, 1), newRightKeys,
            rexBuilder);
    relBuilder.semiJoin(newCondition);
    break;

  case LEFT:
    // The right-hand side produces no more than 1 row (because of the
    // Aggregate) and no fewer than 1 row (because of LEFT), and therefore
    // we can eliminate the semi-join.
    break;

  default:
    throw new AssertionError(join.getJoinType());
  }
  if (project != null) {
    relBuilder.project(project.getProjects(), project.getRowType().getFieldNames());
  }
  final RelNode relNode = relBuilder.build();
  call.transformTo(relNode);
}
 

public static RelNode apply(RelOptRuleCall call, Aggregate aggregate,
    Project project) {
  // Find all fields which we need to be straightforward field projections.
  final Set<Integer> interestingFields = RelOptUtil.getAllFields(aggregate);

  // Build the map from old to new; abort if any entry is not a
  // straightforward field projection.
  final Map<Integer, Integer> map = new HashMap<>();
  for (int source : interestingFields) {
    final RexNode rex = project.getProjects().get(source);
    if (!(rex instanceof RexInputRef)) {
      return null;
    }
    map.put(source, ((RexInputRef) rex).getIndex());
  }

  final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (aggregate.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(aggregate.getGroupSets(), map));
  }

  final ImmutableList.Builder<AggregateCall> aggCalls =
      ImmutableList.builder();
  final int sourceCount = aggregate.getInput().getRowType().getFieldCount();
  final int targetCount = project.getInput().getRowType().getFieldCount();
  final Mappings.TargetMapping targetMapping =
      Mappings.target(map, sourceCount, targetCount);
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    aggCalls.add(aggregateCall.transform(targetMapping));
  }

  final Aggregate newAggregate =
      aggregate.copy(aggregate.getTraitSet(), project.getInput(),
          newGroupSet, newGroupingSets, aggCalls.build());

  // Add a project if the group set is not in the same order or
  // contains duplicates.
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(newAggregate);
  final List<Integer> newKeys =
      Lists.transform(aggregate.getGroupSet().asList(), map::get);
  if (!newKeys.equals(newGroupSet.asList())) {
    final List<Integer> posList = new ArrayList<>();
    for (int newKey : newKeys) {
      posList.add(newGroupSet.indexOf(newKey));
    }
    for (int i = newAggregate.getGroupCount();
         i < newAggregate.getRowType().getFieldCount(); i++) {
      posList.add(i);
    }
    relBuilder.project(relBuilder.fields(posList));
  }

  return relBuilder.build();
}
 

/**
 * Creates the topmost projection that will sit on top of the selected join
 * ordering. The projection needs to match the original join ordering. Also,
 * places any post-join filters on top of the project.
 *
 * @param multiJoin join factors being optimized
 * @param joinTree selected join ordering
 * @param fieldNames field names corresponding to the projection expressions
 *
 * @return created projection
 */
private RelNode createTopProject(
    RelBuilder relBuilder,
    LoptMultiJoin multiJoin,
    LoptJoinTree joinTree,
    List<String> fieldNames) {
  List<RexNode> newProjExprs = new ArrayList<>();
  RexBuilder rexBuilder =
      multiJoin.getMultiJoinRel().getCluster().getRexBuilder();

  // create a projection on top of the joins, matching the original
  // join order
  final List<Integer> newJoinOrder = joinTree.getTreeOrder();
  int nJoinFactors = multiJoin.getNumJoinFactors();
  List<RelDataTypeField> fields = multiJoin.getMultiJoinFields();

  // create a mapping from each factor to its field offset in the join
  // ordering
  final Map<Integer, Integer> factorToOffsetMap = new HashMap<>();
  for (int pos = 0, fieldStart = 0; pos < nJoinFactors; pos++) {
    factorToOffsetMap.put(newJoinOrder.get(pos), fieldStart);
    fieldStart +=
        multiJoin.getNumFieldsInJoinFactor(newJoinOrder.get(pos));
  }

  for (int currFactor = 0; currFactor < nJoinFactors; currFactor++) {
    // if the factor is the right factor in a removable self-join,
    // then where possible, remap references to the right factor to
    // the corresponding reference in the left factor
    Integer leftFactor = null;
    if (multiJoin.isRightFactorInRemovableSelfJoin(currFactor)) {
      leftFactor = multiJoin.getOtherSelfJoinFactor(currFactor);
    }
    for (int fieldPos = 0;
        fieldPos < multiJoin.getNumFieldsInJoinFactor(currFactor);
        fieldPos++) {
      int newOffset = factorToOffsetMap.get(currFactor) + fieldPos;
      if (leftFactor != null) {
        Integer leftOffset =
            multiJoin.getRightColumnMapping(currFactor, fieldPos);
        if (leftOffset != null) {
          newOffset =
              factorToOffsetMap.get(leftFactor) + leftOffset;
        }
      }
      newProjExprs.add(
          rexBuilder.makeInputRef(
              fields.get(newProjExprs.size()).getType(),
              newOffset));
    }
  }

  relBuilder.push(joinTree.getJoinTree());
  relBuilder.project(newProjExprs, fieldNames);

  // Place the post-join filter (if it exists) on top of the final
  // projection.
  RexNode postJoinFilter =
      multiJoin.getMultiJoinRel().getPostJoinFilter();
  if (postJoinFilter != null) {
    relBuilder.filter(postJoinFilter);
  }
  return relBuilder.build();
}
 

/**
 * Tests the ElasticSearch match query. The match query is translated from
 * CONTAINS query which is build using RelBuilder, RexBuilder because the
 * normal SQL query assumes CONTAINS query is for date/period range.
 *
 * <p>Equivalent SQL query:
 *
 * <blockquote>
 * <code>select * from zips where city contains 'waltham'</code>
 * </blockquote>
 *
 * <p>ElasticSearch query for it:
 *
 * <blockquote><code>
 * {"query":{"constant_score":{"filter":{"match":{"city":"waltham"}}}}}
 * </code></blockquote>
 */
@Test void testMatchQuery() throws Exception {

  CalciteConnection con = (CalciteConnection) newConnectionFactory()
      .createConnection();
  SchemaPlus postSchema = con.getRootSchema().getSubSchema("elastic");

  FrameworkConfig postConfig = Frameworks.newConfigBuilder()
      .parserConfig(SqlParser.Config.DEFAULT)
      .defaultSchema(postSchema)
      .build();

  final RelBuilder builder = RelBuilder.create(postConfig);
  builder.scan(ZIPS);

  final RelDataTypeFactory typeFactory =
      new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
  final RexBuilder rexBuilder = new RexBuilder(typeFactory);

  RexNode nameRexNode = rexBuilder.makeCall(SqlStdOperatorTable.ITEM,
      rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.ANY), 0),
      rexBuilder.makeCharLiteral(
          new NlsString("city", typeFactory.getDefaultCharset().name(),
              SqlCollation.COERCIBLE)));

  RelDataType mapType = typeFactory.createMapType(
      typeFactory.createSqlType(SqlTypeName.VARCHAR),
      typeFactory.createTypeWithNullability(
          typeFactory.createSqlType(SqlTypeName.ANY), true));

  List<RexNode> namedList =
      ImmutableList.of(rexBuilder.makeInputRef(mapType, 0),
          nameRexNode);

  // Add fields in builder stack so it is accessible while filter preparation
  builder.projectNamed(namedList, Arrays.asList("_MAP", "city"), true);

  RexNode filterRexNode = builder
      .call(SqlStdOperatorTable.CONTAINS, builder.field("city"),
          builder.literal("waltham"));
  builder.filter(filterRexNode);

  String builderExpected = ""
      + "LogicalFilter(condition=[CONTAINS($1, 'waltham')])\n"
      + "  LogicalProject(_MAP=[$0], city=[ITEM($0, 'city')])\n"
      + "    ElasticsearchTableScan(table=[[elastic, " + ZIPS + "]])\n";

  RelNode root = builder.build();

  RelRunner ru = (RelRunner) con.unwrap(Class.forName("org.apache.calcite.tools.RelRunner"));
  try (PreparedStatement preparedStatement = ru.prepare(root)) {

    String s = CalciteAssert.toString(preparedStatement.executeQuery());
    final String result = ""
        + "_MAP={id=02154, city=NORTH WALTHAM, loc=[-71.236497, 42.382492], "
        + "pop=57871, state=MA}; city=NORTH WALTHAM\n";

    // Validate query prepared
    assertThat(root, hasTree(builderExpected));

    // Validate result returned from ES
    assertThat(s, is(result));
  }
}
 

@Override
public RelNode visit(LogicalAggregate aggregate) {
  RelNode input = aggregate.getInput().accept(this);

  // Create a new project with null UPDATE_COLUMN below aggregate
  final RelBuilder relBuilder = newCalciteRelBuilderWithoutContext(aggregate.getCluster());
  relBuilder.push(input);
  List<RexNode> nodes = input.getRowType().getFieldList().stream().map(q -> {
    if (UPDATE_COLUMN.equals(q.getName())) {
      return relBuilder.getRexBuilder().makeNullLiteral(q.getType());
    } else{
      return relBuilder.getRexBuilder().makeInputRef(q.getType(), q.getIndex());
    }
  }).collect(Collectors.toList());
  relBuilder.project(nodes, input.getRowType().getFieldNames());

  // create a new aggregate with null UPDATE_COLUMN in groupSet
  RelDataType incomingRowType = relBuilder.peek().getRowType();
  RelDataTypeField modField = incomingRowType.getField(UPDATE_COLUMN, false, false);
  ImmutableBitSet newGroupSet = aggregate.getGroupSet().rebuild().set(modField.getIndex()).build();
  GroupKey groupKey = relBuilder.groupKey(newGroupSet, aggregate.indicator, null);

  final int groupCount = aggregate.getGroupCount();
  final Pointer<Integer> ind = new Pointer<>(groupCount-1);
  final List<String> fieldNames = aggregate.getRowType().getFieldNames();
  final List<AggregateCall> aggCalls = aggregate.getAggCallList().stream().map(q -> {
    ind.value++;
    if (q.getName() == null) {
      return q.rename(fieldNames.get(ind.value));
    }
    return q;
  }).collect(Collectors.toList());

  relBuilder.aggregate(groupKey, aggCalls);

  // create a new project on top to preserve rowType
  Iterable<RexInputRef> projects = FluentIterable.from(aggregate.getRowType().getFieldNames())
    .transform(new Function<String, RexInputRef>() {
      @Override
      public RexInputRef apply(String fieldName) {
        return relBuilder.field(fieldName);
      }
    })
    .append(relBuilder.field(UPDATE_COLUMN));

  relBuilder.project(projects);

  return relBuilder.build();
}
 

public RelNode updateScan(IncrementallyUpdateable scan) {
  RelNode newScan = scan;
  RelDataTypeField refreshField = scan.getRowType().getField(refreshColumn, false, false);

  if (refreshField == null) {
    // Check if the field exist as part of the schema
    newScan = scan.projectInvisibleColumn(refreshColumn);
    if (newScan == null) {
      throw UserException.dataReadError()
          .message("Table does not include column identified for incremental update of name '%s'.", refreshColumn)
          .build(logger);
    }
    refreshField = newScan.getRowType().getField(refreshColumn, false, false);
  }

  switch(refreshField.getType().getSqlTypeName()) {
    case INTEGER:
    case BIGINT:
    case FLOAT:
    case DOUBLE:
    case VARCHAR:
    case TIMESTAMP:
    case DATE:
    case DECIMAL:
      break;
    default:
      throw UserException.dataReadError()
          .message("Dremio only supports incremental column update on INTEGER, BIGINT, FLOAT, DOUBLE, VARCHAR, TIMESTAMP, DATE, and DECIMAL types. The identified column was of type %s.", refreshField.getType().getSqlTypeName())
          .build(logger);
  }

  // No need to add a project if field name is correct
  if (UPDATE_COLUMN.equals(refreshColumn)) {
    return newScan;
  }

  final RelBuilder relBuilder = newCalciteRelBuilderWithoutContext(newScan.getCluster());

  relBuilder.push(newScan);

  List<String> newFieldNames = ImmutableList.<String>builder().addAll(newScan.getRowType().getFieldNames()).add(UPDATE_COLUMN).build();

  Iterable<RexNode> projects = Stream.concat(
      scan.getRowType().getFieldNames().stream().map(relBuilder::field),
      Stream.of(refreshRex(relBuilder)))
      .collect(Collectors.toList());
  relBuilder.project(projects, newFieldNames);

  return relBuilder.build();
}
 

@Override
public RelNode visit(LogicalCorrelate correlate) {
  RelNode left = correlate.getLeft().accept(this);
  leftInputs.put(correlate.getCorrelationId(), left);

  RelNode right = correlate.getRight().accept(this);
  // if right input wasn't changed or left input wasn't changed
  // after rewriting right input, no need to create Correlate with new CorrelationId
  if (correlate.getRight() == right
      || left == leftInputs.get(correlate.getCorrelationId())) {
    if (correlate.getLeft() == left) {
      return correlate;
    }
    // changed only inputs, but CorrelationId left the same
    return correlate.copy(correlate.getTraitSet(), Arrays.asList(left, right));
  }

  Correlate newCorrelate = correlate.copy(correlate.getTraitSet(),
      leftInputs.get(correlate.getCorrelationId()), right,
      updatedCorrelationIds.get(correlate.getCorrelationId()),
      ImmutableBitSet.of(left.getRowType().getFieldCount()), correlate.getJoinType());

  RelBuilder builder = DrillRelFactories.LOGICAL_BUILDER.create(correlate.getCluster(), null);
  builder.push(newCorrelate);

  List<RexNode> topProjectExpressions = left.getRowType().getFieldList().stream()
          .map(field -> builder.getRexBuilder().makeInputRef(newCorrelate, field.getIndex()))
          .collect(Collectors.toList());

  //Accommodate the new $COMPLEX_FIELD_NAME column.
  int rightStartIndex = left.getRowType().getFieldList().size() + 1;
  switch (correlate.getJoinType()) {
    case LEFT:
    case INNER:
      // adds field from the right input of correlate to the top project
      topProjectExpressions.addAll(right.getRowType().getFieldList().stream()
              .map(field -> builder.getRexBuilder().makeInputRef(newCorrelate, field.getIndex() + rightStartIndex))
              .collect(Collectors.toList()));
      // fall through
    case ANTI:
    case SEMI:
      builder.project(topProjectExpressions, correlate.getRowType().getFieldNames());
  }
  return builder.build();
}
 

/**
 * Creates the topmost projection that will sit on top of the selected join
 * ordering. The projection needs to match the original join ordering. Also,
 * places any post-join filters on top of the project.
 *
 * @param multiJoin join factors being optimized
 * @param joinTree selected join ordering
 * @param fieldNames field names corresponding to the projection expressions
 *
 * @return created projection
 */
private RelNode createTopProject(
    RelBuilder relBuilder,
    LoptMultiJoin multiJoin,
    LoptJoinTree joinTree,
    List<String> fieldNames) {
  List<RexNode> newProjExprs = new ArrayList<>();
  RexBuilder rexBuilder =
      multiJoin.getMultiJoinRel().getCluster().getRexBuilder();

  // create a projection on top of the joins, matching the original
  // join order
  final List<Integer> newJoinOrder = joinTree.getTreeOrder();
  int nJoinFactors = multiJoin.getNumJoinFactors();
  List<RelDataTypeField> fields = multiJoin.getMultiJoinFields();

  // create a mapping from each factor to its field offset in the join
  // ordering
  final Map<Integer, Integer> factorToOffsetMap = new HashMap<>();
  for (int pos = 0, fieldStart = 0; pos < nJoinFactors; pos++) {
    factorToOffsetMap.put(newJoinOrder.get(pos), fieldStart);
    fieldStart +=
        multiJoin.getNumFieldsInJoinFactor(newJoinOrder.get(pos));
  }

  for (int currFactor = 0; currFactor < nJoinFactors; currFactor++) {
    // if the factor is the right factor in a removable self-join,
    // then where possible, remap references to the right factor to
    // the corresponding reference in the left factor
    Integer leftFactor = null;
    if (multiJoin.isRightFactorInRemovableSelfJoin(currFactor)) {
      leftFactor = multiJoin.getOtherSelfJoinFactor(currFactor);
    }
    for (int fieldPos = 0;
        fieldPos < multiJoin.getNumFieldsInJoinFactor(currFactor);
        fieldPos++) {
      int newOffset = factorToOffsetMap.get(currFactor) + fieldPos;
      if (leftFactor != null) {
        Integer leftOffset =
            multiJoin.getRightColumnMapping(currFactor, fieldPos);
        if (leftOffset != null) {
          newOffset =
              factorToOffsetMap.get(leftFactor) + leftOffset;
        }
      }
      newProjExprs.add(
          rexBuilder.makeInputRef(
              fields.get(newProjExprs.size()).getType(),
              newOffset));
    }
  }

  relBuilder.push(joinTree.getJoinTree());
  relBuilder.project(newProjExprs, fieldNames);

  // Place the post-join filter (if it exists) on top of the final
  // projection.
  RexNode postJoinFilter =
      multiJoin.getMultiJoinRel().getPostJoinFilter();
  if (postJoinFilter != null) {
    relBuilder.filter(postJoinFilter);
  }
  return relBuilder.build();
}