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

private static void populateEquivalences(Map<Integer, BitSet> equivalence,
		RexNode predicate) {
	switch (predicate.getKind()) {
		case EQUALS:
			RexCall call = (RexCall) predicate;
			final List<RexNode> operands = call.getOperands();
			if (operands.get(0) instanceof RexInputRef) {
				final RexInputRef ref0 = (RexInputRef) operands.get(0);
				if (operands.get(1) instanceof RexInputRef) {
					final RexInputRef ref1 = (RexInputRef) operands.get(1);
					populateEquivalence(equivalence, ref0.getIndex(), ref1.getIndex());
					populateEquivalence(equivalence, ref1.getIndex(), ref0.getIndex());
				}
			}
	}
}
 

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

/** Translate a binary relation. */
private String translateMatch2(RexNode node) {
  // We currently only use equality, but inequalities on clustering keys
  // should be possible in the future
  switch (node.getKind()) {
  case EQUALS:
    return translateBinary("=", "=", (RexCall) node);
  case LESS_THAN:
    return translateBinary("<", ">", (RexCall) node);
  case LESS_THAN_OR_EQUAL:
    return translateBinary("<=", ">=", (RexCall) node);
  case GREATER_THAN:
    return translateBinary(">", "<", (RexCall) node);
  case GREATER_THAN_OR_EQUAL:
    return translateBinary(">=", "<=", (RexCall) node);
  default:
    throw new AssertionError("cannot translate " + node);
  }
}
 

@Override public RexNode visitCall(RexCall call) {
  if (call.getOperator().getKind() == SqlKind.NOT) {
    RexNode child = call.getOperands().get(0);
    if (child.getKind() == SqlKind.LIKE) {
      return rexBuilder.makeCall(SqlStdOperatorTable.NOT_LIKE,
          visitList(((RexCall) child).getOperands()));
    }
  }
  return super.visitCall(call);
}
 

private RexNode createCollation(
		RexNode node,
		RelFieldCollation.Direction direction,
		RelFieldCollation.NullDirection nullDirection,
		Set<SqlKind> kinds) {
	switch (node.getKind()) {
		case DESCENDING:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), RelFieldCollation.Direction.DESCENDING,
				nullDirection, kinds);
		case NULLS_FIRST:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), direction,
				RelFieldCollation.NullDirection.FIRST, kinds);
		case NULLS_LAST:
			kinds.add(node.getKind());
			return createCollation(((RexCall) node).getOperands().get(0), direction,
				RelFieldCollation.NullDirection.LAST, kinds);
		default:
			if (nullDirection == null) {
				// Set the null direction if not specified.
				// Consistent with HIVE/SPARK/MYSQL/BLINK-RUNTIME.
				if (FlinkPlannerImpl.defaultNullCollation()
					.last(direction.equals(RelFieldCollation.Direction.DESCENDING))) {
					kinds.add(SqlKind.NULLS_LAST);
				} else {
					kinds.add(SqlKind.NULLS_FIRST);
				}
			}
			return node;
	}
}
 

private boolean translateBinary(SqlKind operator, RexNode left, RexNode right) {
    if (right.getKind() != SqlKind.LITERAL) return false;
    final RexLiteral rightLiteral = (RexLiteral) right;
    switch (left.getKind())
    {
        case INPUT_REF:
            String name = getRowType().getFieldNames().get(((RexInputRef) left).getIndex());
            translateOp(operator, name, rightLiteral);
            return true;
        case CAST:
            return translateBinary(operator, ((RexCall) left).operands.get(0), right);
        default: return false;
    }
}
 

boolean isFloorCeilCall(RexNode e) {
  switch (e.getKind()) {
  case FLOOR:
  case CEIL:
    final RexCall call = (RexCall) e;
    return call.getOperands().size() == 2;
  default:
    return false;
  }
}
 

/** Removes cast from string.
 *
 * <p>For example, {@code x > CAST('2015-01-07' AS DATE)}
 * becomes {@code x > '2015-01-07'}.
 */
private static RexNode stripCastFromString(RexNode node) {
    switch (node.getKind()) {
    case EQUALS:
    case IS_NOT_DISTINCT_FROM:
    case NOT_EQUALS:
    case GREATER_THAN:
    case GREATER_THAN_OR_EQUAL:
    case LESS_THAN:
    case LESS_THAN_OR_EQUAL:
        final RexCall call = (RexCall) node;
        final RexNode o0 = call.getOperands().get(0);
        final RexNode o1 = call.getOperands().get(1);
        if (o0.getKind() == SqlKind.CAST && o1.getKind() != SqlKind.CAST) {
            final RexNode o0b = ((RexCall) o0).getOperands().get(0);
            switch (o0b.getType().getSqlTypeName()) {
            case CHAR:
            case VARCHAR:
                return call.clone(call.getType(), ImmutableList.of(o0b, o1));
            }
        }
        if (o1.getKind() == SqlKind.CAST && o0.getKind() != SqlKind.CAST) {
            final RexNode o1b = ((RexCall) o1).getOperands().get(0);
            switch (o1b.getType().getSqlTypeName()) {
            case CHAR:
            case VARCHAR:
                return call.clone(call.getType(), ImmutableList.of(o0, o1b));
            }
        }
    }
    return node;
}
 

/**
 * Whether this Project requires a final column re-ordering. Returns False for all cases except when
 * convert_fromjson function is present.  For convert_fromjson function, the Project operator at
 * run-time produces an output schema with convert_fromjson expr appended to the end of the schema.
 * We need a final column re-ordering to ensure the correct column order.
*/
@Override
public boolean needsFinalColumnReordering() {
  for (RexNode expr : this.exps) {
    // TODO: a convert_fromjson nested within other convert functions currently does not work.
    // When it is supported, we should enhance this check by using a visitor to find the nested function.
    if (expr.getKind() == SqlKind.OTHER_FUNCTION &&
        expr instanceof RexCall &&
        ((RexCall) expr).getOperator().getName().equalsIgnoreCase("CONVERT_FROMJSON")) {
      return true;
    }
  }
  return false;
}
 

/** Returns whether an expression is definitely null.
 *
 * <p>The answer is based on calls to {@link #isNull} for its constituent
 * expressions, and you may override methods to test hypotheses such as
 * "if {@code x} is null, is {@code x + y} null? */
public boolean isNull(RexNode node) {
  final Policy policy = policy(node.getKind());
  switch (policy) {
  case NOT_NULL:
    return false;
  case ANY:
    return anyNull(((RexCall) node).getOperands());
  default:
    break;
  }

  switch (node.getKind()) {
  case LITERAL:
    return ((RexLiteral) node).isNull();
  // We can only guarantee AND to return NULL if both inputs are NULL  (similar for OR)
  // AND(NULL, FALSE) = FALSE
  case AND:
  case OR:
  case COALESCE:
    return allNull(((RexCall) node).getOperands());
  case NULLIF:
    // NULLIF(null, X) where X can be NULL, returns NULL
    // NULLIF(X, Y) where X is not NULL, then this may return NULL if X = Y, otherwise X.
    return allNull(ImmutableList.of(((RexCall) node).getOperands().get(0)));
  case INPUT_REF:
    return isNull((RexInputRef) node);
  case CASE:
    final RexCall caseCall = (RexCall) node;
    final List<RexNode> caseValues = new ArrayList<>();
    for (int i = 0; i < caseCall.getOperands().size(); i++) {
      if (!RexUtil.isCasePredicate(caseCall, i)) {
        caseValues.add(caseCall.getOperands().get(i));
      }
    }
    return allNull(caseValues);
  default:
    return false;
  }
}
 

private String getSingleFilterCondition(Implementor implementor, String op, RexCall call) {
  final String fieldName;
  final String literal;
  final RexNode left = call.operands.get(0);
  final RexNode right = call.operands.get(1);
  if (left.getKind() == LITERAL) {
    if (right.getKind() != INPUT_REF) {
      throw new IllegalArgumentException(
          "Expected a RexCall with a single field and single literal");
    } else {
      fieldName = implementor.getFieldName(this, ((RexInputRef) right).getIndex());
      literal = getLiteralAsString((RexLiteral) left);
    }
  } else if (right.getKind() == LITERAL) {
    if (left.getKind() != INPUT_REF) {
      throw new IllegalArgumentException(
          "Expected a RexCall with a single field and single literal");
    } else {
      fieldName = implementor.getFieldName(this, ((RexInputRef) left).getIndex());
      literal = getLiteralAsString((RexLiteral) right);
    }
  } else {
    throw new IllegalArgumentException(
        "Expected a RexCall with a single field and single literal");
  }

  return '(' + fieldName + ' ' + op + ' ' + literal + ')';
}
 

@Nullable
private static DruidJsonFilter toBetweenDruidFilter(RexNode rexNode, RelDataType rowType,
    DruidQuery query) {
  if (rexNode.getKind() != SqlKind.BETWEEN) {
    return null;
  }
  final RexCall rexCall = (RexCall) rexNode;
  if (rexCall.getOperands().size() < 4) {
    return null;
  }
  // BETWEEN (ASYMMETRIC, REF, 'lower-bound', 'upper-bound')
  final RexNode refNode = rexCall.getOperands().get(1);
  final RexNode lhs = rexCall.getOperands().get(2);
  final RexNode rhs = rexCall.getOperands().get(3);

  final String lhsLiteralValue = toDruidLiteral(lhs, rowType, query);
  final String rhsLiteralValue = toDruidLiteral(rhs, rowType, query);
  if (lhsLiteralValue == null || rhsLiteralValue == null) {
    return null;
  }
  final boolean isNumeric = lhs.getType().getFamily() == SqlTypeFamily.NUMERIC
      || lhs.getType().getFamily() == SqlTypeFamily.NUMERIC;
  final Pair<String, ExtractionFunction> druidColumn = DruidQuery
      .toDruidColumn(refNode, rowType, query);
  final String columnName = druidColumn.left;
  final ExtractionFunction extractionFunction = druidColumn.right;

  if (columnName == null) {
    return null;
  }
  return new JsonBound(columnName, lhsLiteralValue, false, rhsLiteralValue,
      false, isNumeric,
      extractionFunction);

}
 

@Deprecated // to be removed before 2.0
public static boolean analyzeSimpleEquiJoin(
    LogicalJoin join,
    int[] joinFieldOrdinals) {
  RexNode joinExp = join.getCondition();
  if (joinExp.getKind() != SqlKind.EQUALS) {
    return false;
  }
  RexCall binaryExpression = (RexCall) joinExp;
  RexNode leftComparand = binaryExpression.operands.get(0);
  RexNode rightComparand = binaryExpression.operands.get(1);
  if (!(leftComparand instanceof RexInputRef)) {
    return false;
  }
  if (!(rightComparand instanceof RexInputRef)) {
    return false;
  }

  final int leftFieldCount =
      join.getLeft().getRowType().getFieldCount();
  RexInputRef leftFieldAccess = (RexInputRef) leftComparand;
  if (!(leftFieldAccess.getIndex() < leftFieldCount)) {
    // left field must access left side of join
    return false;
  }

  RexInputRef rightFieldAccess = (RexInputRef) rightComparand;
  if (!(rightFieldAccess.getIndex() >= leftFieldCount)) {
    // right field must access right side of join
    return false;
  }

  joinFieldOrdinals[0] = leftFieldAccess.getIndex();
  joinFieldOrdinals[1] = rightFieldAccess.getIndex() - leftFieldCount;
  return true;
}
 

/** Converts an "t1.c1 = t2.c2" expression into two (input, field) pairs. */
private static int[][] grab(List<RelNode> leaves, RexNode rex) {
  switch (rex.getKind()) {
  case EQUALS:
    break;
  default:
    throw new AssertionError("only equi-join allowed");
  }
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  return new int[][] {
      inputField(leaves, operands.get(0)),
      inputField(leaves, operands.get(1))};
}
 

/** Converts an "t1.c1 = t2.c2" expression into two (input, field) pairs. */
private static int[][] grab(List<RelNode> leaves, RexNode rex) {
  switch (rex.getKind()) {
  case EQUALS:
    break;
  default:
    throw new AssertionError("only equi-join allowed");
  }
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  return new int[][] {
      inputField(leaves, operands.get(0)),
      inputField(leaves, operands.get(1))};
}
 

/**
 * Decomposes a predicate into a list of expressions that are AND'ed
 * together, and a list of expressions that are preceded by NOT.
 *
 * <p>For example, {@code a AND NOT b AND NOT (c and d) AND TRUE AND NOT
 * FALSE} returns {@code rexList = [a], notList = [b, c AND d]}.</p>
 *
 * <p>TRUE and NOT FALSE expressions are ignored. FALSE and NOT TRUE
 * expressions are placed on {@code rexList} and {@code notList} as other
 * expressions.</p>
 *
 * <p>For example, {@code a AND TRUE AND NOT TRUE} returns
 * {@code rexList = [a], notList = [TRUE]}.</p>
 *
 * @param rexPredicate predicate to be analyzed
 * @param rexList      list of decomposed RexNodes (except those with NOT)
 * @param notList      list of decomposed RexNodes that were prefixed NOT
 */
public static void decomposeConjunction(
    RexNode rexPredicate,
    List<RexNode> rexList,
    List<RexNode> notList) {
  if (rexPredicate == null || rexPredicate.isAlwaysTrue()) {
    return;
  }
  switch (rexPredicate.getKind()) {
  case AND:
    for (RexNode operand : ((RexCall) rexPredicate).getOperands()) {
      decomposeConjunction(operand, rexList, notList);
    }
    break;
  case NOT:
    final RexNode e = ((RexCall) rexPredicate).getOperands().get(0);
    if (e.isAlwaysFalse()) {
      return;
    }
    switch (e.getKind()) {
    case OR:
      final List<RexNode> ors = new ArrayList<>();
      decomposeDisjunction(e, ors);
      for (RexNode or : ors) {
        switch (or.getKind()) {
        case NOT:
          rexList.add(((RexCall) or).operands.get(0));
          break;
        default:
          notList.add(or);
        }
      }
      break;
    default:
      notList.add(e);
    }
    break;
  case LITERAL:
    if (!RexLiteral.isNullLiteral(rexPredicate)
        && RexLiteral.booleanValue(rexPredicate)) {
      return; // ignore TRUE
    }
    // fall through
  default:
    rexList.add(rexPredicate);
    break;
  }
}
 

/**
 * Convert {@link RexNode} condition into {@link SqlNode}
 *
 * @param node           condition Node
 * @param leftContext    LeftContext
 * @param rightContext   RightContext
 * @param leftFieldCount Number of field on left result
 * @return SqlJoin which represent the condition
 */
private SqlNode convertConditionToSqlNode(RexNode node,
                                          Context leftContext,
                                          Context rightContext,
                                          int leftFieldCount) {
  if (!(node instanceof RexCall)) {
    throw new AssertionError(node);
  }
  final List<RexNode> operands;
  final SqlOperator op;
  switch (node.getKind()) {
    case AND:
    case OR:
      operands = ((RexCall) node).getOperands();
      op = ((RexCall) node).getOperator();
      SqlNode sqlCondition = null;
      for (RexNode operand : operands) {
        SqlNode x = convertConditionToSqlNode(operand, leftContext,
            rightContext, leftFieldCount);
        if (sqlCondition == null) {
          sqlCondition = x;
        } else {
          sqlCondition = op.createCall(POS, sqlCondition, x);
        }
      }
      return sqlCondition;

    case EQUALS:
    case IS_NOT_DISTINCT_FROM:
    case NOT_EQUALS:
    case GREATER_THAN:
    case GREATER_THAN_OR_EQUAL:
    case LESS_THAN:
    case LESS_THAN_OR_EQUAL:
      operands = ((RexCall) node).getOperands();
      op = ((RexCall) node).getOperator();
      if (operands.get(0) instanceof RexInputRef
          && operands.get(1) instanceof RexInputRef) {
        final RexInputRef op0 = (RexInputRef) operands.get(0);
        final RexInputRef op1 = (RexInputRef) operands.get(1);

        if (op0.getIndex() < leftFieldCount
            && op1.getIndex() >= leftFieldCount) {
          // Arguments were of form 'op0 = op1'
          return op.createCall(POS,
              leftContext.field(op0.getIndex()),
              rightContext.field(op1.getIndex() - leftFieldCount));
        }
        if (op1.getIndex() < leftFieldCount
            && op0.getIndex() >= leftFieldCount) {
          // Arguments were of form 'op1 = op0'
          return reverseOperatorDirection(op).createCall(POS,
              leftContext.field(op1.getIndex()),
              rightContext.field(op0.getIndex() - leftFieldCount));
        }
      }
  }
  throw new AssertionError(node);
}
 

/**
 * @param rexNode    leaf Input Ref to Druid Column
 * @param rowType    row type
 * @param druidQuery druid query
 *
 * @return {@link Pair} of Column name and Extraction Function on the top of the input ref or
 * {@link Pair of(null, null)} when can not translate to valid Druid column
 */
protected static Pair<String, ExtractionFunction> toDruidColumn(RexNode rexNode,
    RelDataType rowType, DruidQuery druidQuery) {
  final String columnName;
  final ExtractionFunction extractionFunction;
  final Granularity granularity;
  switch (rexNode.getKind()) {
  case INPUT_REF:
    columnName = extractColumnName(rexNode, rowType, druidQuery);
    if (needUtcTimeExtract(rexNode)) {
      extractionFunction = TimeExtractionFunction.createDefault(
          DateTimeUtils.UTC_ZONE.getID());
    } else {
      extractionFunction = null;
    }
    break;
  case EXTRACT:
    granularity = DruidDateTimeUtils
        .extractGranularity(rexNode, druidQuery.getConnectionConfig().timeZone());
    if (granularity == null) {
      // unknown Granularity
      return Pair.of(null, null);
    }
    if (!TimeExtractionFunction.isValidTimeExtract((RexCall) rexNode)) {
      return Pair.of(null, null);
    }
    RexNode extractValueNode = ((RexCall) rexNode).getOperands().get(1);
    if (extractValueNode.getType().getSqlTypeName() == SqlTypeName.DATE
        || extractValueNode.getType().getSqlTypeName() == SqlTypeName.TIMESTAMP) {
      // Use 'UTC' at the extraction level
      extractionFunction =
          TimeExtractionFunction.createExtractFromGranularity(
              granularity, DateTimeUtils.UTC_ZONE.getID());
      columnName = extractColumnName(extractValueNode, rowType, druidQuery);
    } else if (extractValueNode.getType().getSqlTypeName()
        == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE) {
      // Use local time zone at the extraction level
      extractionFunction =
        TimeExtractionFunction.createExtractFromGranularity(
            granularity, druidQuery.getConnectionConfig().timeZone());
      columnName = extractColumnName(extractValueNode, rowType, druidQuery);
    } else {
      return Pair.of(null, null);
    }
    break;
  case FLOOR:
    granularity = DruidDateTimeUtils
        .extractGranularity(rexNode, druidQuery.getConnectionConfig().timeZone());
    if (granularity == null) {
      // unknown Granularity
      return Pair.of(null, null);
    }
    if (!TimeExtractionFunction.isValidTimeFloor((RexCall) rexNode)) {
      return Pair.of(null, null);
    }
    RexNode floorValueNode = ((RexCall) rexNode).getOperands().get(0);
    if (needUtcTimeExtract(floorValueNode)) {
      // Use 'UTC' at the extraction level, since all datetime types
      // are represented in 'UTC'
      extractionFunction =
          TimeExtractionFunction.createFloorFromGranularity(
              granularity, DateTimeUtils.UTC_ZONE.getID());
      columnName = extractColumnName(floorValueNode, rowType, druidQuery);
    } else {
      return Pair.of(null, null);
    }
    break;
  case CAST:
    // CASE we have a cast over InputRef. Check that cast is valid
    if (!isValidLeafCast(rexNode)) {
      return Pair.of(null, null);
    }
    RexNode operand0 = ((RexCall) rexNode).getOperands().get(0);
    columnName =
        extractColumnName(operand0, rowType, druidQuery);
    if (needUtcTimeExtract(rexNode)) {
      // CASE CAST to TIME/DATE need to make sure that we have valid extraction fn
      extractionFunction = TimeExtractionFunction.translateCastToTimeExtract(rexNode,
          TimeZone.getTimeZone(druidQuery.getConnectionConfig().timeZone()));
      if (extractionFunction == null) {
        // no extraction Function means cast is not valid thus bail out
        return Pair.of(null, null);
      }
    } else {
      extractionFunction = null;
    }
    break;
  default:
    return Pair.of(null, null);
  }
  return Pair.of(columnName, extractionFunction);
}
 

/** Pushes predicates into a CASE.
 *
 * <p>We have a loose definition of 'predicate': any boolean expression will
 * do, except CASE. For example '(CASE ...) = 5' or '(CASE ...) IS NULL'.
 */
public static RexCall pushPredicateIntoCase(RexCall call) {
  if (call.getType().getSqlTypeName() != SqlTypeName.BOOLEAN) {
    return call;
  }
  switch (call.getKind()) {
  case CASE:
  case AND:
  case OR:
    return call; // don't push CASE into CASE!
  case EQUALS: {
    // checks that the EQUALS operands may be splitted and
    // doesn't push EQUALS into CASE
    List<RexNode> equalsOperands = call.getOperands();
    ImmutableBitSet left = RelOptUtil.InputFinder.bits(equalsOperands.get(0));
    ImmutableBitSet right = RelOptUtil.InputFinder.bits(equalsOperands.get(1));
    if (!left.isEmpty() && !right.isEmpty() && left.intersect(right).isEmpty()) {
      return call;
    }
  }
  }
  int caseOrdinal = -1;
  final List<RexNode> operands = call.getOperands();
  for (int i = 0; i < operands.size(); i++) {
    RexNode operand = operands.get(i);
    if (operand.getKind() == SqlKind.CASE) {
      caseOrdinal = i;
    }
  }
  if (caseOrdinal < 0) {
    return call;
  }
  // Convert
  //   f(CASE WHEN p1 THEN v1 ... END, arg)
  // to
  //   CASE WHEN p1 THEN f(v1, arg) ... END
  final RexCall case_ = (RexCall) operands.get(caseOrdinal);
  final List<RexNode> nodes = new ArrayList<>();
  for (int i = 0; i < case_.getOperands().size(); i++) {
    RexNode node = case_.getOperands().get(i);
    if (!RexUtil.isCasePredicate(case_, i)) {
      node = substitute(call, caseOrdinal, node);
    }
    nodes.add(node);
  }
  return case_.clone(call.getType(), nodes);
}
 

private static RexNode stripFilter(RexNode node) {
  if (node.getKind() == SqlKind.IS_TRUE) {
    return ((RexCall) node).getOperands().get(0);
  }
  return node;
}