Java Code Examples for org.apache.calcite.linq4j.Ord

/**
 * Returns a permutation describing where the Project's fields come from
 * after the Project is pushed down.
 */
public static Mappings.TargetMapping permutationPushDownProject(
    List<RexNode> nodes,
    RelDataType inputRowType,
    int sourceOffset,
    int targetOffset) {
  final Mappings.TargetMapping mapping =
      Mappings.create(MappingType.PARTIAL_FUNCTION,
          inputRowType.getFieldCount() + sourceOffset,
          nodes.size() + targetOffset);
  for (Ord<RexNode> node : Ord.zip(nodes)) {
    if (node.e instanceof RexInputRef) {
      mapping.set(
          ((RexInputRef) node.e).getIndex() + sourceOffset,
          node.i + targetOffset);
    }
  }
  return mapping;
}
 

private List<GenericRecord> toArray(NavigableMap<Long, VersionedValue> object) {
    List<GenericRecord> records = new ArrayList<>();
    Schema recordSchema = avroSchema.getElementType();
    for (VersionedValue versionedValue : object.values()) {
        Comparable[] value = versionedValue.getValue();
        GenericRecordBuilder builder = new GenericRecordBuilder(recordSchema);
        for (Ord<Field> field : Ord.zip(recordSchema.getFields())) {
            if (field.i == 0) {
                builder.set(field.e, versionedValue.getVersion());
            } else if (field.i == 1) {
                builder.set(field.e, versionedValue.getCommit());
            } else if (field.i == 2) {
                builder.set(field.e, versionedValue.isDeleted());
            } else {
                if (!versionedValue.isDeleted()) {
                    Object v = AvroSchema.toAvroValue(field.e.schema(), value[field.i - 3]);
                    if (v != null) {
                        builder.set(field.e, v);
                    }
                }
            }
        }
        records.add(builder.build());
    }
    return records;
}
 

/**
 * Projects all {@code input} output fields plus the additional expressions.
 *
 * @param input        Input relational expression
 * @param additionalExprs Additional expressions and names
 * @return the new Project
 */
private RelNode createProjectWithAdditionalExprs(
    RelNode input,
    List<Pair<RexNode, String>> additionalExprs) {
  final List<RelDataTypeField> fieldList =
      input.getRowType().getFieldList();
  List<Pair<RexNode, String>> projects = new ArrayList<>();
  for (Ord<RelDataTypeField> field : Ord.zip(fieldList)) {
    projects.add(
        Pair.of(
            (RexNode) relBuilder.getRexBuilder().makeInputRef(
                field.e.getType(), field.i),
            field.e.getName()));
  }
  projects.addAll(additionalExprs);
  return relBuilder.push(input)
      .projectNamed(Pair.left(projects), Pair.right(projects), true)
      .build();
}
 

/**
 * Projects all {@code input} output fields plus the additional expressions.
 *
 * @param input        Input relational expression
 * @param additionalExprs Additional expressions and names
 * @return the new Project
 */
private RelNode createProjectWithAdditionalExprs(
    RelNode input,
    List<Pair<RexNode, String>> additionalExprs) {
  final List<RelDataTypeField> fieldList =
      input.getRowType().getFieldList();
  List<Pair<RexNode, String>> projects = new ArrayList<>();
  for (Ord<RelDataTypeField> field : Ord.zip(fieldList)) {
    projects.add(
        Pair.of(
            (RexNode) relBuilder.getRexBuilder().makeInputRef(
                field.e.getType(), field.i),
            field.e.getName()));
  }
  projects.addAll(additionalExprs);
  return relBuilder.push(input)
      .projectNamed(Pair.left(projects), Pair.right(projects), true)
      .build();
}
 

public boolean checkOperandTypes(
	SqlCallBinding callBinding,
	boolean throwOnFailure) {

	for (Ord<SqlNode> op : Ord.zip(callBinding.operands())) {
		if (!checkSingleOperandType(
			callBinding,
			op.e,
			false)) {
			// TODO: check type coercion when we support implicit type conversion
			// recheck to validate.
			for (Ord<SqlNode> op1 : Ord.zip(callBinding.operands())) {
				if (!checkSingleOperandType(
					callBinding,
					op1.e,
					throwOnFailure)) {
					return false;
				}
			}
			return false;
		}
	}
	return true;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  FilterPrel filter = call.rel(0);
  NestedLoopJoinPrel join = call.rel(1);

  if ((join.getProjectedFields() == null) || join.getProjectedFields().cardinality() == join.getInputRowType().getFieldCount()) {
    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), filter.getCondition()), join.getProjectedFields()));
  } else {
    // Current filter condition is written based on projected fields on join. In order to push this filter down we need to rewrite filter condition
    final ImmutableBitSet topProjectedColumns = RelOptUtil.InputFinder.bits(filter.getCondition());
    final ImmutableBitSet bottomProjectedColumns = join.getProjectedFields();

    Mapping mapping = Mappings.create(MappingType.SURJECTION, join.getRowType().getFieldCount(), join.getInputRowType().getFieldCount());
    for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
      if (topProjectedColumns.get(ord.i)) {
        mapping.set(ord.i, ord.e);
      }
    }

    RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
    RexNode updatedCondition = shuttle.apply(filter.getCondition());

    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), updatedCondition), join.getProjectedFields()));
  }
}
 

public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
  if (checkers.size() != callBinding.getOperandCount()) {
    // assume this is an inapplicable sub-rule of a composite rule;
    // don't throw
    return false;
  }

  for (Ord<SqlNode> op : Ord.zip(callBinding.operands())) {
    if (!checkSingleOperandType(
        callBinding,
        op.e,
        op.i,
        throwOnFailure)) {
      return false;
    }
  }
  return true;
}
 

public Measure(SqlAggFunction agg, boolean distinct, @Nullable String name,
    Iterable<Column> args) {
  this.agg = Objects.requireNonNull(agg);
  this.distinct = distinct;
  this.name = name;
  this.args = ImmutableList.copyOf(args);

  final StringBuilder b = new StringBuilder()
      .append(agg)
      .append(distinct ? "(DISTINCT " : "(");
  for (Ord<Column> arg : Ord.zip(this.args)) {
    if (arg.i > 0) {
      b.append(", ");
    }
    if (arg.e instanceof BaseColumn) {
      b.append(((BaseColumn) arg.e).table);
      b.append('.');
      b.append(((BaseColumn) arg.e).column);
    } else {
      b.append(arg.e.alias);
    }
  }
  b.append(')');
  this.digest = b.toString();
}
 

/**
 * Creates a new projection based on the original projection, adjusting all
 * input refs using an adjustment array passed in. If there was no original
 * projection, create a new one that selects every field from the underlying
 * rel.
 *
 * <p>If the resulting projection would be trivial, return the child.
 *
 * @param projChild   child of the new project
 * @param adjustments array indicating how much each input reference should
 *                    be adjusted by
 * @return the created projection
 */
public RelNode createNewProject(RelNode projChild, int[] adjustments) {
  final List<Pair<RexNode, String>> projects = new ArrayList<>();

  if (origProj != null) {
    for (Pair<RexNode, String> p : origProj.getNamedProjects()) {
      projects.add(
          Pair.of(
              convertRefsAndExprs(
                  p.left,
                  projChild.getRowType().getFieldList(),
                  adjustments),
              p.right));
    }
  } else {
    for (Ord<RelDataTypeField> field : Ord.zip(childFields)) {
      projects.add(
          Pair.of(
              rexBuilder.makeInputRef(
                  field.e.getType(), field.i), field.e.getName()));
    }
  }
  return relBuilder.push(projChild)
      .project(Pair.left(projects), Pair.right(projects))
      .build();
}
 

/**
 * Returns a permutation describing where output fields come from. In
 * the returned map, value of {@code map.getTargetOpt(i)} is {@code n} if
 * field {@code i} projects input field {@code n}, -1 if it is an
 * expression.
 */
public static Mappings.TargetMapping permutation(
    List<RexNode> nodes,
    RelDataType inputRowType) {
  final Mappings.TargetMapping mapping =
      Mappings.create(
          MappingType.PARTIAL_FUNCTION,
          nodes.size(),
          inputRowType.getFieldCount());
  for (Ord<RexNode> node : Ord.zip(nodes)) {
    if (node.e instanceof RexInputRef) {
      mapping.set(
          node.i,
          ((RexInputRef) node.e).getIndex());
    }
  }
  return mapping;
}
 

Run(final List<Column> columns) {
  for (Ord<Column> column : Ord.zip(columns)) {
    if (column.e.ordinal != column.i) {
      throw new IllegalArgumentException();
    }
  }
  this.columns = columns;
  this.singletonSpaces =
      new ArrayList<>(Collections.nCopies(columns.size(), null));
  for (ImmutableBitSet ordinals
      : ImmutableBitSet.range(columns.size()).powerSet()) {
    final Space space = new Space(ordinals, toColumns(ordinals));
    spaces.add(space);
    if (ordinals.cardinality() == 1) {
      singletonSpaces.set(ordinals.nth(0), space);
    }
  }
}
 

/** Generates code for a pattern.
 *
 * <p>For example, for the pattern {@code (A B)}, generates
 * {@code patternBuilder.symbol("A").symbol("B").seq()}. */
private Expression implementPattern(Expression patternBuilder_,
    RexNode pattern) {
  switch (pattern.getKind()) {
  case LITERAL:
    final String symbol = ((RexLiteral) pattern).getValueAs(String.class);
    return Expressions.call(patternBuilder_,
        BuiltInMethod.PATTERN_BUILDER_SYMBOL.method,
        Expressions.constant(symbol));

  case PATTERN_CONCAT:
    final RexCall concat = (RexCall) pattern;
    for (Ord<RexNode> operand : Ord.zip(concat.operands)) {
      patternBuilder_ = implementPattern(patternBuilder_, operand.e);
      if (operand.i > 0) {
        patternBuilder_ = Expressions.call(patternBuilder_,
            BuiltInMethod.PATTERN_BUILDER_SEQ.method);
      }
    }
    return patternBuilder_;

  default:
    throw new AssertionError("unknown kind: " + pattern);
  }
}
 

public List<FunctionParameter> getParameters() {
  final List<FunctionParameter> parameters = new ArrayList<>();
  for (final Ord<RelProtoDataType> type : Ord.zip(getParams())) {
    parameters.add(
        new FunctionParameter() {
          public int getOrdinal() {
            return type.i;
          }

          public String getName() {
            return "arg" + type.i;
          }

          public RelDataType getType(RelDataTypeFactory typeFactory) {
            return type.e.apply(typeFactory);
          }

          public boolean isOptional() {
            return false;
          }
        });
  }
  return parameters;
}
 

/**
 * Projects all {@code input} output fields plus the additional expressions.
 *
 * @param input        Input relational expression
 * @param additionalExprs Additional expressions and names
 * @return the new Project
 */
private RelNode createProjectWithAdditionalExprs(
    RelNode input,
    List<Pair<RexNode, String>> additionalExprs) {
  final List<RelDataTypeField> fieldList =
      input.getRowType().getFieldList();
  List<Pair<RexNode, String>> projects = new ArrayList<>();
  Ord.forEach(fieldList, (field, i) ->
      projects.add(
          Pair.of(relBuilder.getRexBuilder().makeInputRef(field.getType(), i),
              field.getName())));
  projects.addAll(additionalExprs);
  return relBuilder.push(input)
      .projectNamed(Pair.left(projects), Pair.right(projects), true)
      .build();
}
 

/**
 * Finds type and new ordinal relative to new inputs by oldOrdinal and
 * innerOrdinal indexes.
 *
 * @param oldOrdinal   ordinal of the field relative to old inputs
 * @param innerOrdinal when oldOrdinal points to struct and target field
 *                     is inner field of struct, this argument should contain
 *                     calculated field's ordinal within struct after flattening.
 *                     Otherwise when oldOrdinal points to primitive field, this
 *                     argument should be zero.
 * @return flat type with new ordinal relative to new inputs
 */
private Ord<RelDataType> getNewFieldForOldInput(int oldOrdinal, int innerOrdinal) {
  assert currentRel != null;
  // sum of predecessors post flatten sizes points to new ordinal
  // of flat field or first field of flattened struct
  final int postFlatteningOrdinal = currentRel.getInputs().stream()
      .flatMap(node -> node.getRowType().getFieldList().stream())
      .limit(oldOrdinal)
      .map(RelDataTypeField::getType)
      .mapToInt(this::postFlattenSize)
      .sum();
  final int newOrdinal = postFlatteningOrdinal + innerOrdinal;
  // NoSuchElementException may be thrown because of two reasons:
  // 1. postFlattenSize() didn't predict well
  // 2. innerOrdinal has wrong value
  RelDataTypeField newField = getNewInputFieldByNewOrdinal(newOrdinal);
  return Ord.of(newOrdinal, newField.getType());
}
 

/** Creates a reference to a field which originated in a relation with the
 * given alias. Searches for the relation starting at the top of the
 * stack. */
public RexNode field(int inputCount, String alias, String fieldName) {
  Objects.requireNonNull(alias);
  Objects.requireNonNull(fieldName);
  final List<String> fields = new ArrayList<>();
  for (int inputOrdinal = 0; inputOrdinal < inputCount; ++inputOrdinal) {
    final Frame frame = peek_(inputOrdinal);
    for (Ord<Field> p
        : Ord.zip(frame.fields)) {
      // If alias and field name match, reference that field.
      if (p.e.left.contains(alias)
          && p.e.right.getName().equals(fieldName)) {
        return field(inputCount, inputCount - 1 - inputOrdinal, p.i);
      }
      fields.add(
          String.format(Locale.ROOT, "{aliases=%s,fieldName=%s}", p.e.left,
              p.e.right.getName()));
    }
  }
  throw new IllegalArgumentException("{alias=" + alias + ",fieldName=" + fieldName + "} "
      + "field not found; fields are: " + fields);
}
 

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

private static int lookupGroupExpr(GroupAnalyzer groupAnalyzer,
    SqlNode expr) {
  for (Ord<SqlNode> node : Ord.zip(groupAnalyzer.groupExprs)) {
    if (node.e.equalsDeep(expr, Litmus.IGNORE)) {
      return node.i;
    }
  }

  switch (expr.getKind()) {
  case HOP:
  case TUMBLE:
  case SESSION:
    groupAnalyzer.extraExprs.add(expr);
    break;
  }
  groupAnalyzer.groupExprs.add(expr);
  return groupAnalyzer.groupExprs.size() - 1;
}
 

public String getAllowedSignatures(String opName) {
  StringBuilder ret = new StringBuilder();
  for (Ord<SqlTypeName> typeName : Ord.zip(SqlTypeName.STRING_TYPES)) {
    if (typeName.i > 0) {
      ret.append(NL);
    }
    ret.append(
        SqlUtil.getAliasedSignature(this, opName,
            ImmutableList.of(typeName.e, SqlTypeName.INTEGER)));
    ret.append(NL);
    ret.append(
        SqlUtil.getAliasedSignature(this, opName,
            ImmutableList.of(typeName.e, SqlTypeName.INTEGER,
                SqlTypeName.INTEGER)));
  }
  return ret.toString();
}
 

private boolean argTypesValid(SqlCallBinding callBinding) {
  if (callBinding.getOperandCount() < 2) {
    return true; // nothing to compare
  }
  RelDataType firstType = null;
  for (Ord<SqlNode> operand : Ord.zip(callBinding.operands())) {
    RelDataType type =
        callBinding.getValidator().deriveType(
            callBinding.getScope(),
            operand.e);
    if (operand.i == 0) {
      firstType = type;
    } else {
      if (!SqlTypeUtil.sameNamedType(firstType, type)) {
        return false;
      }
    }
  }
  return true;
}
 

/**
 * Accepts a {@link SqlVisitor}, and tells it to visit each child.
 *
 * @param visitor Visitor
 */
public <R> void acceptCall(
    SqlVisitor<R> visitor,
    SqlCall call,
    boolean onlyExpressions,
    SqlBasicVisitor.ArgHandler<R> argHandler) {
  if (onlyExpressions) {
    for (Ord<SqlNode> operand : Ord.zip(call.getOperandList())) {
      // if the second param is an Identifier then it's supposed to
      // be a name from a window clause and isn't part of the
      // group by check
      if (operand == null) {
        continue;
      }
      if (operand.i == 1 && operand.e instanceof SqlIdentifier) {
        continue;
      }
      argHandler.visitChild(visitor, call, operand.i, operand.e);
    }
  } else {
    super.acceptCall(visitor, call, onlyExpressions, argHandler);
  }
}
 

/**
 * Checks that the operand values in a {@link SqlCall} to this operator are
 * valid. Subclasses must either override this method or supply an instance
 * of {@link SqlOperandTypeChecker} to the constructor.
 *
 * @param callBinding    description of call
 * @param throwOnFailure whether to throw an exception if check fails
 *                       (otherwise returns false in that case)
 * @return whether check succeeded
 */
public boolean checkOperandTypes(
    SqlCallBinding callBinding,
    boolean throwOnFailure) {
  // Check that all of the operands are of the right type.
  if (null == operandTypeChecker) {
    // If you see this you must either give operandTypeChecker a value
    // or override this method.
    throw Util.needToImplement(this);
  }

  if (kind != SqlKind.ARGUMENT_ASSIGNMENT) {
    for (Ord<SqlNode> operand : Ord.zip(callBinding.operands())) {
      if (operand.e != null
          && operand.e.getKind() == SqlKind.DEFAULT
          && !operandTypeChecker.isOptional(operand.i)) {
        throw callBinding.newValidationError(RESOURCE.defaultForOptionalParameter());
      }
    }
  }

  return operandTypeChecker.checkOperandTypes(
      callBinding,
      throwOnFailure);
}
 

public String getAllowedSignatures(SqlOperator op, String opName) {
  if (allowedSignatures != null) {
    return allowedSignatures;
  }
  if (composition == Composition.SEQUENCE) {
    throw new AssertionError(
        "specify allowedSignatures or override getAllowedSignatures");
  }
  StringBuilder ret = new StringBuilder();
  for (Ord<SqlOperandTypeChecker> ord
      : Ord.<SqlOperandTypeChecker>zip(allowedRules)) {
    if (ord.i > 0) {
      ret.append(SqlOperator.NL);
    }
    ret.append(ord.e.getAllowedSignatures(op, opName));
    if (composition == Composition.AND) {
      break;
    }
  }
  return ret.toString();
}
 

@Override public boolean checkOperandTypesWithoutTypeCoercion(SqlCallBinding callBinding,
    boolean throwOnFailure) {
  if (families.size() != callBinding.getOperandCount()) {
    // assume this is an inapplicable sub-rule of a composite rule;
    // don't throw exception.
    return false;
  }

  for (Ord<SqlNode> op : Ord.zip(callBinding.operands())) {
    if (!checkSingleOperandType(
        callBinding,
        op.e,
        op.i,
        throwOnFailure)) {
      return false;
    }
  }
  return true;
}
 

public String getAllowedSignatures(SqlOperator op, String opName) {
  StringBuilder sb = new StringBuilder();
  sb.append(opName);
  sb.append("(");
  for (Ord<RelDataType> paramType : Ord.zip(paramTypes)) {
    if (paramType.i > 0) {
      sb.append(", ");
    }
    if (paramNames != null) {
      sb.append(paramNames.get(paramType.i))
          .append(" => ");
    }
    sb.append("<");
    sb.append(paramType.e.getFamily());
    sb.append(">");
  }
  sb.append(")");
  return sb.toString();
}
 

public <R> void acceptCall(
    SqlVisitor<R> visitor,
    SqlCall call,
    boolean onlyExpressions,
    SqlBasicVisitor.ArgHandler<R> argHandler) {
  if (onlyExpressions) {
    for (Ord<SqlNode> operand : Ord.zip(call.getOperandList())) {
      // if the second param is an Identifier then it's supposed to
      // be a name from a window clause and isn't part of the
      // group by check
      if (operand.e == null) {
        continue;
      }
      if (operand.i == 1 && operand.e instanceof SqlIdentifier) {
        // skip refName
        continue;
      }
      argHandler.visitChild(visitor, call, operand.i, operand.e);
    }
  } else {
    super.acceptCall(visitor, call, onlyExpressions, argHandler);
  }
}
 

/**
 * Permutes argument types to correspond to the order of parameter names.
 */
private static List<RelDataType> permuteArgTypes(SqlFunction function,
    List<String> argNames, List<RelDataType> argTypes) {
  // Arguments passed by name. Make sure that the function has
  // parameters of all of these names.
  Map<Integer, Integer> map = new HashMap<>();
  for (Ord<String> argName : Ord.zip(argNames)) {
    int i = function.getParamNames().indexOf(argName.e);
    if (i < 0) {
      return null;
    }
    map.put(i, argName.i);
  }
  return Functions.generate(function.getParamTypes().size(), index -> {
    if (map.containsKey(index)) {
      return argTypes.get(map.get(index));
    } else {
      return null;
    }
  });
}
 

private GenericRecord toRecord(Comparable[] object) {
    GenericRecordBuilder builder = new GenericRecordBuilder(avroSchema);
    for (Ord<Field> field : Ord.zip(avroSchema.getFields())) {
        Comparable c = object[field.i];
        builder.set(field.e, AvroSchema.toAvroValue(field.e.schema(), c));
    }
    return builder.build();
}
 

@Override
public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
    final BlockBuilder builder = new BlockBuilder();
    Expression unionExp = null;
    for (Ord<RelNode> ord : Ord.zip(inputs)) {
        EnumerableRel input = (EnumerableRel) ord.e;
        final Result result = implementor.visitChild(this, ord.i, input, pref);
        Expression childExp =
                builder.append(
                        "child" + ord.i,
                        result.block);

        if (unionExp == null) {
            unionExp = childExp;
        } else {
            unionExp = createUnionExpression(unionExp, childExp, result.format == JavaRowFormat.ARRAY);
        }
    }

    builder.add(unionExp);
    final PhysType physType =
            PhysTypeImpl.of(
                    implementor.getTypeFactory(),
                    getRowType(),
                    pref.prefer(JavaRowFormat.CUSTOM));
    return implementor.result(physType, builder.toBlock());
}
 

public static List<NamedExpression> aggsToExpr(
    RelDataType rowType, RelNode input, ImmutableBitSet groupSet, List<AggregateCall> aggCalls) {
  final List<String> fields = rowType.getFieldNames();
  final List<String> childFields = input.getRowType().getFieldNames();
  final List<NamedExpression> aggExprs = Lists.newArrayList();
  for (Ord<AggregateCall> aggCall : Ord.zip(aggCalls)) {
    int aggExprOrdinal = groupSet.cardinality() + aggCall.i;
    FieldReference ref = FieldReference.getWithQuotedRef(fields.get(aggExprOrdinal));
    LogicalExpression expr = toExpr(aggCall.e, childFields);
    NamedExpression ne = new NamedExpression(expr, ref);
    aggExprs.add(ne);
  }
  return aggExprs;
}