Java Code Examples for org.apache.calcite.rel.metadata.RelMetadataQuery#getColumnOrigins()

/**
 * extract origins of the fields
 * @param graph the root node of the query
 * @param rowType the rowType after validation
 * @return the origins of the fields in the query. at the same index as the original rowType
 */
public static List<FieldOrigin> getFieldOrigins(RelNode graph, RelDataType rowType) {
  if (graph.getRowType().getFieldCount() != rowType.getFieldCount()) {
    throw new IllegalArgumentException(format(
        "graph and rowType should have the same field count:\ngraph: %s\nrowType: %s",
        graph, rowType));
  }
  List<FieldOrigin> definitions = new ArrayList<>();
  RelMetadataQuery query = graph.getCluster().getMetadataQuery();
  for (int i = 0; i < graph.getRowType().getFieldCount(); i++) {
    Set<RelColumnOrigin> origins = query.getColumnOrigins(graph, i);
    List<Origin> namedOrigins = new ArrayList<>();
    for (RelColumnOrigin relColumnOrigin : origins) {
      List<String> table = Origins.getTable(relColumnOrigin);
      String colName = Origins.getColName(relColumnOrigin);
      namedOrigins.add(new Origin(colName, relColumnOrigin.isDerived()).setTableList(table));
    }
    // we need the rowtype after validation and before planning
    // graph.getRowType() may not have to be the user facing rowtype anymore
    // (even if it often does)
    String name = rowType.getFieldList().get(i).getName();
    definitions.add(i, new FieldOrigin(name).setOriginsList(namedOrigins));
  }
  return definitions;
}
 

public Set<RelColumnOrigin> getColumnOrigins(ConvertFromJsonPrel rel, RelMetadataQuery mq, int iOutputColumn) {
  final List<ConversionColumn> conversions = rel.getConversions();
  if (iOutputColumn < conversions.size()) {
    return null; // just return null for now, it should work for ConvertFromJsonConverter
  }
  return mq.getColumnOrigins(rel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(RelSubset rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(MoreObjects.firstNonNull(rel.getBest(), rel.getOriginal()), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(LimitRelBase rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(rel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(ExpansionNode rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(rel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(CustomPrel rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(rel.getOriginPrel(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(SampleRelBase rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(rel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(DictionaryLookupPrel rel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(rel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(JdbcRelBase jdbc, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(jdbc.getSubTree(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(ExchangePrel exchangePrel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(exchangePrel.getInput(), iOutputColumn);
}
 

public Set<RelColumnOrigin> getColumnOrigins(SelectionVectorRemoverPrel selectionVectorRemoverPrel, RelMetadataQuery mq, int iOutputColumn) {
  return mq.getColumnOrigins(selectionVectorRemoverPrel.getInput(), iOutputColumn);
}
 

@SuppressWarnings("unused") // Called through reflection
public Set<RelColumnOrigin> getColumnOrigins(LogicalWindow window, RelMetadataQuery mq, int iOutputColumn) {
  final RelNode inputRel = window.getInput();
  final int numFieldsInInput = inputRel.getRowType().getFieldCount();
  if (iOutputColumn < numFieldsInInput) {
    return mq.getColumnOrigins(inputRel, iOutputColumn);
  }

  if (iOutputColumn >= window.getRowType().getFieldCount()) {
    return Collections.emptySet();
  }

  int startGroupIdx = iOutputColumn - numFieldsInInput;
  int curentIdx = 0;
  Group finalGroup = null;
  for (Group group : window.groups) {
    curentIdx += group.aggCalls.size();
    if (curentIdx > startGroupIdx) {
      // this is the group
      finalGroup = group;
      break;
    }
  }
  Preconditions.checkNotNull(finalGroup);
  // calculating index of the aggCall within a group
  // currentIdx = through idx within groups/aggCalls (max currentIdx = sum(groups size * aggCals_per_group) )
  // since currentIdx at this moment points to the end of the group substracting aggCals_per_group
  // to get to the beginning of the group and have startGroupIdx substract the diff
  final int aggCalIdx = startGroupIdx - (curentIdx - finalGroup.aggCalls.size());
  Preconditions.checkElementIndex(aggCalIdx, finalGroup.aggCalls.size());

  final Set<RelColumnOrigin> set = new HashSet<>();
  // Add aggregation column references
  final RexWinAggCall aggCall = finalGroup.aggCalls.get(aggCalIdx);
  for (RexNode operand : aggCall.operands) {
    if (operand instanceof RexInputRef) {
      final RexInputRef opInputRef = (RexInputRef) operand;
      if (opInputRef.getIndex() < numFieldsInInput) {
        Set<RelColumnOrigin> inputSet =
          mq.getColumnOrigins(inputRel, opInputRef.getIndex());
        inputSet = createDerivedColumnOrigins(inputSet);
        if (inputSet != null) {
          set.addAll(inputSet);
        }
      }
    }
  }

  return set;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LogicalAggregate agg = call.rel(0);
  final RelNode input = agg.getInput();

  boolean distinctReplaced = false;
  List<AggregateCall> calls = new ArrayList<>();

  RelMetadataQuery query = null;

  final Boolean[] memo = new Boolean[agg.getInput().getRowType().getFieldCount()];
  for (AggregateCall c : agg.getAggCallList()) {
    final boolean candidate = c.isDistinct() && c.getArgList().size() == 1 && "COUNT".equals(c.getAggregation().getName());

    if(!candidate) {
      calls.add(c);
      continue;
    }

    final int inputOrdinal = c.getArgList().get(0);
    boolean allowed = false;
    if(memo[inputOrdinal] != null) {
      allowed = memo[inputOrdinal];
    } else {
      if(query == null) {
        query = agg.getCluster().getMetadataQuery();
      }

      Set<RelColumnOrigin> origins = query.getColumnOrigins(input, inputOrdinal);

      // see if any column origin allowed a transformation.
      for(RelColumnOrigin o : origins) {
        RelOptTable table = o.getOriginTable();
        NamespaceTable namespaceTable = table.unwrap(NamespaceTable.class);
        if(namespaceTable == null) {
          // unable to decide, no way to transform.
          return;
        }

        if(namespaceTable.isApproximateStatsAllowed()) {
          allowed = true;
        }
      }

      memo[inputOrdinal] = allowed;

    }


    if(allowed) {
      calls.add(AggregateCall.create(HyperLogLog.NDV, false, c.getArgList(), -1, c.getType(), c.getName()));
      distinctReplaced = true;
    } else {
      calls.add(c);
    }

  }

  if(!distinctReplaced) {
    return;
  }

  final RelBuilder builder = relBuilderFactory.create(agg.getCluster(), null);
  builder.push(agg.getInput());
  builder.aggregate(builder.groupKey(agg.getGroupSet().toArray()), calls);
  call.transformTo(builder.build());
}
 

private Set<RelColumnOrigin> checkColumnOrigin(String sql) {
  RelNode rel = convertSql(sql);
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
  return mq.getColumnOrigins(rel, 0);
}