org.apache.calcite.rex.RexInputRef Java Examples

private static RexInputRef findRexInputRef(final RexNode node) {
  try {
    RexVisitor<Void> visitor =
        new RexVisitorImpl<Void>(true) {
          public Void visitCall(RexCall call) {
            for (RexNode child : call.getOperands()) {
              child.accept(this);
            }
            return super.visitCall(call);
          }

          public Void visitInputRef(RexInputRef inputRef) {
            throw new Util.FoundOne(inputRef);
          }
        };
    node.accept(visitor);
    return null;
  } catch (Util.FoundOne e) {
    Util.swallow(e, null);
    return (RexInputRef) e.getNode();
  }
}
 

/**
     * Translates a call to a binary operator. Returns whether successful.
     */
    private Pair<String, RexLiteral> translateBinary2(RexNode left, RexNode right) {
      switch (right.getKind()) {
        case LITERAL:
          break;
        default:
          return null;
      }

      final RexLiteral rightLiteral = (RexLiteral) right;
      switch (left.getKind()) {
        case INPUT_REF:
          final RexInputRef left1 = (RexInputRef) left;
          String name = fieldNames.get(left1.getIndex());
          return new Pair<>(name, rightLiteral);
        case CAST:
          return translateBinary2(((RexCall) left).operands.get(0), right);
//        case OTHER_FUNCTION:
//          String itemName = SolrRules.isItem((RexCall) left);
//          if (itemName != null) {
//            return translateOp2(op, itemName, rightLiteral);
//          }
        default:
          return null;
      }
    }
 

public static boolean isSimpleColumnSelection(Project project) {
  HashSet<Integer> inputRefReferenced = new HashSet<>();
  for (Pair<RexNode, String> proj : project.getNamedProjects()) {
    if (proj.getKey().getKind() != SqlKind.INPUT_REF) {
      return false;
    }
    RexInputRef inputRef = (RexInputRef) proj.getKey();
    // If the input reference is again referenced, then it is not a simple column selection (since it is not a permutation).
    if (inputRefReferenced.contains(inputRef.getIndex())) {
      return false;
    }
    final String nameOfProjectField = proj.getValue();
    final String nameOfInput = project.getInput().getRowType().getFieldNames().get(inputRef.getIndex());
    // Renaming a column is not a simple column selection
    if (nameOfProjectField == null || !nameOfProjectField.equals(nameOfInput)) {
      return false;
    }
    inputRefReferenced.add(inputRef.getIndex());
  }
  return true;
}
 

/**
 * Expression lineage from {@link TableScan}.
 *
 * <p>We extract the fields referenced by the expression and we express them
 * using {@link RexTableInputRef}.
 */
public Set<RexNode> getExpressionLineage(TableScan rel,
    RelMetadataQuery mq, RexNode outputExpression) {
  final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();

  // Extract input fields referenced by expression
  final ImmutableBitSet inputFieldsUsed = extractInputRefs(outputExpression);

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexNode inputRef = RexTableInputRef.of(
        RelTableRef.of(rel.getTable(), 0),
        RexInputRef.of(idx, rel.getRowType().getFieldList()));
    final RexInputRef ref = RexInputRef.of(idx, rel.getRowType().getFieldList());
    mapping.put(ref, ImmutableSet.of(inputRef));
  }

  // Return result
  return createAllPossibleExpressions(rexBuilder, outputExpression, mapping);
}
 

/** Check if join condition only references RexInputRef. */
private static boolean referenceByMapping(
    RexNode joinCondition, List<RexNode>... projectsOfInputs) {
  List<RexNode> projects = new ArrayList<>();
  for (List<RexNode> projectsOfInput: projectsOfInputs) {
    projects.addAll(projectsOfInput);
  }

  try {
    RexVisitor rexVisitor = new RexVisitorImpl<Void>(true) {
      @Override public Void visitInputRef(RexInputRef inputRef) {
        if (!(projects.get(inputRef.getIndex()) instanceof RexInputRef)) {
          throw Util.FoundOne.NULL;
        }
        return super.visitInputRef(inputRef);
      }
    };
    joinCondition.accept(rexVisitor);
  } catch (Util.FoundOne e) {
    return false;
  }
  return true;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
    InstanceByClassTableScan scan = call.rel(0);
    RelOptTable table = scan.getTable();
    RelOptSchema schema = table.getRelOptSchema();
    List<String> indexName = new ArrayList<String>(table.getQualifiedName());
    indexName.set(indexName.size() - 1, "$ids$:" + indexName.get(indexName.size() - 1));
    LogicalTableScan ids = LogicalTableScan.create(scan.getCluster(), schema.getTableForMember(indexName));

    InstanceByClassTable instanceByClassTable = table.unwrap(InstanceByClassTable.class);
    int snapshotId = SnapshotHolder.put(instanceByClassTable.snapshot);

    RelOptCluster cluster = scan.getCluster();
    RexInputRef objectId = cluster.getRexBuilder().makeInputRef(ids, 0);
    RexBuilderContext rexContext = new ExecutionRexBuilderContext(cluster, snapshotId, objectId);

    List<Function<RexBuilderContext, RexNode>> resolvers = instanceByClassTable.getResolvers();
    List<RexNode> exprs = new ArrayList<RexNode>(resolvers.size());
    for (Function<RexBuilderContext, RexNode> resolver : resolvers) {
        exprs.add(resolver.apply(rexContext));
    }
    call.transformTo(RelOptUtil.createProject(ids, exprs, table.getRowType().getFieldNames(), true));
}
 

@Override
public RexNode visitInputRef(RexInputRef var) {
    int srcIndex = var.getIndex();
    int destIndex = srcIndex + adjustments[srcIndex];

    RelDataType type;
    if (destFields != null) {
        type = destFields.get(destIndex).getType();
    } else if (leftDestFields != null) {
        if (destIndex < nLeftDestFields) {
            type = leftDestFields.get(destIndex).getType();
        } else {
            type = rightDestFields.get(destIndex - nLeftDestFields).getType();
        }
    } else {
        type = srcFields.get(srcIndex).getType();
    }
    if ((adjustments[srcIndex] != 0) || (srcFields == null) || (type != srcFields.get(srcIndex).getType())) {
        return rexBuilder.makeInputRef(type, destIndex);
    } else {
        return var;
    }
}
 

public Double averageRexSize(RexNode node, List<Double> inputColumnSizes) {
  switch (node.getKind()) {
  case INPUT_REF:
    return inputColumnSizes.get(((RexInputRef) node).getIndex());
  case LITERAL:
    return typeValueSize(node.getType(),
        ((RexLiteral) node).getValueAs(Comparable.class));
  default:
    if (node instanceof RexCall) {
      RexCall call = (RexCall) node;
      for (RexNode operand : call.getOperands()) {
        // It's a reasonable assumption that a function's result will have
        // similar size to its argument of a similar type. For example,
        // UPPER(c) has the same average size as c.
        if (operand.getType().getSqlTypeName()
            == node.getType().getSqlTypeName()) {
          return averageRexSize(operand, inputColumnSizes);
        }
      }
    }
    return averageTypeValueSize(node.getType());
  }
}
 

@Override public RelNode toRel(ToRelContext context) {
  RelNode rel = LogicalTableScan.create(context.getCluster(), fromTable,
      context.getTableHints());
  final RexBuilder rexBuilder = context.getCluster().getRexBuilder();
  rel = LogicalFilter.create(
      rel, getConstraint(rexBuilder, rel.getRowType()));
  final List<RelDataTypeField> fieldList =
      rel.getRowType().getFieldList();
  final List<Pair<RexNode, String>> projects =
      new AbstractList<Pair<RexNode, String>>() {
        @Override public Pair<RexNode, String> get(int index) {
          return RexInputRef.of2(mapping.get(index), fieldList);
        }

        @Override public int size() {
          return mapping.size();
        }
      };
  return LogicalProject.create(rel,
      ImmutableList.of(),
      Pair.left(projects),
      Pair.right(projects));
}
 

@Test void testExpressionLineageTwoColumnsSwapped() {
  // deptno is column 7 in catalog.sales.emp
  // mgr is column 3 in catalog.sales.emp
  final RelNode rel = convertSql("select deptno, mgr from emp");
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();

  final RexNode ref1 = RexInputRef.of(0, rel.getRowType().getFieldList());
  final Set<RexNode> r1 = mq.getExpressionLineage(rel, ref1);
  assertThat(r1.size(), is(1));
  final RexTableInputRef result1 = (RexTableInputRef) r1.iterator().next();
  assertThat(result1.getQualifiedName(), is(EMP_QNAME));
  assertThat(result1.getIndex(), is(7));

  final RexNode ref2 = RexInputRef.of(1, rel.getRowType().getFieldList());
  final Set<RexNode> r2 = mq.getExpressionLineage(rel, ref2);
  assertThat(r2.size(), is(1));
  final RexTableInputRef result2 = (RexTableInputRef) r2.iterator().next();
  assertThat(result2.getQualifiedName(), is(EMP_QNAME));
  assertThat(result2.getIndex(), is(3));

  assertThat(result1.getIdentifier(), is(result2.getIdentifier()));
}
 

public static DataContext of(RelDataType rowType,
    List<Pair<RexInputRef, RexNode>> usageList) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  for (Pair<RexInputRef, RexNode> elem : usageList) {
    Pair<Integer, ?> value = getValue(elem.getKey(), elem.getValue());
    if (value == null) {
      LOGGER.warn("{} is not handled for {} for checking implication",
          elem.getKey(), elem.getValue());
      return null;
    }
    int index = value.getKey();
    values[index] = value.getValue();
  }
  return new VisitorDataContext(values);
}
 

void translateJoinColumn(RexCall condition, Multimap<TblColRef, TblColRef> joinColumns) {
    SqlKind kind = condition.getOperator().getKind();
    if (kind == SqlKind.AND) {
        for (RexNode operand : condition.getOperands()) {
            RexCall subCond = (RexCall) operand;
            translateJoinColumn(subCond, joinColumns);
        }
    } else if (kind == SqlKind.EQUALS) {
        List<RexNode> operands = condition.getOperands();
        RexInputRef op0 = (RexInputRef) operands.get(0);
        TblColRef col0 = columnRowType.getColumnByIndex(op0.getIndex());
        RexInputRef op1 = (RexInputRef) operands.get(1);
        TblColRef col1 = columnRowType.getColumnByIndex(op1.getIndex());
        // map left => right
        if (op0.getIndex() < columnRowTypeLeftRightCut)
            joinColumns.put(col0, col1);
        else
            joinColumns.put(col1, col0);
    }
}
 

/**
 *  Get the field name for the left node to use for IN SET query
 */
private String getLeftNodeFieldName(RexNode left) {
  switch (left.getKind()) {
  case INPUT_REF:
    final RexInputRef left1 = (RexInputRef) left;
    return fieldNames.get(left1.getIndex());
  case CAST:
    // FIXME This will not work in all cases (for example, we ignore string encoding)
    return getLeftNodeFieldName(((RexCall) left).operands.get(0));
  case ITEM:
  case OTHER_FUNCTION:
    return left.accept(new GeodeRules.RexToGeodeTranslator(this.fieldNames));
  default:
    return null;
  }
}
 

private static RelFieldCollation collation(RexNode node,
    RelFieldCollation.Direction direction,
    RelFieldCollation.NullDirection nullDirection, List<RexNode> extraNodes) {
  switch (node.getKind()) {
  case INPUT_REF:
    return new RelFieldCollation(((RexInputRef) node).getIndex(), direction,
        Util.first(nullDirection, direction.defaultNullDirection()));
  case DESCENDING:
    return collation(((RexCall) node).getOperands().get(0),
        RelFieldCollation.Direction.DESCENDING,
        nullDirection, extraNodes);
  case NULLS_FIRST:
    return collation(((RexCall) node).getOperands().get(0), direction,
        RelFieldCollation.NullDirection.FIRST, extraNodes);
  case NULLS_LAST:
    return collation(((RexCall) node).getOperands().get(0), direction,
        RelFieldCollation.NullDirection.LAST, extraNodes);
  default:
    final int fieldIndex = extraNodes.size();
    extraNodes.add(node);
    return new RelFieldCollation(fieldIndex, direction,
        Util.first(nullDirection, direction.defaultNullDirection()));
  }
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	LogicalJoin join = call.rel(0);
	LogicalProject project = call.rel(1);

	// only accept SEMI/ANTI join
	JoinRelType joinType = join.getJoinType();
	if (joinType != JoinRelType.SEMI && joinType != JoinRelType.ANTI) {
		return false;
	}
	// all expressions in Project should be RexInputRef
	for (RexNode p : project.getProjects()) {
		if (!(p instanceof RexInputRef)) {
			return false;
		}
	}
	return true;
}
 

@Override public RexNode visitInputRef(RexInputRef input) {
  final int oldIndex = input.getIndex();
  final Ord<RelDataType> field = getNewFieldForOldInput(oldIndex);
  RelDataTypeField inputFieldByOldIndex = currentRel.getInputs().stream()
      .flatMap(relInput -> relInput.getRowType().getFieldList().stream())
      .skip(oldIndex)
      .findFirst()
      .orElseThrow(() ->
          new AssertionError("Found input ref with index not found in old inputs"));
  if (inputFieldByOldIndex.getType().isStruct()) {
    iRestructureInput = field.i;
    List<RexNode> rexNodes = restructureFields(inputFieldByOldIndex.getType());
    return rexBuilder.makeCall(
        inputFieldByOldIndex.getType(),
        SqlStdOperatorTable.ROW,
        rexNodes);
  }

  // Use the actual flattened type, which may be different from the current
  // type.
  RelDataType fieldType = removeDistinct(field.e);
  return new RexInputRef(field.i, fieldType);
}
 

/**
 * Test {@link RelOptUtil#splitJoinCondition(RelNode, RelNode, RexNode, List, List, List)}
 * where the join condition contains an expanded version of IS NOT DISTINCT using CASE
 */
@Test void testSplitJoinConditionExpandedIsNotDistinctFromUsingCase() {
  int leftJoinIndex = empScan.getRowType().getFieldNames().indexOf("DEPTNO");
  int rightJoinIndex = deptRow.getFieldNames().indexOf("DEPTNO");

  RexInputRef leftKeyInputRef = RexInputRef.of(leftJoinIndex, empDeptJoinRelFields);
  RexInputRef rightKeyInputRef =
      RexInputRef.of(empRow.getFieldCount() + rightJoinIndex, empDeptJoinRelFields);
  RexNode joinCond = RelOptUtil.isDistinctFrom(
      relBuilder.getRexBuilder(),
      leftKeyInputRef,
      rightKeyInputRef,
      true);


  splitJoinConditionHelper(
      joinCond,
      Collections.singletonList(leftJoinIndex),
      Collections.singletonList(rightJoinIndex),
      Collections.singletonList(false),
      relBuilder.literal(true));
}
 

private boolean isBooleanColumnReference(RexNode node, List<String> fieldNames) {
  // FIXME Ignore casts for rel and assume they aren't really necessary
  if (node.isA(SqlKind.CAST)) {
    node = ((RexCall) node).getOperands().get(0);
  }
  if (node.isA(SqlKind.NOT)) {
    node = ((RexCall) node).getOperands().get(0);
  }
  if (node.isA(SqlKind.INPUT_REF)) {
    if (node.getType().getSqlTypeName() == SqlTypeName.BOOLEAN) {
      final RexInputRef left1 = (RexInputRef) node;
      String name = fieldNames.get(left1.getIndex());
      return name != null;
    }
  }
  return false;
}
 

@Override public RexNode visitFieldAccess(RexFieldAccess fieldAccess) {
  if (cm.mapFieldAccessToCorRef.containsKey(fieldAccess)) {
    // if it is a corVar, change it to be input ref.
    CorRef corVar = cm.mapFieldAccessToCorRef.get(fieldAccess);

    // corVar offset should point to the leftInput of currentRel,
    // which is the Correlate.
    RexNode newRexNode =
        new RexInputRef(corVar.field, fieldAccess.getType());

    if (projectPulledAboveLeftCorrelator
        && (nullIndicator != null)) {
      // need to enforce nullability by applying an additional
      // cast operator over the transformed expression.
      newRexNode =
          createCaseExpression(nullIndicator, null, newRexNode);
    }
    return newRexNode;
  }
  return fieldAccess;
}
 

public static List<RexNode> identityProjects(RelDataType type, ImmutableBitSet selectedColumns) {
  List<RexNode> projects = new ArrayList<>();
  if (selectedColumns == null) {
    selectedColumns = ImmutableBitSet.range(type.getFieldCount());
  }
  for (Pair<Integer,RelDataTypeField> pair : Pair.zip(selectedColumns, type.getFieldList())) {
    projects.add(new RexInputRef(pair.left, pair.right.getType()));
  }
  return projects;
}
 

private void reduceNotNullableFilter(
    RelOptRuleCall call,
    Filter filter,
    RexNode rexNode,
    boolean reverse) {
  // If the expression is a IS [NOT] NULL on a non-nullable
  // column, then we can either remove the filter or replace
  // it with an Empty.
  boolean alwaysTrue;
  switch (rexNode.getKind()) {
  case IS_NULL:
  case IS_UNKNOWN:
    alwaysTrue = false;
    break;
  case IS_NOT_NULL:
    alwaysTrue = true;
    break;
  default:
    return;
  }
  if (reverse) {
    alwaysTrue = !alwaysTrue;
  }
  RexNode operand = ((RexCall) rexNode).getOperands().get(0);
  if (operand instanceof RexInputRef) {
    RexInputRef inputRef = (RexInputRef) operand;
    if (!inputRef.getType().isNullable()) {
      if (alwaysTrue) {
        call.transformTo(filter.getInput());
      } else {
        call.transformTo(createEmptyRelOrEquivalent(call, filter));
      }
      // New plan is absolutely better than old plan.
      call.getPlanner().prune(filter);
    }
  }
}
 

private int[] getProjectFields(List<RexNode> exps) {
    final int[] fields = new int[exps.size()];

    for (int i = 0; i < exps.size(); i++) {
        final RexNode exp = exps.get(i);

        if (exp instanceof RexInputRef) {
            fields[i] = ((RexInputRef) exp).getIndex();
        } else {
            return null; // not a simple projection
        }
    }

    return fields;
}
 

private void updateUsage(SqlOperator op, RexInputRef inputRef,
    RexNode literal) {
  final InputRefUsage<SqlOperator, RexNode> inputRefUse =
      getUsageMap(inputRef);
  Pair<SqlOperator, RexNode> use = Pair.of(op, literal);
  inputRefUse.usageList.add(use);
}
 

RemoveCorrelationRexShuttle(RexBuilder rexBuilder, boolean projectPulledAboveLeftCorrelator,
        RexInputRef nullIndicator, Set<Integer> isCount) {
    this.projectPulledAboveLeftCorrelator = projectPulledAboveLeftCorrelator;
    this.nullIndicator = nullIndicator; // may be null
    this.isCount = ImmutableSet.copyOf(isCount);
    this.rexBuilder = rexBuilder;
    this.typeFactory = rexBuilder.getTypeFactory();
}
 

@Test void testCalc() {
  final FrameworkConfig config = RelBuilderTest.config().build();
  final RelBuilder builder = RelBuilder.create(config);
  final RexBuilder rexBuilder = builder.getRexBuilder();
  final LogicalTableScan scan = (LogicalTableScan) builder.scan("EMP").build();
  final RexProgramBuilder programBuilder =
      new RexProgramBuilder(scan.getRowType(), rexBuilder);
  final RelDataTypeField field = scan.getRowType().getField("SAL", false, false);
  programBuilder.addIdentity();
  programBuilder.addCondition(
      rexBuilder.makeCall(SqlStdOperatorTable.GREATER_THAN,
          new RexInputRef(field.getIndex(), field.getType()),
          builder.literal(10)));
  final LogicalCalc calc = LogicalCalc.create(scan, programBuilder.getProgram());
  String relJson = RelOptUtil.dumpPlan("", calc,
      SqlExplainFormat.JSON, SqlExplainLevel.EXPPLAN_ATTRIBUTES);
  String s =
      Frameworks.withPlanner((cluster, relOptSchema, rootSchema) -> {
        final RelJsonReader reader = new RelJsonReader(
            cluster, getSchema(calc), rootSchema);
        RelNode node;
        try {
          node = reader.read(relJson);
        } catch (IOException e) {
          throw TestUtil.rethrow(e);
        }
        return RelOptUtil.dumpPlan("", node, SqlExplainFormat.TEXT,
            SqlExplainLevel.EXPPLAN_ATTRIBUTES);
      });
  final String expected =
      "LogicalCalc(expr#0..7=[{inputs}], expr#8=[10], expr#9=[>($t5, $t8)],"
          + " proj#0..7=[{exprs}], $condition=[$t9])\n"
          + "  LogicalTableScan(table=[[scott, EMP]])\n";
  assertThat(s, isLinux(expected));
}
 

private int[] getProjectFields(List<RexNode> exps) {
    final int[] fields = new int[exps.size()];

    for (int i = 0; i < exps.size(); i++) {
        final RexNode exp = exps.get(i);

        if (exp instanceof RexInputRef) {
            fields[i] = ((RexInputRef) exp).getIndex();
        } else {
            return null; // not a simple projection
        }
    }

    return fields;
}
 

private Map<Integer, Integer> getCollationMap(DrillProjectRel project) {
  Map<Integer, Integer> m = new HashMap<>();

  for (Ord<RexNode> node : Ord.zip(project.getProjects())) {
    // For collation, only $0 will keep the sort-ness after projection.
    if (node.e instanceof RexInputRef) {
      m.put( ((RexInputRef) node.e).getIndex(), node.i);
    }
  }
  return m;

}
 

/**
 * For each field creates row expression.
 *
 * @param rowType row type
 * @return list of row expressions
 */
public static List<RexNode> prepareFieldExpressions(RelDataType rowType) {
  List<RexNode> expressions = new ArrayList<>();
  for (int i = 0; i < rowType.getFieldCount(); i++) {
    expressions.add(RexInputRef.of(i, rowType));
  }
  return expressions;
}
 

/**
 * Pulls a {@link Project} above a {@link Correlate} from its RHS input.
 * Enforces nullability for join output.
 *
 * @param correlate  Correlate
 * @param project the original project as the RHS input of the join
 * @param isCount Positions which are calls to the <code>COUNT</code>
 *                aggregation function
 * @return the subtree with the new Project at the root
 */
private RelNode aggregateCorrelatorOutput(
    Correlate correlate,
    Project project,
    Set<Integer> isCount) {
  final RelNode left = correlate.getLeft();
  final JoinRelType joinType = correlate.getJoinType();

  // now create the new project
  final List<Pair<RexNode, String>> newProjects = new ArrayList<>();

  // Project everything from the LHS and then those from the original
  // project
  final List<RelDataTypeField> leftInputFields =
      left.getRowType().getFieldList();

  for (int i = 0; i < leftInputFields.size(); i++) {
    newProjects.add(RexInputRef.of2(i, leftInputFields));
  }

  // Marked where the projected expr is coming from so that the types will
  // become nullable for the original projections which are now coming out
  // of the nullable side of the OJ.
  boolean projectPulledAboveLeftCorrelator =
      joinType.generatesNullsOnRight();

  for (Pair<RexNode, String> pair : project.getNamedProjects()) {
    RexNode newProjExpr =
        removeCorrelationExpr(
            pair.left,
            projectPulledAboveLeftCorrelator,
            isCount);
    newProjects.add(Pair.of(newProjExpr, pair.right));
  }

  return relBuilder.push(correlate)
      .projectNamed(Pair.left(newProjects), Pair.right(newProjects), true)
      .build();
}
 

@Override
public RexNode visitInputRef(RexInputRef inputRef) {
    Integer index = mapper.get(inputRef);
    if (index != null) {
        return RexBuilder.getRexFactory().makeInputRef(index, inputRef.getType());
    }
    return super.visitInputRef(inputRef);
}