Java Code Examples for org.apache.calcite.rel.core.Window#Group

/**
 * @see #dispatch
 * https://issues.apache.org/jira/projects/CALCITE/issues/CALCITE-3112
 */
public Result visit(Window e) {
  Result x = visitChild(0, e.getInput());
  Builder builder = x.builder(e);
  RelNode input = e.getInput();
  int inputFieldCount = input.getRowType().getFieldCount();
  final List<SqlNode> rexOvers = new ArrayList<>();
  for (Window.Group group : e.groups) {
    rexOvers.addAll(builder.context.toSql(group, e.constants, inputFieldCount));
  }
  final List<SqlNode> selectList = new ArrayList<>();
  for (RelDataTypeField field : input.getRowType().getFieldList()) {
    addSelect(selectList, builder.context.field(field.getIndex()), e.getRowType());
  }
  for (SqlNode rexOver : rexOvers) {
    addSelect(selectList, rexOver, e.getRowType());
  }
  builder.setSelect(new SqlNodeList(selectList, POS));
  return builder.result();
}
 

/** @see #dispatch */
public Result visit(Window e) {
  Result x = visitChild(0, e.getInput());
  Builder builder = x.builder(e);
  RelNode input = e.getInput();
  int inputFieldCount = input.getRowType().getFieldCount();
  final List<SqlNode> rexOvers = new ArrayList<>();
  for (Window.Group group: e.groups) {
    rexOvers.addAll(builder.context.toSql(group, e.constants, inputFieldCount));
  }
  final List<SqlNode> selectList = new ArrayList<>();

  for (RelDataTypeField field: input.getRowType().getFieldList()) {
    addSelect(selectList, builder.context.field(field.getIndex()), e.getRowType());
  }

  for (SqlNode rexOver: rexOvers) {
    addSelect(selectList, rexOver, e.getRowType());
  }

  builder.setSelect(new SqlNodeList(selectList, POS));
  return builder.result();
}
 

@Override
public boolean matches(RelOptRuleCall call) {
  final DrillWindowRel oldWinRel = (DrillWindowRel) call.rels[0];
  for (Window.Group group : oldWinRel.groups) {
    for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
      if (isConversionToSumZeroNeeded(rexWinAggCall.getOperator(), rexWinAggCall.getType())) {
        return true;
      }
    }
  }
  return false;
}
 

/**
 * Create a RelCollation that has partition-by as the leading keys followed by order-by keys
 * @param window The window specification
 * @return a RelCollation with {partition-by keys, order-by keys}
 */
private RelCollation getCollation(Window.Group window) {
    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : BitSets.toIter(window.keys)) {
        fields.add(new RelFieldCollation(group));
    }

    fields.addAll(window.orderKeys.getFieldCollations());

    return RelCollations.of(fields);
}
 

private List<DrillDistributionTrait.DistributionField> getDistributionFields(Window.Group window) {
    List<DrillDistributionTrait.DistributionField> groupByFields = Lists.newArrayList();
    for (int group : BitSets.toIter(window.keys)) {
        DrillDistributionTrait.DistributionField field = new DrillDistributionTrait.DistributionField(group);
        groupByFields.add(field);
    }

    return groupByFields;
}
 

private List<DistributionField> getDistributionFieldsFromCollation(Window.Group window) {
    List<DistributionField> distFields = Lists.newArrayList();

    for (RelFieldCollation relField : window.collation().getFieldCollations()) {
        DistributionField field = new DistributionField(relField.getFieldIndex());
        distFields.add(field);
    }

    return distFields;
}
 

/**
 * Create a RelCollation that has partition-by as the leading keys followed by order-by keys
 * @param window The window specification
 * @return a RelCollation with {partition-by keys, order-by keys}
 */
private RelCollation getCollation(Window.Group window) {
  List<RelFieldCollation> fields = Lists.newArrayList();
  for (int group : BitSets.toIter(window.keys)) {
    fields.add(new RelFieldCollation(group));
  }

  fields.addAll(window.orderKeys.getFieldCollations());

  return RelCollations.of(fields);
}
 

private List<DistributionTrait.DistributionField> getDistributionFields(Window.Group window) {
  List<DistributionTrait.DistributionField> groupByFields = Lists.newArrayList();
  for (int group : BitSets.toIter(window.keys)) {
    DistributionTrait.DistributionField field = new DistributionTrait.DistributionField(group);
    groupByFields.add(field);
  }

  return groupByFields;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final DrillWindowRel oldWinRel = (DrillWindowRel) call.rels[0];
  final ImmutableList.Builder<Window.Group> builder = ImmutableList.builder();

  for (Window.Group group : oldWinRel.groups) {
    final List<Window.RexWinAggCall> aggCalls = Lists.newArrayList();
    for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
      if (isConversionToSumZeroNeeded(rexWinAggCall.getOperator(), rexWinAggCall.getType())) {
        final RelDataType argType = rexWinAggCall.getType();
        final RelDataType sumType = oldWinRel.getCluster().getTypeFactory()
            .createTypeWithNullability(argType, argType.isNullable());
        final SqlAggFunction sumZeroAgg = new DrillCalciteSqlAggFunctionWrapper(
            new SqlSumEmptyIsZeroAggFunction(), sumType);
        final Window.RexWinAggCall sumZeroCall =
              RexBuilder.getRexFactory().makeWinAggCall(
                sumZeroAgg,
                sumType,
                rexWinAggCall.getOperands(),
                rexWinAggCall.getOrdinal(),
                rexWinAggCall.isDistinct());
        aggCalls.add(sumZeroCall);
      } else {
        aggCalls.add(rexWinAggCall);
      }
    }

    final Window.Group newGroup = new Window.Group(
        group.keys,
        group.isRows,
        group.lowerBound,
        group.upperBound,
        group.orderKeys,
        aggCalls);
    builder.add(newGroup);
  }

  call.transformTo(new DrillWindowRel(
      oldWinRel.getCluster(),
      oldWinRel.getTraitSet(),
      oldWinRel.getInput(),
      oldWinRel.constants,
      oldWinRel.getRowType(),
      builder.build()));
}
 

public static List<RexOver> getOver(Window w) {
  RelNode input = w.getInput();
  int inputFieldCount = input.getRowType().getFieldCount();
  final List<RexOver> rexOvers = new ArrayList<>();
  for (Window.Group group : w.groups) {
    final List<RexNode> partitionKeys = new ArrayList<>();
    final List<RexFieldCollation> orderKeys = new ArrayList<>();

    // Convert RelFieldCollation to RexFieldCollation
    for (RelFieldCollation collation : group.orderKeys.getFieldCollations()) {
      Set<SqlKind> directions = new HashSet<>();
      if (collation.direction.isDescending()) {
        directions.add(SqlKind.DESCENDING);
      }
      if (collation.nullDirection == RelFieldCollation.NullDirection.LAST) {
        directions.add(SqlKind.NULLS_LAST);

      } else if (collation.nullDirection == RelFieldCollation.NullDirection.FIRST) {
        directions.add(SqlKind.NULLS_FIRST);
      }

      RexFieldCollation rexCollation = new RexFieldCollation(w.getCluster().getRexBuilder().makeInputRef(input, collation.getFieldIndex()), directions);
      orderKeys.add(rexCollation);
    }

    // Add partition keys
    for (int partition : group.keys) {
      partitionKeys.add(w.getCluster().getRexBuilder().makeInputRef(input, partition));
    }

    // Create RexWindow
    RexWindow window = new RexWindow(partitionKeys, orderKeys, group.lowerBound, group.upperBound, group.isRows);

    // For each window agg call, create rex over
    for (Window.RexWinAggCall winAggCall : group.aggCalls) {

      RexShuttle replaceConstants = new RexShuttle() {
        @Override public RexNode visitInputRef(RexInputRef inputRef) {
          int index = inputRef.getIndex();
          RexNode ref;
          if (index > inputFieldCount - 1) {
            ref = w.constants.get(index - inputFieldCount);
          } else {
            ref = inputRef;
          }
          return ref;
        }
      };
      RexCall aggCall = (RexCall) winAggCall.accept(replaceConstants);
      SqlAggFunction aggFunction = (SqlAggFunction) winAggCall.getOperator();
      RexOver over = new RexOver(winAggCall.getType(), aggFunction, aggCall.operands, window, winAggCall.distinct);
      rexOvers.add(over);
    }
  }
  return rexOvers;
}
 

public List<SqlNode> toSql(Window.Group group, ImmutableList<RexLiteral> constants,
                           int inputFieldCount) {
  final List<SqlNode> rexOvers = new ArrayList<>();
  final List<SqlNode> partitionKeys = new ArrayList<>();
  final List<SqlNode> orderByKeys = new ArrayList<>();
  for (int partition : group.keys) {
    partitionKeys.add(this.field(partition));
  }
  for (RelFieldCollation collation : group.orderKeys.getFieldCollations()) {
    this.addOrderItem(orderByKeys, collation);
  }
  SqlLiteral isRows = SqlLiteral.createBoolean(group.isRows, POS);
  SqlNode lowerBound = null;
  SqlNode upperBound = null;

  final SqlLiteral allowPartial = null;

  for (Window.RexWinAggCall winAggCall : group.aggCalls) {
    SqlAggFunction aggFunction = (SqlAggFunction) winAggCall.getOperator();
    final SqlWindow sqlWindow = SqlWindow.create(null, null,
      new SqlNodeList(partitionKeys, POS), new SqlNodeList(orderByKeys, POS),
      isRows, lowerBound, upperBound, allowPartial, POS);
    if (aggFunction.allowsFraming()) {
      lowerBound = createSqlWindowBound(group.lowerBound);
      upperBound = createSqlWindowBound(group.upperBound);
      sqlWindow.setLowerBound(lowerBound);
      sqlWindow.setUpperBound(upperBound);
    }

    RexShuttle replaceConstants = new RexShuttle() {
      @Override
      public RexNode visitInputRef(RexInputRef inputRef) {
        int index = inputRef.getIndex();
        RexNode ref;
        if (index > inputFieldCount - 1) {
          ref = constants.get(index - inputFieldCount);
        } else {
          ref = inputRef;
        }
        return ref;
      }
    };
    RexCall aggCall = (RexCall) winAggCall.accept(replaceConstants);
    List<SqlNode> operands = toSql(null, aggCall.operands);
    rexOvers.add(createOverCall(aggFunction, operands, sqlWindow));
  }
  return rexOvers;
}
 

/** @see #dispatch */
public Result visit(Window e) {
  final Result x = visitChild(0, e.getInput());
  final Builder builder = x.builder(e, Clause.SELECT);
  final List<SqlNode> selectList =
      new ArrayList<>(builder.context.fieldList());

  for (Window.Group winGroup : e.groups) {
    final List<SqlNode> partitionList = Expressions.list();
    for (int i : winGroup.keys) {
      partitionList.add(builder.context.field(i));
    }

    final List<SqlNode> orderList = Expressions.list();
    for (RelFieldCollation orderKey : winGroup.collation().getFieldCollations()) {
      orderList.add(builder.context.toSql(orderKey));
    }

    final SqlNode sqlWindow =  SqlWindow.create(
        null, // Window declaration name
        null, // Window reference name
        new SqlNodeList(partitionList, POS),
        new SqlNodeList(orderList, POS),
        SqlLiteral.createBoolean(winGroup.isRows, POS),
        builder.context.toSql(winGroup.lowerBound),
        builder.context.toSql(winGroup.upperBound),
        null, // allowPartial
        POS);

    for (Window.RexWinAggCall winFunc : winGroup.aggCalls) {
      final List<SqlNode> winFuncOperands = Expressions.list();
      for (RexNode operand : winFunc.getOperands()) {
        winFuncOperands.add(builder.context.toSql(null, operand));
      }
      SqlNode aggFunc = winFunc.getOperator().createCall(new SqlNodeList(winFuncOperands, POS));
      selectList.add(SqlStdOperatorTable.OVER.createCall(POS, aggFunc, sqlWindow));
    }
    builder.setSelect(new SqlNodeList(selectList, POS));
  }
  return builder.result();
}
 

public List<SqlNode> toSql(Window.Group group, ImmutableList<RexLiteral> constants,
    int inputFieldCount) {
  final List<SqlNode> rexOvers = new ArrayList<>();
  final List<SqlNode> partitionKeys = new ArrayList<>();
  final List<SqlNode> orderByKeys = new ArrayList<>();
  for (int partition: group.keys) {
    partitionKeys.add(this.field(partition));
  }
  for (RelFieldCollation collation: group.orderKeys.getFieldCollations()) {
    this.addOrderItem(orderByKeys, collation);
  }
  SqlLiteral isRows = SqlLiteral.createBoolean(group.isRows, POS);
  SqlNode lowerBound = null;
  SqlNode upperBound = null;

  final SqlLiteral allowPartial = null;

  for (Window.RexWinAggCall winAggCall: group.aggCalls) {
    SqlAggFunction aggFunction = (SqlAggFunction) winAggCall.getOperator();
    final SqlWindow sqlWindow = SqlWindow.create(null, null,
            new SqlNodeList(partitionKeys, POS), new SqlNodeList(orderByKeys, POS),
            isRows, lowerBound, upperBound, allowPartial, POS);
    if (aggFunction.allowsFraming()) {
      lowerBound = createSqlWindowBound(group.lowerBound);
      upperBound = createSqlWindowBound(group.upperBound);
      sqlWindow.setLowerBound(lowerBound);
      sqlWindow.setUpperBound(upperBound);
    }

    RexShuttle replaceConstants = new RexShuttle() {
      @Override public RexNode visitInputRef(RexInputRef inputRef) {
        int index = inputRef.getIndex();
        RexNode ref;
        if (index > inputFieldCount - 1) {
          ref = constants.get(index - inputFieldCount);
        } else {
          ref = inputRef;
        }
        return ref;
      }
    };
    RexCall aggCall = (RexCall) winAggCall.accept(replaceConstants);
    List<SqlNode> operands = toSql(null, aggCall.operands);
    rexOvers.add(createOverCall(aggFunction, operands, sqlWindow));
  }
  return rexOvers;
}
 

/**
 * Helper method to determine a
 * {@link org.apache.calcite.rel.core.Window}'s collation.
 *
 * <p>A Window projects the fields of its input first, followed by the output
 * from each of its windows. Assuming (quite reasonably) that the
 * implementation does not re-order its input rows, then any collations of its
 * input are preserved.
 */
public static List<RelCollation> window(
		RelMetadataQuery mq,
		RelNode input,
		com.google.common.collect.ImmutableList<Window.Group> groups) {
	return mq.collations(input);
}
 

/** Helper method to determine a
 * {@link org.apache.calcite.rel.core.Window}'s collation.
 *
 * <p>A Window projects the fields of its input first, followed by the output
 * from each of its windows. Assuming (quite reasonably) that the
 * implementation does not re-order its input rows, then any collations of its
 * input are preserved. */
public static List<RelCollation> window(RelMetadataQuery mq, RelNode input,
    ImmutableList<Window.Group> groups) {
  return mq.collations(input);
}
 

/** Helper method to determine a
 * {@link org.apache.calcite.rel.core.Window}'s collation.
 *
 * <p>A Window projects the fields of its input first, followed by the output
 * from each of its windows. Assuming (quite reasonably) that the
 * implementation does not re-order its input rows, then any collations of its
 * input are preserved. */
public static List<RelCollation> window(RelMetadataQuery mq, RelNode input,
    ImmutableList<Window.Group> groups) {
  return mq.collations(input);
}