Java Code Examples for org.apache.calcite.rex.RexNode#accept()

@Test
public void equalityWithCast(){

  final RexNode node = builder.makeCall(SqlStdOperatorTable.EQUALS,
      builder.makeCast(
          factory.createSqlType(SqlTypeName.ANY),
          builder.makeBigintLiteral(BigDecimal.ONE)
      ),
      builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 0)
  );

  FindSimpleFilters finder = new FindSimpleFilters(builder);
  StateHolder holder = node.accept(finder);
  ImmutableList<RexCall> conditions = holder.getConditions();

  assertEquals(1, conditions.size());
  assertEquals(SqlKind.EQUALS, conditions.get(0).getKind());
  // Make sure CAST was removed
  assertEquals(SqlKind.LITERAL, conditions.get(0).getOperands().get(0).getKind());
  assertFalse(holder.hasRemainingExpression());

}
 

@Test
public void doubleAnd(){
  final RexNode node =
      builder.makeCall(SqlStdOperatorTable.AND,
      builder.makeCall(SqlStdOperatorTable.EQUALS,
          builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 0),
          builder.makeBigintLiteral(BigDecimal.ONE)
          ),
      builder.makeCall(SqlStdOperatorTable.EQUALS,
          builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 0),
          builder.makeBigintLiteral(BigDecimal.ONE)
          )
      );

  FindSimpleFilters finder = new FindSimpleFilters(builder);
  StateHolder holder = node.accept(finder);
  ImmutableList<RexCall> conditions = holder.getConditions();

  assertEquals(2, conditions.size());
  assertEquals(SqlKind.EQUALS, conditions.get(0).getKind());
  assertEquals(SqlKind.EQUALS, conditions.get(1).getKind());
  assertFalse(holder.hasRemainingExpression());
}
 

/**
 * Shifts a filter originating from the right child of the LogicalJoin to the
 * right, to reflect the filter now being applied on the resulting
 * MultiJoin.
 *
 * @param joinRel     the original LogicalJoin
 * @param left        the left child of the LogicalJoin
 * @param right       the right child of the LogicalJoin
 * @param rightFilter the filter originating from the right child
 * @return the adjusted right filter
 */
private RexNode shiftRightFilter(
    Join joinRel,
    RelNode left,
    MultiJoin right,
    RexNode rightFilter) {
  if (rightFilter == null) {
    return null;
  }

  int nFieldsOnLeft = left.getRowType().getFieldList().size();
  int nFieldsOnRight = right.getRowType().getFieldList().size();
  int[] adjustments = new int[nFieldsOnRight];
  for (int i = 0; i < nFieldsOnRight; i++) {
    adjustments[i] = nFieldsOnLeft;
  }
  rightFilter =
      rightFilter.accept(
          new RelOptUtil.RexInputConverter(
              joinRel.getCluster().getRexBuilder(),
              right.getRowType().getFieldList(),
              joinRel.getRowType().getFieldList(),
              adjustments));
  return rightFilter;
}
 

protected RexNode removeCorrelationExpr(
    RexNode exp,
    boolean projectPulledAboveLeftCorrelator,
    Set<Integer> isCount) {
  RemoveCorrelationRexShuttle shuttle =
      new RemoveCorrelationRexShuttle(relBuilder.getRexBuilder(),
          projectPulledAboveLeftCorrelator, null, isCount);
  return exp.accept(shuttle);
}
 

private static boolean userDefinedFunctionInProject(Project project) {
  CheckingUserDefinedFunctionVisitor visitor = new CheckingUserDefinedFunctionVisitor();
  for (RexNode node : project.getProjects()) {
    node.accept(visitor);
    if (visitor.containsUserDefinedFunction()) {
      return true;
    }
  }
  return false;
}
 

/**
 * Adjusts a filter to reflect swapping of join inputs
 *
 * @param rexBuilder rexBuilder
 * @param multiJoin join factors being optimized
 * @param origLeft original LHS of the join tree (before swap)
 * @param origRight original RHS of the join tree (before swap)
 * @param condition original join condition
 *
 * @return join condition reflect swap of join inputs
 */
private RexNode swapFilter(
    RexBuilder rexBuilder,
    LoptMultiJoin multiJoin,
    LoptJoinTree origLeft,
    LoptJoinTree origRight,
    RexNode condition) {
  int nFieldsOnLeft =
      origLeft.getJoinTree().getRowType().getFieldCount();
  int nFieldsOnRight =
      origRight.getJoinTree().getRowType().getFieldCount();
  int [] adjustments = new int[nFieldsOnLeft + nFieldsOnRight];

  for (int i = 0; i < nFieldsOnLeft; i++) {
    adjustments[i] = nFieldsOnRight;
  }
  for (int i = nFieldsOnLeft; i < (nFieldsOnLeft + nFieldsOnRight); i++) {
    adjustments[i] = -nFieldsOnLeft;
  }

  condition =
      condition.accept(
          new RelOptUtil.RexInputConverter(
              rexBuilder,
              multiJoin.getJoinFields(origLeft, origRight),
              multiJoin.getJoinFields(origRight, origLeft),
              adjustments));

  return condition;
}
 

@Override public RelNode visit(RelNode other) {
  if (other instanceof Join) {
    Join join = (Join) other;
    try {
      stack.push(join);
      join.getCondition().accept(rexVisitor(join));
    } finally {
      stack.pop();
    }
    return visitJoin(join);
  } else if (other instanceof Correlate) {
    Correlate correlate = (Correlate) other;
    mapCorToCorRel.put(correlate.getCorrelationId(), correlate);
    return visitJoin(correlate);
  } else if (other instanceof Filter) {
    Filter filter = (Filter) other;
    try {
      stack.push(filter);
      filter.getCondition().accept(rexVisitor(filter));
    } finally {
      stack.pop();
    }
  } else if (other instanceof Project) {
    Project project = (Project) other;
    try {
      stack.push(project);
      for (RexNode node : project.getProjects()) {
        node.accept(rexVisitor(project));
      }
    } finally {
      stack.pop();
    }
  }
  return super.visit(other);
}
 

/**
 * Determines if the equality join filters between two factors that map to
 * the same table consist of unique, identical keys.
 *
 * @param multiJoin join factors being optimized
 * @param leftFactor left factor in the join
 * @param rightFactor right factor in the join
 * @param joinFilterList list of join filters between the two factors
 *
 * @return true if the criteria are met
 */
private boolean isSelfJoinFilterUnique(
    RelMetadataQuery mq,
    LoptMultiJoin multiJoin,
    int leftFactor,
    int rightFactor,
    List<RexNode> joinFilterList) {
  RexBuilder rexBuilder =
      multiJoin.getMultiJoinRel().getCluster().getRexBuilder();
  RelNode leftRel = multiJoin.getJoinFactor(leftFactor);
  RelNode rightRel = multiJoin.getJoinFactor(rightFactor);
  RexNode joinFilters =
      RexUtil.composeConjunction(rexBuilder, joinFilterList, true);

  // Adjust the offsets in the filter by shifting the left factor
  // to the left and shifting the right factor to the left and then back
  // to the right by the number of fields in the left
  int [] adjustments = new int[multiJoin.getNumTotalFields()];
  int leftAdjust = multiJoin.getJoinStart(leftFactor);
  int nLeftFields = leftRel.getRowType().getFieldCount();
  for (int i = 0; i < nLeftFields; i++) {
    adjustments[leftAdjust + i] = -leftAdjust;
  }
  int rightAdjust = multiJoin.getJoinStart(rightFactor);
  for (int i = 0; i < rightRel.getRowType().getFieldCount(); i++) {
    adjustments[rightAdjust + i] = -rightAdjust + nLeftFields;
  }
  joinFilters =
      joinFilters.accept(
          new RelOptUtil.RexInputConverter(
              rexBuilder,
              multiJoin.getMultiJoinFields(),
              leftRel.getRowType().getFieldList(),
              rightRel.getRowType().getFieldList(),
              adjustments));

  return areSelfJoinKeysUnique(mq, leftRel, rightRel, joinFilters);
}
 

public static ProjectPushInfo getColumns(RelDataType rowType, List<RexNode> projects) {
  final List<String> fieldNames = rowType.getFieldNames();
  if (fieldNames.isEmpty()) {
    return null;
  }

  RefFieldsVisitor v = new RefFieldsVisitor(rowType);
  for (RexNode exp : projects) {
    PathSegment segment = exp.accept(v);
    v.addColumn(segment);
  }

  return v.getInfo();

}
 

/**
 * Given a list of Index Expressions(usually indexed fields/functions from one or a set of indexes),
 * separate a filter condition into
 *     1), relevant subset of conditions (by relevant, it means at least one given index Expression was found) and,
 *     2), the rest in remainderCondition
 * @param relevantPaths
 * @param condition
 * @return
 */
public IndexConditionInfo indexConditionRelatedToFields(List<LogicalExpression> relevantPaths, RexNode condition) {
  // Use the same filter analyzer that is used for partitioning columns
  RewriteCombineBinaryOperators reverseVisitor =
      new RewriteCombineBinaryOperators(true, builder);

  condition = condition.accept(reverseVisitor);

  RexSeparator separator = new RexSeparator(relevantPaths, scan, builder);
  RexNode indexCondition = separator.getSeparatedCondition(condition);

  if (indexCondition == null) {
    return new IndexConditionInfo(null, null, false);
  }

  List<RexNode> conjuncts = RelOptUtil.conjunctions(condition);
  List<RexNode> indexConjuncts = RelOptUtil.conjunctions(indexCondition);
  for (RexNode indexConjunction: indexConjuncts) {
    RexUtil.removeAll(conjuncts, indexConjunction);
  }

  RexNode remainderCondition = RexUtil.composeConjunction(builder, conjuncts, false);

  indexCondition = indexCondition.accept(reverseVisitor);

  return new IndexConditionInfo(indexCondition, remainderCondition, true);
}
 

RexNode visit(RexNode e) {
  final Integer i = map.get(e);
  if (i != null) {
    return relBuilder.field(i);
  }
  return e.accept(this);
}
 

@Test
public void multipleConditions() {
  final RexNode node = builder.makeCall(
    SqlStdOperatorTable.AND,
    builder.makeCall(
      SqlStdOperatorTable.GREATER_THAN_OR_EQUAL,
      builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 0),
      builder.makeBigintLiteral(new BigDecimal("1"))
    ),
    builder.makeCall(
      SqlStdOperatorTable.AND,
      builder.makeCall(
        SqlStdOperatorTable.LESS_THAN_OR_EQUAL,
        builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 0),
        builder.makeBigintLiteral(new BigDecimal("1"))
      ),
      builder.makeCall(
        SqlStdOperatorTable.LIKE,
        builder.makeInputRef(factory.createSqlType(SqlTypeName.BIGINT), 1),
        builder.makeLiteral("%1mg")
      )
    )
  );
  FindSimpleFilters finder = new FindSimpleFilters((builder));
  StateHolder holder = node.accept(finder);
  assertEquals(holder.getConditions().size(), 2);
  assertTrue(holder.hasRemainingExpression());
}
 

@Override
public FunctionRender render(FunctionRenderer renderer, RexCall call) {
  checkArity(call, 1);
  RexNode input = call.getOperands().get(0);
  FunctionRender op1 = input.accept(renderer.getVisitor());
  return new FunctionRender(getCastScript(input.getType(), call.getType(), op1.getScript(), renderer.isUsingPainless()), op1.getNulls());
}
 

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
  final RelNode input = agg.getInput();
  final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> aggPullUpPredicates = new ArrayList<>();

  ImmutableBitSet groupKeys = agg.getGroupSet();
  if (groupKeys.isEmpty()) {
    // "GROUP BY ()" can convert an empty relation to a non-empty relation, so
    // it is not valid to pull up predicates. In particular, consider the
    // predicate "false": it is valid on all input rows (trivially - there are
    // no rows!) but not on the output (there is one row).
    return RelOptPredicateList.EMPTY;
  }
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());

  int i = 0;
  for (int j : groupKeys) {
    m.set(j, i++);
  }

  for (RexNode r : inputInfo.pulledUpPredicates) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
    if (groupKeys.contains(rCols)) {
      r = r.accept(new RexPermuteInputsShuttle(m, input));
      aggPullUpPredicates.add(r);
    }
  }
  return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}
 

protected RexNode removeCorrelationExpr(RexNode exp, boolean projectPulledAboveLeftCorrelator,
        RexInputRef nullIndicator) {
    RemoveCorrelationRexShuttle shuttle = new RemoveCorrelationRexShuttle(relBuilder.getRexBuilder(),
            projectPulledAboveLeftCorrelator, nullIndicator, ImmutableSet.of());
    return exp.accept(shuttle);
}
 

public RexNode copyOf(RexNode expr) {
  return expr == null ? null : expr.accept(new RexCopierWithOver(cluster.getRexBuilder()));
}
 

public Set<String> go(List<RexNode> rexNodeList) {
    for (RexNode rex : rexNodeList) {
        rex.accept(this);
    }
    return patternVars;
}
 

@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));
}
 

/**
 * Returns whether a given expression is strong.
 *
 * <p>Examples:</p>
 * <ul>
 *   <li>Returns true for {@code c = 1} since it returns null if and only if
 *   c is null
 *   <li>Returns false for {@code c IS NULL} since it always returns TRUE
 *   or FALSE
 *</ul>
 *
 * @param e Expression
 * @return true if the expression is strong, false otherwise
 */
public static boolean isStrong(RexNode e) {
  final ImmutableBitSet.Builder nullColumns = ImmutableBitSet.builder();
  e.accept(
      new RexVisitorImpl<Void>(true) {
        public Void visitInputRef(RexInputRef inputRef) {
          nullColumns.set(inputRef.getIndex());
          return super.visitInputRef(inputRef);
        }
      });
  return isNull(e, nullColumns.build());
}
 

/**
 * Given a subtree, return an updated rexnode tree that replaces item operatos
 * and input refs with a schema aware node. If conversion fails, an
 * IllegalStateException is thrown as one or more ITEM operators are expected
 * to always be directly above input refs
 *
 * @param node
 *          Node to convert
 * @param scan
 *          Elastic scan that is below this expression.
 * @return A new node where RexInputRefs and ITEM operations are all converted
 *         to RexInputRef subclass SchemaField.
 */
public static RexNode convert(RexNode node, ElasticIntermediateScanPrel scan){
  return node.accept(new SchemaingShuttle(scan, ImmutableSet.<ElasticSpecialType>of()));
}