org.apache.calcite.sql.SqlSplittableAggFunction Java Examples

private static boolean isAggregateSupported(Aggregate aggregate, boolean allowFunctions) {
  if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
    return false;
  }
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out
  // If any aggregate call has a filter or is distinct, bail out
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class)
        == null) {
      return false;
    }
    if (aggregateCall.filterArg >= 0 || aggregateCall.isDistinct()) {
      return false;
    }
  }
  return true;
}
 

private static boolean isAggregateSupported(Aggregate aggregate, boolean allowFunctions) {
  if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
    return false;
  }
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out
  // If any aggregate call has a filter or is distinct, bail out
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class)
        == null) {
      return false;
    }
    if (aggregateCall.filterArg >= 0 || aggregateCall.isDistinct()) {
      return false;
    }
  }
  return true;
}
 

/** Creates a {@link org.apache.calcite.sql.SqlSplittableAggFunction.Registry}
 * that is a view of a list. */
private static <E> SqlSplittableAggFunction.Registry<E> registry(
    final List<E> list) {
  return e -> {
    int i = list.indexOf(e);
    if (i < 0) {
      i = list.size();
      list.add(e);
    }
    return i;
  };
}
 

private static boolean isAggregateSupported(Aggregate aggregate) {
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE
      || aggregate.getGroupCount() == 0) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out.
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.filterArg >= 0
        || aggregateCall.getAggregation()
            .unwrap(SqlSplittableAggFunction.class) == null) {
      return false;
    }
  }
  return true;
}
 

private boolean isAggregateSupported(AggregateCall aggCall) {
  if (aggCall.isDistinct()
      || aggCall.hasFilter()
      || aggCall.isApproximate()
      || aggCall.getArgList().size() > 1) {
    return false;
  }
  SqlSplittableAggFunction splitter = aggCall.getAggregation()
      .unwrap(SqlSplittableAggFunction.class);
  return splitter != null;
}
 

/** Creates a {@link org.apache.calcite.sql.SqlSplittableAggFunction.Registry}
 * that is a view of a list. */
private static <E> SqlSplittableAggFunction.Registry<E> registry(
    final List<E> list) {
  return e -> {
    int i = list.indexOf(e);
    if (i < 0) {
      i = list.size();
      list.add(e);
    }
    return i;
  };
}
 

/**
 * Creates a {@link org.apache.calcite.sql.SqlSplittableAggFunction.Registry}
 * that is a view of a list.
 */
private static <E> SqlSplittableAggFunction.Registry<E> registry(
		final List<E> list) {
	return new SqlSplittableAggFunction.Registry<E>() {
		public int register(E e) {
			int i = list.indexOf(e);
			if (i < 0) {
				i = list.size();
				list.add(e);
			}
			return i;
		}
	};
}
 

/**
 * Creates a {@link org.apache.calcite.sql.SqlSplittableAggFunction.Registry}
 * that is a view of a list.
 */
private static <E> SqlSplittableAggFunction.Registry<E> registry(
		final List<E> list) {
	return new SqlSplittableAggFunction.Registry<E>() {
		public int register(E e) {
			int i = list.indexOf(e);
			if (i < 0) {
				i = list.size();
				list.add(e);
			}
			return i;
		}
	};
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SumSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.Sum0Splitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SelfSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.CountSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate topAgg = call.rel(0);
  final Aggregate bottomAgg = call.rel(1);
  if (topAgg.getGroupCount() > bottomAgg.getGroupCount()) {
    return;
  }

  final ImmutableBitSet bottomGroupSet = bottomAgg.getGroupSet();
  final Map<Integer, Integer> map = new HashMap<>();
  bottomGroupSet.forEach(v -> map.put(map.size(), v));
  for (int k : topAgg.getGroupSet()) {
    if (!map.containsKey(k)) {
      return;
    }
  }

  // top aggregate keys must be subset of lower aggregate keys
  final ImmutableBitSet topGroupSet = topAgg.getGroupSet().permute(map);
  if (!bottomGroupSet.contains(topGroupSet)) {
    return;
  }

  boolean hasEmptyGroup = topAgg.getGroupSets()
      .stream().anyMatch(n -> n.isEmpty());

  final List<AggregateCall> finalCalls = new ArrayList<>();
  for (AggregateCall topCall : topAgg.getAggCallList()) {
    if (!isAggregateSupported(topCall)
        || topCall.getArgList().size() == 0) {
      return;
    }
    // Make sure top aggregate argument refers to one of the aggregate
    int bottomIndex = topCall.getArgList().get(0) - bottomGroupSet.cardinality();
    if (bottomIndex >= bottomAgg.getAggCallList().size()
        || bottomIndex < 0) {
      return;
    }
    AggregateCall bottomCall = bottomAgg.getAggCallList().get(bottomIndex);
    // Should not merge if top agg with empty group keys and the lower agg
    // function is COUNT, because in case of empty input for lower agg,
    // the result is empty, if we merge them, we end up with 1 result with
    // 0, which is wrong.
    if (!isAggregateSupported(bottomCall)
        || (bottomCall.getAggregation() == SqlStdOperatorTable.COUNT
             && hasEmptyGroup)) {
      return;
    }
    SqlSplittableAggFunction splitter = Objects.requireNonNull(
        bottomCall.getAggregation().unwrap(SqlSplittableAggFunction.class));
    AggregateCall finalCall = splitter.merge(topCall, bottomCall);
    // fail to merge the aggregate call, bail out
    if (finalCall == null) {
      return;
    }
    finalCalls.add(finalCall);
  }

  // re-map grouping sets
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (topAgg.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(topAgg.getGroupSets(), map));
  }

  final Aggregate finalAgg =
      topAgg.copy(topAgg.getTraitSet(), bottomAgg.getInput(), topGroupSet,
          newGroupingSets, finalCalls);
  call.transformTo(finalAgg);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SelfSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final RelNode input = aggregate.getInput();
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()))) {
    return;
  }

  final RelBuilder relBuilder = call.builder();
  final RexBuilder rexBuilder = relBuilder.getRexBuilder();
  final List<RexNode> projects = new ArrayList<>();
  for (AggregateCall aggCall : aggregate.getAggCallList()) {
    final SqlAggFunction aggregation = aggCall.getAggregation();
    if (aggregation.getKind() == SqlKind.SUM0) {
      // Bail out for SUM0 to avoid potential infinite rule matching,
      // because it may be generated by transforming SUM aggregate
      // function to SUM0 and COUNT.
      return;
    }
    final SqlSplittableAggFunction splitter =
        Objects.requireNonNull(
            aggregation.unwrap(SqlSplittableAggFunction.class));
    final RexNode singleton = splitter.singleton(
        rexBuilder, input.getRowType(), aggCall);
    projects.add(singleton);
  }

  final RelNode newInput = convert(input, aggregate.getTraitSet().simplify());
  relBuilder.push(newInput);
  if (!projects.isEmpty()) {
    projects.addAll(0, relBuilder.fields(aggregate.getGroupSet()));
    relBuilder.project(projects);
  } else if (newInput.getRowType().getFieldCount()
      > aggregate.getRowType().getFieldCount()) {
    // If aggregate was projecting a subset of columns, and there were no
    // aggregate functions, add a project for the same effect.
    relBuilder.project(relBuilder.fields(aggregate.getGroupSet()));
  }
  call.getPlanner().prune(aggregate);
  call.transformTo(relBuilder.build());
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SelfSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.CountSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SelfSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.Sum0Splitter.INSTANCE);
  }
  return super.unwrap(clazz);
}
 

@Override public <T> T unwrap(Class<T> clazz) {
  if (clazz == SqlSplittableAggFunction.class) {
    return clazz.cast(SqlSplittableAggFunction.SumSplitter.INSTANCE);
  }
  return super.unwrap(clazz);
}