pig/src/main/java/org/apache/calcite/adapter/pig/PigAggregate.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to you under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.calcite.adapter.pig;

import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Aggregate;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.util.ImmutableBitSet;

import org.apache.pig.scripting.Pig;

import com.google.common.collect.ImmutableList;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/** Implementation of {@link org.apache.calcite.rel.core.Aggregate} in
 * {@link PigRel#CONVENTION Pig calling convention}. */
public class PigAggregate extends Aggregate implements PigRel {

  public static final String DISTINCT_FIELD_SUFFIX = "_DISTINCT";

  /** Creates a PigAggregate. */
  public PigAggregate(RelOptCluster cluster, RelTraitSet traitSet,
      RelNode input, ImmutableBitSet groupSet,
      List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
    super(cluster, traitSet, ImmutableList.of(), input, groupSet, groupSets, aggCalls);
    assert getConvention() == PigRel.CONVENTION;
  }

  @Deprecated // to be removed before 2.0
  public PigAggregate(RelOptCluster cluster, RelTraitSet traitSet,
      RelNode input, boolean indicator, ImmutableBitSet groupSet,
      List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
    this(cluster, traitSet, input, groupSet, groupSets, aggCalls);
    checkIndicator(indicator);
  }

  @Override public Aggregate copy(RelTraitSet traitSet, RelNode input,
      ImmutableBitSet groupSet, List<ImmutableBitSet> groupSets,
      List<AggregateCall> aggCalls) {
    return new PigAggregate(input.getCluster(), traitSet, input, groupSet,
        groupSets, aggCalls);
  }

  @Override public void implement(Implementor implementor) {
    implementor.visitChild(0, getInput());
    implementor.addStatement(getPigAggregateStatement(implementor));
  }

  /**
   * Generates a GROUP BY statement, followed by an optional FOREACH statement
   * for all aggregate functions used. e.g.
   * <pre>
   * {@code
   * A = GROUP A BY owner;
   * A = FOREACH A GENERATE group, SUM(A.pet_num);
   * }
   * </pre>
   */
  private String getPigAggregateStatement(Implementor implementor) {
    return getPigGroupBy(implementor) + '\n' + getPigForEachGenerate(implementor);
  }

  /**
   * Override this method so it looks down the tree to find the table this node
   * is acting on.
   */
  @Override public RelOptTable getTable() {
    return getInput().getTable();
  }

  /**
   * Generates the GROUP BY statement, e.g.
   * <code>A = GROUP A BY (f1, f2);</code>
   */
  private String getPigGroupBy(Implementor implementor) {
    final String relAlias = implementor.getPigRelationAlias(this);
    final List<RelDataTypeField> allFields = getInput().getRowType().getFieldList();
    final List<Integer> groupedFieldIndexes = groupSet.asList();
    if (groupedFieldIndexes.size() < 1) {
      return relAlias + " = GROUP " + relAlias + " ALL;";
    } else {
      final List<String> groupedFieldNames = new ArrayList<>(groupedFieldIndexes.size());
      for (int fieldIndex : groupedFieldIndexes) {
        groupedFieldNames.add(allFields.get(fieldIndex).getName());
      }
      return relAlias + " = GROUP " + relAlias + " BY ("
          + String.join(", ", groupedFieldNames) + ");";
    }
  }

  /**
   * Generates a FOREACH statement containing invocation of aggregate functions
   * and projection of grouped fields. e.g.
   * <code>A = FOREACH A GENERATE group, SUM(A.pet_num);</code>
   * @see Pig documentation for special meaning of the "group" field after GROUP
   *      BY.
   */
  private String getPigForEachGenerate(Implementor implementor) {
    final String relAlias = implementor.getPigRelationAlias(this);
    final String generateCall = getPigGenerateCall(implementor);
    final List<String> distinctCalls = getDistinctCalls(implementor);
    return relAlias + " = FOREACH " + relAlias + " {\n"
        + String.join(";\n", distinctCalls) + generateCall + "\n};";
  }

  private String getPigGenerateCall(Implementor implementor) {
    final List<Integer> groupedFieldIndexes = groupSet.asList();
    Set<String> groupFields = new HashSet<>(groupedFieldIndexes.size());
    for (int fieldIndex : groupedFieldIndexes) {
      final String fieldName = getInputFieldName(fieldIndex);
      // Pig appends group field name if grouping by multiple fields
      final String groupField = (groupedFieldIndexes.size() == 1 ? "group" : ("group." + fieldName))
          + " AS " + fieldName;

      groupFields.add(groupField);
    }
    final List<String> pigAggCalls = getPigAggregateCalls(implementor);
    List<String> allFields = new ArrayList<>(groupFields.size() + pigAggCalls.size());
    allFields.addAll(groupFields);
    allFields.addAll(pigAggCalls);
    return "  GENERATE " + String.join(", ", allFields) + ';';
  }

  private List<String> getPigAggregateCalls(Implementor implementor) {
    final String relAlias = implementor.getPigRelationAlias(this);
    final List<String> result = new ArrayList<>(aggCalls.size());
    for (AggregateCall ac : aggCalls) {
      result.add(getPigAggregateCall(relAlias, ac));
    }
    return result;
  }

  private String getPigAggregateCall(String relAlias, AggregateCall aggCall) {
    final PigAggFunction aggFunc = toPigAggFunc(aggCall);
    final String alias = aggCall.getName();
    final String fields = String.join(", ", getArgNames(relAlias, aggCall));
    return aggFunc.name() + "(" + fields + ") AS " + alias;
  }

  private static PigAggFunction toPigAggFunc(AggregateCall aggCall) {
    return PigAggFunction.valueOf(aggCall.getAggregation().getKind(),
        aggCall.getArgList().size() < 1);
  }

  private List<String> getArgNames(String relAlias, AggregateCall aggCall) {
    final List<String> result = new ArrayList<>(aggCall.getArgList().size());
    for (int fieldIndex : aggCall.getArgList()) {
      result.add(getInputFieldNameForAggCall(relAlias, aggCall, fieldIndex));
    }
    return result;
  }

  private String getInputFieldNameForAggCall(String relAlias, AggregateCall aggCall,
      int fieldIndex) {
    final String inputField = getInputFieldName(fieldIndex);
    return aggCall.isDistinct() ? (inputField + DISTINCT_FIELD_SUFFIX)
        : (relAlias + '.' + inputField);
  }

  /**
   * Returns the calls to aggregate functions that have the {@code DISTINT} flag.
   *
   * <p>An aggregate function call like <code>COUNT(DISTINCT COL)</code> in Pig
   * is achieved via two statements in a {@code FOREACH} that follows a
   * {@code GROUP} statement:
   *
   * <blockquote>
   * <code>
   * TABLE = GROUP TABLE ALL;<br>
   * TABLE = FOREACH TABLE {<br>
   * &nbsp;&nbsp;<b>COL.DISTINCT = DISTINCT COL;<br>
   * &nbsp;&nbsp;GENERATE COUNT(COL.DISTINCT) AS C;</b><br>
   * }</code>
   * </blockquote>
   */
  private List<String> getDistinctCalls(Implementor implementor) {
    final String relAlias = implementor.getPigRelationAlias(this);
    final List<String> result = new ArrayList<>();
    for (AggregateCall aggCall : aggCalls) {
      if (aggCall.isDistinct()) {
        for (int fieldIndex : aggCall.getArgList()) {
          String fieldName = getInputFieldName(fieldIndex);
          result.add("  " + fieldName + DISTINCT_FIELD_SUFFIX + " = DISTINCT "
              + relAlias + '.' + fieldName + ";\n");
        }
      }
    }
    return result;
  }

  private String getInputFieldName(int fieldIndex) {
    return getInput().getRowType().getFieldList().get(fieldIndex).getName();
  }
}